1303 lines · 43.3 KB · TypeScript
RawDownload
import { encoding_for_model } from "@dqbd/tiktoken";import { TiktokenModel } from "@dqbd/tiktoken";import {  Document,  JsonFormatWithOptions,  TokenUsage,} from "../../../controllers/v2/types";import { Logger } from "winston";import { Meta } from "..";import { logger } from "../../../lib/logger";import { modelPrices } from "../../../lib/extract/usage/model-prices";import {  AISDKError,  generateObject,  generateText,  LanguageModel,  NoObjectGeneratedError,  jsonSchema,} from "ai";import { getModel } from "../../../lib/generic-ai";import { z } from "zod";import fs from "fs/promises";import Ajv from "ajv";import { extractData } from "../lib/extractSmartScrape";import { CostTracking } from "../../../lib/extract/extraction-service";import { isAgentExtractModelValid } from "../../../controllers/v1/types";import { hasFormatOfType } from "../../../lib/format-utils";// TODO: fix this, it's horribletype LanguageModelV1ProviderMetadata = {  anthropic?: {    thinking?: {      type: "enabled" | "disabled";      budgetTokens?: number;    };    tool_choice?: "auto" | "none" | "required";  };}; // Get max tokens from model pricesconst getModelLimits = (model: string) => {  const modelConfig = modelPrices[model];  if (!modelConfig) {    // Default fallback values    return {      maxInputTokens: 8192,      maxOutputTokens: 4096,      maxTokens: 12288,    };  }  return {    maxInputTokens: modelConfig.max_input_tokens || modelConfig.max_tokens,    maxOutputTokens: modelConfig.max_output_tokens || modelConfig.max_tokens,    maxTokens: modelConfig.max_tokens,  };}; export class LLMRefusalError extends Error {  public refusal: string;   constructor(refusal: string) {    super("LLM refused to extract the website's content");    this.refusal = refusal;  }} function normalizeSchema(x: any): any {  if (typeof x !== "object" || x === null) return x;   if (x["$defs"] !== null && typeof x["$defs"] === "object") {    x["$defs"] = Object.fromEntries(      Object.entries(x["$defs"]).map(([name, schema]) => [        name,        normalizeSchema(schema),      ]),    );  }   if (x && x.anyOf) {    x.anyOf = x.anyOf.map(x => normalizeSchema(x));  }   if (x && x.oneOf) {    x.oneOf = x.oneOf.map(x => normalizeSchema(x));  }   if (x && x.allOf) {    x.allOf = x.allOf.map(x => normalizeSchema(x));  }   if (x && x.not) {    x.not = normalizeSchema(x.not);  }   if (x && x.type === "object") {    return {      ...x,      properties: Object.fromEntries(        Object.entries(x.properties || {}).map(([k, v]) => [          k,          normalizeSchema(v),        ]),      ),      required: Object.keys(x.properties || {}),      additionalProperties: false,    };  } else if (x && x.type === "array") {    return {      ...x,      items: normalizeSchema(x.items),    };  } else {    return x;  }} interface TrimResult {  text: string;  numTokens: number;  warning?: string;} export function trimToTokenLimit(  text: string,  maxTokens: number,  modelId: string = "gpt-4o",  previousWarning?: string,): TrimResult {  try {    const encoder = encoding_for_model(modelId as TiktokenModel);    try {      const tokens = encoder.encode(text);      const numTokens = tokens.length;       if (numTokens <= maxTokens) {        return { text, numTokens };      }       const modifier = 3;      // Start with 3 chars per token estimation      let currentText = text.slice(0, Math.floor(maxTokens * modifier) - 1);       // Keep trimming until we're under the token limit      while (true) {        const currentTokens = encoder.encode(currentText);        if (currentTokens.length <= maxTokens) {          const warning = `The extraction content would have used more tokens (${numTokens}) than the maximum we allow (${maxTokens}). -- the input has been automatically trimmed.`;          return {            text: currentText,            numTokens: currentTokens.length,            warning: previousWarning              ? `${warning} ${previousWarning}`              : warning,          };        }        const overflow = currentTokens.length * modifier - maxTokens - 1;        // If still over limit, remove another chunk        currentText = currentText.slice(          0,          Math.floor(currentText.length - overflow),        );      }    } catch (e) {      throw e;    } finally {      encoder.free();    }  } catch (error) {    // Fallback to a more conservative character-based approach    const estimatedCharsPerToken = 2.8;    const safeLength = maxTokens * estimatedCharsPerToken;    const trimmedText = text.slice(0, Math.floor(safeLength));     const warning = `Failed to derive number of LLM tokens the extraction might use -- the input has been automatically trimmed to the maximum number of tokens (${maxTokens}) we support.`;     return {      text: trimmedText,      numTokens: maxTokens, // We assume we hit the max in this fallback case      warning: previousWarning ? `${warning} ${previousWarning}` : warning,    };  }} export function calculateCost(  model: string,  inputTokens: number,  outputTokens: number,) {  const modelCosts = {    "openai/o3-mini": { input_cost: 1.1, output_cost: 4.4 },    "gpt-4o-mini": { input_cost: 0.15, output_cost: 0.6 },    "openai/gpt-4o-mini": { input_cost: 0.15, output_cost: 0.6 },    "openai/gpt-4o": { input_cost: 2.5, output_cost: 10 },    "google/gemini-2.0-flash-001": { input_cost: 0.15, output_cost: 0.6 },    "gemini-2.0-flash": { input_cost: 0.15, output_cost: 0.6 },    "deepseek/deepseek-r1": { input_cost: 0.55, output_cost: 2.19 },    "google/gemini-2.0-flash-thinking-exp:free": {      input_cost: 0.55,      output_cost: 2.19,    },    "google/gemini-2.5-flash-lite": { input_cost: 0.1, output_cost: 0.4 },  };  let modelCost = modelCosts[model] || { input_cost: 0, output_cost: 0 };  //gemini-2.5-pro-exp-03-25 pricing  if (model.includes("gemini-2.5-pro")) {    let inputCost = 0;    let outputCost = 0;    if (inputTokens <= 200000) {      inputCost = 1.25;      outputCost = 10.0;    } else {      inputCost = 2.5;      outputCost = 15.0;    }    modelCost = { input_cost: inputCost, output_cost: outputCost };  }  const totalCost =    (inputTokens * modelCost.input_cost +      outputTokens * modelCost.output_cost) /    1_000_000;   return totalCost;} export type GenerateCompletionsOptions = {  model?: LanguageModel;  logger: Logger;  options: Omit<JsonFormatWithOptions, "type"> & {    systemPrompt?: string;    temperature?: number;  };  markdown?: string;  previousWarning?: string;  isExtractEndpoint?: boolean;  mode?: "object" | "no-object";  providerOptions?: LanguageModelV1ProviderMetadata;  retryModel?: LanguageModel;  costTrackingOptions: {    costTracking: CostTracking;    metadata: Record<string, any>;  };  metadata: {    teamId: string;    functionId?: string;    extractId?: string;    scrapeId?: string;    deepResearchId?: string;    llmsTxtId?: string;  };};export async function generateCompletions({  logger,  options,  markdown,  previousWarning,  isExtractEndpoint,  model = getModel("gpt-4o-mini", "openai"),  mode = "object",  providerOptions,  retryModel = getModel("claude-3-5-sonnet-20240620", "anthropic"),  costTrackingOptions,  metadata,}: GenerateCompletionsOptions): Promise<{  extract: any;  numTokens: number;  warning: string | undefined;  totalUsage: TokenUsage;  model: string;}> {  let extract: any;  let warning: string | undefined;  let currentModel = model;  let lastError: Error | null = null;   if (markdown === undefined) {    throw new Error("document.markdown is undefined -- this is unexpected");  }   try {    const prompt =      options.prompt !== undefined        ? `Transform the following content into structured JSON output based on the provided schema and this user request: ${options.prompt}. If schema is provided, strictly follow it.\n\n${markdown}`        : `Transform the following content into structured JSON output based on the provided schema if any.\n\n${markdown}`;     if (mode === "no-object") {      try {        const result = await generateText({          model: currentModel,          prompt: options.prompt + (markdown ? `\n\nData:${markdown}` : ""),          system: options.systemPrompt,          providerOptions: {            anthropic: {              thinking: { type: "enabled", budgetTokens: 12000 },            },            google: {              labels: {                teamId: metadata.teamId,                functionId: metadata.functionId ?? "unspecified",                extractId: metadata.extractId ?? "unspecified",                scrapeId: metadata.scrapeId ?? "unspecified",                deepResearchId: metadata.deepResearchId ?? "unspecified",                llmsTxtId: metadata.llmsTxtId ?? "unspecified",              },            },          },          experimental_telemetry: {            isEnabled: true,            functionId: metadata.functionId              ? metadata.functionId + "/generateText"              : "generateText",            metadata: {              teamId: metadata.teamId,              ...(metadata.extractId                ? {                    langfuseTraceId: "extract:" + metadata.extractId,                    extractId: metadata.extractId,                  }                : {}),              ...(metadata.scrapeId                ? {                    langfuseTraceId: "scrape:" + metadata.scrapeId,                    scrapeId: metadata.scrapeId,                  }                : {}),              ...(metadata.deepResearchId                ? {                    langfuseTraceId: "deepResearch:" + metadata.deepResearchId,                    deepResearchId: metadata.deepResearchId,                  }                : {}),              ...(metadata.llmsTxtId                ? {                    langfuseTraceId: "llmsTxt:" + metadata.llmsTxtId,                    llmsTxtId: metadata.llmsTxtId,                  }                : {}),            },          },        });         costTrackingOptions.costTracking.addCall({          type: "other",          metadata: {            ...costTrackingOptions.metadata,            gcDetails: "no-object",          },          model: currentModel.modelId,          cost: calculateCost(            currentModel.modelId,            result.usage?.promptTokens ?? 0,            result.usage?.completionTokens ?? 0,          ),          tokens: {            input: result.usage?.promptTokens ?? 0,            output: result.usage?.completionTokens ?? 0,          },        });         extract = result.text;         return {          extract,          warning,          numTokens: result.usage?.promptTokens ?? 0,          totalUsage: {            promptTokens: result.usage?.promptTokens ?? 0,            completionTokens: result.usage?.completionTokens ?? 0,            totalTokens:              result.usage?.promptTokens ??              0 + (result.usage?.completionTokens ?? 0),          },          model: currentModel.modelId,        };      } catch (error) {        lastError = error as Error;        if (          error.message?.includes("Quota exceeded") ||          error.message?.includes("You exceeded your current quota") ||          error.message?.includes("rate limit")        ) {          logger.warn("Quota exceeded, retrying with fallback model", {            error: lastError.message,          });          currentModel = retryModel;          try {            const result = await generateText({              model: currentModel,              prompt: options.prompt + (markdown ? `\n\nData:${markdown}` : ""),              system: options.systemPrompt,              providerOptions: {                anthropic: {                  thinking: { type: "enabled", budgetTokens: 12000 },                },                google: {                  labels: {                    teamId: metadata.teamId,                    functionId: metadata.functionId ?? "unspecified",                    extractId: metadata.extractId ?? "unspecified",                    scrapeId: metadata.scrapeId ?? "unspecified",                    deepResearchId: metadata.deepResearchId ?? "unspecified",                    llmsTxtId: metadata.llmsTxtId ?? "unspecified",                  },                },              },              experimental_telemetry: {                isEnabled: true,                functionId: metadata.functionId                  ? metadata.functionId + "/generateText"                  : "generateText",                metadata: {                  teamId: metadata.teamId,                  ...(metadata.extractId                    ? {                        langfuseTraceId: "extract:" + metadata.extractId,                        extractId: metadata.extractId,                      }                    : {}),                  ...(metadata.scrapeId                    ? {                        langfuseTraceId: "scrape:" + metadata.scrapeId,                        scrapeId: metadata.scrapeId,                      }                    : {}),                  ...(metadata.deepResearchId                    ? {                        langfuseTraceId:                          "deepResearch:" + metadata.deepResearchId,                        deepResearchId: metadata.deepResearchId,                      }                    : {}),                  ...(metadata.llmsTxtId                    ? {                        langfuseTraceId: "llmsTxt:" + metadata.llmsTxtId,                        llmsTxtId: metadata.llmsTxtId,                      }                    : {}),                },              },            });             extract = result.text;             costTrackingOptions.costTracking.addCall({              type: "other",              metadata: {                ...costTrackingOptions.metadata,                gcDetails: "no-object fallback",              },              model: currentModel.modelId,              cost: calculateCost(                currentModel.modelId,                result.usage?.promptTokens ?? 0,                result.usage?.completionTokens ?? 0,              ),              tokens: {                input: result.usage?.promptTokens ?? 0,                output: result.usage?.completionTokens ?? 0,              },            });             return {              extract,              warning,              numTokens: result.usage?.promptTokens ?? 0,              totalUsage: {                promptTokens: result.usage?.promptTokens ?? 0,                completionTokens: result.usage?.completionTokens ?? 0,                totalTokens:                  result.usage?.promptTokens ??                  0 + (result.usage?.completionTokens ?? 0),              },              model: currentModel.modelId,            };          } catch (retryError) {            lastError = retryError as Error;            logger.error("Failed with fallback model", {              originalError: lastError.message,              model: currentModel.modelId,            });            throw lastError;          }        }        throw lastError;      }    }     let schema = options.schema;    // Normalize the bad json schema users write (mogery)    if (schema && !(schema instanceof z.ZodType)) {      // let schema = options.schema;      if (schema) {        schema = removeDefaultProperty(schema);      }       if (schema && schema.type === "array") {        schema = {          type: "object",          properties: {            items: options.schema,          },          required: ["items"],

Lines 1–500 of 1303.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t