浏览代码

minor adjustments

Mish Ushakov 1 年之前
父节点
当前提交
16f161fb5b
共有 2 个文件被更改,包括 12 次插入13 次删除
  1. 9 9
      src/index.ts
  2. 3 4
      tests/lib.ts

+ 9 - 9
src/index.ts

@@ -14,8 +14,8 @@ type ScraperLoadResult = {
   mode: ScraperLoadOptions['mode']
 }
 
-type ScraperRunOptions = {
-  schema: z.ZodSchema<any>
+type ScraperRunOptions<Z extends z.ZodSchema<any>> = {
+  schema: Z
   model?: OpenAI.Chat.ChatModel
 } & ScraperLoadOptions
 
@@ -93,11 +93,11 @@ export default class LLMScraper {
   }
 
   // Generate completion using OpenAI
-  private async generateCompletions(
+  private async generateCompletions<T extends z.ZodSchema<any>>(
     model: OpenAI.Chat.ChatModel = 'gpt-4-turbo',
-    schema: z.ZodSchema<any>,
+    schema: T,
     pages: ScraperLoadResult[]
-  ) {
+  ): Promise<z.infer<typeof schema>[]> {
     const openai = new OpenAI()
     return pages.map(async (page) => {
       const content = this.preparePage(page)
@@ -120,11 +120,11 @@ export default class LLMScraper {
   }
 
   // Load pages and generate completion
-  async run(
+  async run<T extends z.ZodSchema<any>>(
     url: string | string[],
-    options: ScraperRunOptions
-  ): Promise<z.infer<(typeof options)['schema']>> {
+    options: ScraperRunOptions<T>
+  ) {
     const pages = await this.load(url, options)
-    return this.generateCompletions(options.model, options.schema, pages)
+    return this.generateCompletions<T>(options.model, options.schema, pages)
   }
 }

+ 3 - 4
tests/lib.ts

@@ -9,11 +9,10 @@ const schema = z.object({
   title:z.string().describe('Title of the page'),
 })
 
-type schema = z.infer<typeof schema>
-
 const pages = await scraper.run(['https://example.com'], {
   schema,
-  mode: 'text'
+  mode: 'text',
+  closeOnFinish: true,
 })
 
-console.log(await Promise.all(pages))
+console.log(pages[0])