Parcourir la source

added test timeou

Mish Ushakov il y a 2 mois
Parent
commit
36d2bc40ce
6 fichiers modifiés avec 12 ajouts et 15 suppressions
  1. 0 1
      examples/streaming.ts
  2. 2 2
      package-lock.json
  3. 1 1
      package.json
  4. 7 10
      src/models.ts
  5. 1 1
      tests/streaming.test.ts
  6. 1 0
      vitest.config.ts

+ 0 - 1
examples/streaming.ts

@@ -27,7 +27,6 @@ const schema = z.object({
 // Run the scraper in streaming mode
 const { stream } = await scraper.stream(page, schema, {
   format: 'html',
-  output: 'array',
 })
 
 // Stream the result from LLM

+ 2 - 2
package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "llm-scraper",
-  "version": "1.5.1",
+  "version": "1.6.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "llm-scraper",
-      "version": "1.5.1",
+      "version": "1.6.0",
       "license": "MIT",
       "dependencies": {
         "@ai-sdk/provider": "^1.1.3",

+ 1 - 1
package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "llm-scraper",
-  "version": "1.5.1",
+  "version": "1.6.0",
   "description": "Turn any webpage intro structured data using LLMs",
   "main": "dist/index.js",
   "scripts": {

+ 7 - 10
src/models.ts

@@ -11,16 +11,11 @@ import { ScraperLLMOptions } from './index.js'
 import { PreProcessResult } from './preprocess.js'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 
-export type ScraperCompletionResult<T extends z.ZodSchema<any>> = {
-  data: z.infer<T>
-  url: string
-}
-
 const defaultPrompt =
   'You are a sophisticated web scraper. Extract the contents of the webpage'
 
-const defaultCodePrompt = `Provide a scraping function in JavaScript that extracts and formats data according to a schema from the current page.
-The function must be IIFE. No comments or imports. The code you generate will be executed straight away, you shouldn't output anything besides runnable code.`
+const defaultCodePrompt =
+  "Provide a scraping function in JavaScript that extracts and formats data according to a schema from the current page. The function must be IIFE. No comments or imports. The code you generate will be executed straight away, you shouldn't output anything besides runnable code."
 
 function stripMarkdownBackticks(text: string) {
   const match = text.match(/^```(.*)\n(.*)/)
@@ -70,14 +65,14 @@ export async function generateAISDKCompletions<T>(
   }
 }
 
-export async function streamAISDKCompletions<T>(
+export function streamAISDKCompletions<T>(
   model: LanguageModelV1,
   page: PreProcessResult,
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   options?: ScraperLLMOptions
 ) {
   const content = prepareAISDKPage(page)
-  const { partialObjectStream } = await streamObject<T>({
+  const { partialObjectStream } = streamObject<T>({
     model,
     messages: [
       { role: 'system', content: options?.prompt || defaultPrompt },
@@ -88,6 +83,7 @@ export async function streamAISDKCompletions<T>(
     temperature: options?.temperature,
     maxTokens: options?.maxTokens,
     topP: options?.topP,
+    mode: options?.mode,
   })
 
   return {
@@ -102,7 +98,8 @@ export async function generateAISDKCode<T>(
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   options?: ScraperLLMOptions
 ) {
-  const parsedSchema = schema instanceof z.ZodType ? zodToJsonSchema(schema) : schema
+  const parsedSchema =
+    schema instanceof z.ZodType ? zodToJsonSchema(schema) : schema
 
   const result = await generateText({
     model,

+ 1 - 1
tests/streaming.test.ts

@@ -24,4 +24,4 @@ test('streaming', async ({ page, scraper }) => {
   }
 
   expect(last).toHaveLength(5)
-}, 10000)
+})

+ 1 - 0
vitest.config.ts

@@ -3,5 +3,6 @@ import { defineConfig } from 'vitest/config'
 export default defineConfig({
   test: {
     include: ['tests/**/*.test.ts'],
+    testTimeout: 30000,
   },
 })