Browse Source

added test timeou

Mish Ushakov 2 months ago
parent
commit
36d2bc40ce
6 changed files with 12 additions and 15 deletions
  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
 // Run the scraper in streaming mode
 const { stream } = await scraper.stream(page, schema, {
 const { stream } = await scraper.stream(page, schema, {
   format: 'html',
   format: 'html',
-  output: 'array',
 })
 })
 
 
 // Stream the result from LLM
 // Stream the result from LLM

+ 2 - 2
package-lock.json

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

+ 1 - 1
package.json

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

+ 7 - 10
src/models.ts

@@ -11,16 +11,11 @@ import { ScraperLLMOptions } from './index.js'
 import { PreProcessResult } from './preprocess.js'
 import { PreProcessResult } from './preprocess.js'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 
 
-export type ScraperCompletionResult<T extends z.ZodSchema<any>> = {
-  data: z.infer<T>
-  url: string
-}
-
 const defaultPrompt =
 const defaultPrompt =
   'You are a sophisticated web scraper. Extract the contents of the webpage'
   '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) {
 function stripMarkdownBackticks(text: string) {
   const match = text.match(/^```(.*)\n(.*)/)
   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,
   model: LanguageModelV1,
   page: PreProcessResult,
   page: PreProcessResult,
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   options?: ScraperLLMOptions
   options?: ScraperLLMOptions
 ) {
 ) {
   const content = prepareAISDKPage(page)
   const content = prepareAISDKPage(page)
-  const { partialObjectStream } = await streamObject<T>({
+  const { partialObjectStream } = streamObject<T>({
     model,
     model,
     messages: [
     messages: [
       { role: 'system', content: options?.prompt || defaultPrompt },
       { role: 'system', content: options?.prompt || defaultPrompt },
@@ -88,6 +83,7 @@ export async function streamAISDKCompletions<T>(
     temperature: options?.temperature,
     temperature: options?.temperature,
     maxTokens: options?.maxTokens,
     maxTokens: options?.maxTokens,
     topP: options?.topP,
     topP: options?.topP,
+    mode: options?.mode,
   })
   })
 
 
   return {
   return {
@@ -102,7 +98,8 @@ export async function generateAISDKCode<T>(
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   schema: z.Schema<T, z.ZodTypeDef, any> | Schema<T>,
   options?: ScraperLLMOptions
   options?: ScraperLLMOptions
 ) {
 ) {
-  const parsedSchema = schema instanceof z.ZodType ? zodToJsonSchema(schema) : schema
+  const parsedSchema =
+    schema instanceof z.ZodType ? zodToJsonSchema(schema) : schema
 
 
   const result = await generateText({
   const result = await generateText({
     model,
     model,

+ 1 - 1
tests/streaming.test.ts

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

+ 1 - 0
vitest.config.ts

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