Bläddra i källkod

added test timeou

Mish Ushakov 2 månader sedan
förälder
incheckning
36d2bc40ce
6 ändrade filer med 12 tillägg och 15 borttagningar
  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,
   },
 })