瀏覽代碼

added test timeou

Mish Ushakov 2 月之前
父節點
當前提交
36d2bc40ce
共有 6 個文件被更改,包括 12 次插入15 次删除
  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,
   },
 })