docs: add interactive chat for anthropic

.changeset/fifty-papayas-know.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add quantized parameter to HuggingFaceEmbedding

apps/docs/docs/modules/embeddings/available_embeddings/huggingface.md

@@ -23,3 +23,15 @@ const results = await queryEngine.query({
   query,
 });
 ```
+
+Per default, `HuggingFaceEmbedding` is using the `Xenova/all-MiniLM-L6-v2` model. You can change the model by passing the `modelType` parameter to the constructor.
+If you're not using a quantized model, set the `quantized` parameter to `false`.
+
+For example, to use the not quantized `BAAI/bge-small-en-v1.5` model, you can use the following code:
+
+```
+const embedModel = new HuggingFaceEmbedding({
+  modelType: "BAAI/bge-small-en-v1.5",
+  quantized: false,
+});
+```

examples/anthropic/chat.ts

@@ -3,6 +3,7 @@ import { Anthropic } from "llamaindex";
 (async () => {
   const anthropic = new Anthropic({
     apiKey: process.env.ANTHROPIC_API_KEY,
+    model: "claude-3-opus",
   });
   const result = await anthropic.chat({
     messages: [

examples/anthropic/chat_interactive.ts

@@ -0,0 +1,34 @@
+import { Anthropic, SimpleChatEngine, SimpleChatHistory } from "llamaindex";
+import { stdin as input, stdout as output } from "node:process";
+import readline from "node:readline/promises";
+
+(async () => {
+  const llm = new Anthropic({
+    apiKey: process.env.ANTHROPIC_API_KEY,
+    model: "claude-3-opus",
+  });
+  // chatHistory will store all the messages in the conversation
+  const chatHistory = new SimpleChatHistory({
+    messages: [
+      {
+        content: "You want to talk in rhymes.",
+        role: "system",
+      },
+    ],
+  });
+  const chatEngine = new SimpleChatEngine({
+    llm,
+    chatHistory,
+  });
+  const rl = readline.createInterface({ input, output });
+
+  while (true) {
+    const query = await rl.question("User: ");
+    process.stdout.write("Assistant: ");
+    const stream = await chatEngine.chat({ message: query, stream: true });
+    for await (const chunk of stream) {
+      process.stdout.write(chunk.response);
+    }
+    process.stdout.write("\n");
+  }
+})();

examples/anthropic/stream.ts

@@ -0,0 +1,23 @@
+import { Anthropic } from "llamaindex";
+
+(async () => {
+  const anthropic = new Anthropic({
+    apiKey: process.env.ANTHROPIC_API_KEY,
+    model: "claude-instant-1.2",
+  });
+  const stream = await anthropic.chat({
+    messages: [
+      { content: "You want to talk in rhymes.", role: "system" },
+      {
+        content:
+          "How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
+        role: "user",
+      },
+    ],
+    stream: true,
+  });
+
+  for await (const chunk of stream) {
+    process.stdout.write(chunk.delta);
+  }
+})();

examples/nlsqlTableQueryEngine.ts

@@ -1,81 +0,0 @@
-import knex from "knex";
-import {
-  NLSQLQueryEngine,
-  OpenAI,
-  SQLDatabase,
-  serviceContextFromDefaults,
-} from "llamaindex";
-
-async function main() {
-  const llm = new OpenAI({
-    model: "gpt-4",
-  });
-
-  const engine = knex({
-    client: "sqlite3", // or 'better-sqlite3'
-    connection: {
-      filename: ":memory:",
-    },
-  });
-
-  const db = new SQLDatabase({
-    engine,
-    schema: undefined,
-    metadata: {},
-    ignoreTables: undefined,
-    includeTables: ["test_table_1"],
-    sampleRowsInTableInfo: 3,
-    indexesInTableInfo: true,
-    customTableInfo: undefined,
-    maxStringLength: 100,
-  });
-
-  const tableName = "test_table_1";
-
-  await engine.schema.createTable(tableName, async (table) => {
-    table.increments("id");
-    table.string("comment");
-    table.string("author");
-
-    await db.insertIntoTable(tableName, {
-      comment: "this is a test1",
-      author: "emanuel",
-    });
-    await db.insertIntoTable(tableName, {
-      comment: "this is a test2",
-      author: "alex",
-    });
-    await db.insertIntoTable(tableName, {
-      comment: "this is a test3",
-      author: "yi",
-    });
-    await db.insertIntoTable(tableName, {
-      comment: "this is a test4",
-      author: "alex",
-    });
-
-    const ctx = serviceContextFromDefaults({
-      llm,
-    });
-
-    const engine = new NLSQLQueryEngine({
-      sqlDatabase: db,
-      tables: ["test_table_1"],
-      verbose: true,
-      serviceContext: ctx,
-      synthesizeResponse: true,
-    });
-
-    const response = await engine.query({
-      query: "What's the comment from author yi and emanuel?",
-    });
-
-    console.log({ response });
-
-    process.exit(0);
-  });
-}
-
-main().then(() => [
-  // process.exit(0)
-]);

examples/package.json

@@ -9,10 +9,8 @@
     "chromadb": "^1.8.1",
     "commander": "^11.1.0",
     "dotenv": "^16.4.1",
-    "knex": "^3.1.0",
     "llamaindex": "latest",
-    "mongodb": "^6.2.0",
-    "sqlite3": "^5.1.7"
+    "mongodb": "^6.2.0"
   },
   "devDependencies": {
     "@types/node": "^18.19.10",

examples/test.ts

@@ -1,46 +0,0 @@
-import knex from "knex";
-import { SQLDatabase } from "llamaindex";
-
-async function main() {
-  const engine = knex({
-    client: "sqlite3", // or 'better-sqlite3'
-    connection: {
-      filename: ":memory:",
-    },
-  });
-
-  const db = new SQLDatabase({
-    engine,
-    schema: undefined,
-    metadata: {},
-    ignoreTables: undefined,
-    includeTables: ["test_table"],
-    sampleRowsInTableInfo: 3,
-    indexesInTableInfo: true,
-    customTableInfo: undefined,
-    maxStringLength: 100,
-  });
-
-  const tableName = "test_table";
-
-  await engine.schema.createTable(tableName, () => {});
-
-  await db.insertIntoTable(tableName, {
-    name: "test1",
-    comment: "this is a test1",
-  });
-  await db.insertIntoTable(tableName, {
-    name: "test2",
-    comment: "this is a test2",
-  });
-  await db.insertIntoTable(tableName, {
-    name: "test3",
-    comment: "this is a test3",
-  });
-  await db.insertIntoTable(tableName, {
-    name: "test4",
-    comment: "this is a test4",
-  });
-}
-
-main();

packages/core/package.json

@@ -4,7 +4,7 @@
   "license": "MIT",
   "type": "module",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.13.0",
+    "@anthropic-ai/sdk": "^0.15.0",
     "@aws-crypto/sha256-js": "^5.2.0",
     "@datastax/astra-db-ts": "^0.1.4",
     "@llamaindex/cloud": "0.0.4",
@@ -23,7 +23,6 @@
     "cohere-ai": "^7.7.5",
     "file-type": "^18.7.0",
     "js-tiktoken": "^1.0.10",
-    "knex": "^3.1.0",
     "lodash": "^4.17.21",
     "mammoth": "^1.6.0",
     "md-utils-ts": "^2.0.0",
@@ -95,7 +94,7 @@
     "build": "rm -rf ./dist && pnpm run build:esm && pnpm run build:cjs && pnpm run build:type",
     "build:esm": "swc src -d dist --strip-leading-paths --config-file .swcrc",
     "build:cjs": "swc src -d dist/cjs --strip-leading-paths --config-file .cjs.swcrc",
-    "build:type": "rm -f .tsbuildinfo && tsc -b --diagnostics",
+    "build:type": "pnpm run -w type-check",
     "copy": "cp -r ../../README.md ../../LICENSE .",
     "postbuild": "pnpm run copy && node -e \"require('fs').writeFileSync('./dist/cjs/package.json', JSON.stringify({ type: 'commonjs' }))\"",
     "circular-check": "madge -c ./src/index.ts",

packages/core/src/embeddings/HuggingFaceEmbedding.ts

@@ -20,6 +20,7 @@ export enum HuggingFaceEmbeddingModelType {
  */
 export class HuggingFaceEmbedding extends BaseEmbedding {
   modelType: string = HuggingFaceEmbeddingModelType.XENOVA_ALL_MINILM_L6_V2;
+  quantized: boolean = true;
 
   private extractor: any;
 
@@ -31,7 +32,9 @@ export class HuggingFaceEmbedding extends BaseEmbedding {
   async getExtractor() {
     if (!this.extractor) {
       const { pipeline } = await import("@xenova/transformers");
-      this.extractor = await pipeline("feature-extraction", this.modelType);
+      this.extractor = await pipeline("feature-extraction", this.modelType, {
+        quantized: this.quantized,
+      });
     }
     return this.extractor;
   }

packages/core/src/engines/query/index.ts

@@ -1,4 +1,3 @@
 export * from "./RetrieverQueryEngine.js";
 export * from "./RouterQueryEngine.js";
 export * from "./SubQuestionQueryEngine.js";
-export * from "./sql/index.js";

packages/core/src/engines/query/sql/NLSQLQueryEngine.ts

@@ -1,59 +0,0 @@
-import {
-  NLSQLRetriever,
-  type SQLDatabase,
-  type ServiceContext,
-} from "../../../index.js";
-import type { TextToSQLPrompt } from "../../../retriever/sql/prompts.js";
-import { BaseSQLTableQueryEngine } from "./types.js";
-
-type NLSQLQueryEngineParams = {
-  sqlDatabase: SQLDatabase;
-  textToSQLPrompt?: TextToSQLPrompt;
-  contextQueryKwargs?: any | null;
-  synthesizeResponse?: boolean;
-  responseSynthesisPrompt?: any | null;
-  tables?: any[] | string[] | undefined;
-  serviceContext?: ServiceContext | undefined;
-  contextStrPrefix?: string | undefined;
-  sqlOnly?: boolean;
-  verbose?: boolean;
-};
-
-export class NLSQLQueryEngine extends BaseSQLTableQueryEngine {
-  _sqlRetriever: NLSQLRetriever;
-
-  constructor({
-    sqlDatabase,
-    textToSQLPrompt,
-    contextQueryKwargs = null,
-    synthesizeResponse = true,
-    responseSynthesisPrompt = null,
-    tables,
-    serviceContext,
-    contextStrPrefix,
-    sqlOnly = false,
-    verbose = false,
-  }: NLSQLQueryEngineParams) {
-    super({
-      synthesizeResponse,
-      responseSynthesisPrompt,
-      serviceContext,
-      verbose,
-    });
-
-    this._sqlRetriever = new NLSQLRetriever({
-      sqlDatabase,
-      textToSQLPrompt,
-      contextQueryKwargs,
-      tables,
-      contextStrPrefix,
-      serviceContext,
-      sqlOnly,
-      verbose,
-    });
-  }
-
-  get sqlRetriever(): NLSQLRetriever {
-    return this._sqlRetriever;
-  }
-}

packages/core/src/engines/query/sql/index.ts

@@ -1 +0,0 @@
-export * from "./NLSQLQueryEngine.js";

packages/core/src/engines/query/sql/prompts.ts

@@ -1,17 +0,0 @@
-export const defaultResponseSynthesisPrompt = ({
-  query,
-  context,
-  sqlQuery,
-}: {
-  query?: string;
-  context?: string;
-  sqlQuery: string;
-}) => `
-Given an input question, synthesize a response from the query results.
-Query: ${query}
-SQL: ${sqlQuery}
-SQL Response: ${context}
-Response: 
-`;
-
-export type ResponseSynthesisPrompt = typeof defaultResponseSynthesisPrompt;

packages/core/src/engines/query/sql/types.ts

@@ -1,117 +0,0 @@
-import { Response } from "../../../Response.js";
-import {
-  serviceContextFromDefaults,
-  type ServiceContext,
-} from "../../../ServiceContext.js";
-import {
-  CompactAndRefine,
-  MetadataMode,
-  ResponseSynthesizer,
-} from "../../../index.js";
-import type { SQLRetriever } from "../../../retriever/sql/types.js";
-import type {
-  BaseQueryEngine,
-  QueryEngineParamsNonStreaming,
-  QueryEngineParamsStreaming,
-} from "../../../types.js";
-import {
-  defaultResponseSynthesisPrompt,
-  type ResponseSynthesisPrompt,
-} from "./prompts.js";
-
-export abstract class BaseSQLTableQueryEngine implements BaseQueryEngine {
-  synthesizeResponse: boolean;
-  responseSynthesisPrompt: ResponseSynthesisPrompt;
-  serviceContext: ServiceContext;
-  verbose: boolean;
-
-  constructor(init: {
-    synthesizeResponse?: boolean;
-    responseSynthesisPrompt?: ResponseSynthesisPrompt;
-    serviceContext?: ServiceContext;
-    verbose?: boolean;
-  }) {
-    this.synthesizeResponse = init.synthesizeResponse ?? true;
-    this.responseSynthesisPrompt =
-      init.responseSynthesisPrompt || defaultResponseSynthesisPrompt;
-    this.serviceContext = init.serviceContext || serviceContextFromDefaults({});
-    this.verbose = init.verbose || false;
-  }
-
-  getPrompts(): {
-    responseSynthesisPrompt: ResponseSynthesisPrompt;
-  } {
-    return { responseSynthesisPrompt: this.responseSynthesisPrompt };
-  }
-
-  updatePrompts(prompts: {
-    responseSynthesisPrompt: ResponseSynthesisPrompt;
-  }): void {
-    if ("responseSynthesisPrompt" in prompts) {
-      this.responseSynthesisPrompt = prompts.responseSynthesisPrompt;
-    }
-  }
-
-  getPromptModules(): {
-    sqlRetriever: SQLRetriever;
-  } {
-    return { sqlRetriever: this.sqlRetriever };
-  }
-
-  abstract get sqlRetriever(): SQLRetriever;
-
-  query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
-  query(params: QueryEngineParamsNonStreaming): Promise<Response>;
-  async query(
-    params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
-  ): Promise<Response | AsyncIterable<Response>> {
-    const { query, stream } = params;
-
-    if (stream) {
-      throw new Error("Streaming is not supported");
-    }
-
-    const [retrievedNodes, metadata] =
-      await this.sqlRetriever.retrieveWithMetadata({
-        queryStr: query,
-      });
-
-    const sqlQueryStr = metadata.sqlQuery;
-
-    console.log(`> SQL query: ${sqlQueryStr}`); // TODO: Remove
-
-    console.log(`> Sythesize Response ${this.synthesizeResponse}`);
-
-    if (this.synthesizeResponse) {
-      const responseBuilder = new CompactAndRefine(
-        this.serviceContext,
-        ({ query, context }) =>
-          this.responseSynthesisPrompt({
-            query,
-            context,
-            sqlQuery: sqlQueryStr,
-          }),
-      );
-
-      const responseSynthesizer = new ResponseSynthesizer({
-        serviceContext: this.serviceContext,
-        responseBuilder,
-      });
-
-      const response = await responseSynthesizer.synthesize({
-        query,
-        nodesWithScore: retrievedNodes,
-      });
-
-      response.metadata.sqlQuery = sqlQueryStr;
-
-      return response;
-    }
-
-    const responseStr = retrievedNodes
-      .map((node) => node.node.getContent(MetadataMode.ALL))
-      .join("\n");
-
-    return new Response(responseStr, []);
-  }
-}

packages/core/src/index.ts

@@ -26,9 +26,7 @@ export * from "./objects/index.js";
 export * from "./postprocessors/index.js";
 export * from "./prompts/index.js";
 export * from "./readers/index.js";
-export * from "./retriever/index.js";
 export * from "./selectors/index.js";
 export * from "./storage/index.js";
 export * from "./synthesizers/index.js";
 export * from "./tools/index.js";
-export * from "./utilities/index.js";

packages/core/src/llm/LLM.ts

@@ -1,7 +1,6 @@
 import type OpenAILLM from "openai";
 import type { ClientOptions as OpenAIClientOptions } from "openai";
 import type {
-  AnthropicStreamToken,
   CallbackManager,
   Event,
   EventType,
@@ -13,11 +12,7 @@ import type { ChatCompletionMessageParam } from "openai/resources/index.js";
 import type { LLMOptions } from "portkey-ai";
 import { Tokenizers, globalsHelper } from "../GlobalsHelper.js";
 import type { AnthropicSession } from "./anthropic.js";
-import {
-  ANTHROPIC_AI_PROMPT,
-  ANTHROPIC_HUMAN_PROMPT,
-  getAnthropicSession,
-} from "./anthropic.js";
+import { getAnthropicSession } from "./anthropic.js";
 import type { AzureOpenAIConfig } from "./azure.js";
 import {
   getAzureBaseUrl,
@@ -613,12 +608,30 @@ If a question does not make any sense, or is not factually coherent, explain why
   }
 }
 
-export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
-  // both models have 100k context window, see https://docs.anthropic.com/claude/reference/selecting-a-model
-  "claude-2": { contextWindow: 200000 },
-  "claude-instant-1": { contextWindow: 100000 },
+export const ALL_AVAILABLE_ANTHROPIC_LEGACY_MODELS = {
+  "claude-2.1": {
+    contextWindow: 200000,
+  },
+  "claude-instant-1.2": {
+    contextWindow: 100000,
+  },
 };
 
+export const ALL_AVAILABLE_V3_MODELS = {
+  "claude-3-opus": { contextWindow: 200000 },
+  "claude-3-sonnet": { contextWindow: 200000 },
+};
+
+export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
+  ...ALL_AVAILABLE_ANTHROPIC_LEGACY_MODELS,
+  ...ALL_AVAILABLE_V3_MODELS,
+};
+
+const AVAILABLE_ANTHROPIC_MODELS_WITHOUT_DATE: { [key: string]: string } = {
+  "claude-3-opus": "claude-3-opus-20240229",
+  "claude-3-sonnet": "claude-3-sonnet-20240229",
+} as { [key in keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS]: string };
+
 /**
  * Anthropic LLM implementation
  */
@@ -640,7 +653,7 @@ export class Anthropic extends BaseLLM {
 
   constructor(init?: Partial<Anthropic>) {
     super();
-    this.model = init?.model ?? "claude-2";
+    this.model = init?.model ?? "claude-3-opus";
     this.temperature = init?.temperature ?? 0.1;
     this.topP = init?.topP ?? 0.999; // Per Ben Mann
     this.maxTokens = init?.maxTokens ?? undefined;
@@ -674,21 +687,24 @@ export class Anthropic extends BaseLLM {
     };
   }
 
-  mapMessagesToPrompt(messages: ChatMessage[]) {
-    return (
-      messages.reduce((acc, message) => {
-        return (
-          acc +
-          `${
-            message.role === "system"
-              ? ""
-              : message.role === "assistant"
-                ? ANTHROPIC_AI_PROMPT + " "
-                : ANTHROPIC_HUMAN_PROMPT + " "
-          }${message.content.trim()}`
-        );
-      }, "") + ANTHROPIC_AI_PROMPT
-    );
+  getModelName = (model: string): string => {
+    if (Object.keys(AVAILABLE_ANTHROPIC_MODELS_WITHOUT_DATE).includes(model)) {
+      return AVAILABLE_ANTHROPIC_MODELS_WITHOUT_DATE[model];
+    }
+    return model;
+  };
+
+  formatMessages(messages: ChatMessage[]) {
+    return messages.map((message) => {
+      if (message.role !== "user" && message.role !== "assistant") {
+        throw new Error("Unsupported Anthropic role");
+      }
+
+      return {
+        content: message.content,
+        role: message.role,
+      };
+    });
   }
 
   chat(
@@ -698,49 +714,67 @@ export class Anthropic extends BaseLLM {
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream } = params;
+    let { messages } = params;
+
+    const { parentEvent, stream } = params;
+
+    let systemPrompt: string | null = null;
+
+    const systemMessages = messages.filter(
+      (message) => message.role === "system",
+    );
+
+    if (systemMessages.length > 0) {
+      systemPrompt = systemMessages
+        .map((message) => message.content)
+        .join("\n");
+      messages = messages.filter((message) => message.role !== "system");
+    }
+
     //Streaming
     if (stream) {
-      return this.streamChat(messages, parentEvent);
+      return this.streamChat(messages, parentEvent, systemPrompt);
     }
 
     //Non-streaming
-    const response = await this.session.anthropic.completions.create({
-      model: this.model,
-      prompt: this.mapMessagesToPrompt(messages),
-      max_tokens_to_sample: this.maxTokens ?? 100000,
+    const response = await this.session.anthropic.messages.create({
+      model: this.getModelName(this.model),
+      messages: this.formatMessages(messages),
+      max_tokens: this.maxTokens ?? 4096,
       temperature: this.temperature,
       top_p: this.topP,
+      ...(systemPrompt && { system: systemPrompt }),
     });
 
     return {
-      message: { content: response.completion.trimStart(), role: "assistant" },
-      //^ We're trimming the start because Anthropic often starts with a space in the response
-      // That space will be re-added when we generate the next prompt.
+      message: { content: response.content[0].text, role: "assistant" },
     };
   }
 
   protected async *streamChat(
     messages: ChatMessage[],
     parentEvent?: Event | undefined,
+    systemPrompt?: string | null,
   ): AsyncIterable<ChatResponseChunk> {
-    // AsyncIterable<AnthropicStreamToken>
-    const stream: AsyncIterable<AnthropicStreamToken> =
-      await this.session.anthropic.completions.create({
-        model: this.model,
-        prompt: this.mapMessagesToPrompt(messages),
-        max_tokens_to_sample: this.maxTokens ?? 100000,
-        temperature: this.temperature,
-        top_p: this.topP,
-        stream: true,
-      });
+    const stream = await this.session.anthropic.messages.create({
+      model: this.getModelName(this.model),
+      messages: this.formatMessages(messages),
+      max_tokens: this.maxTokens ?? 4096,
+      temperature: this.temperature,
+      top_p: this.topP,
+      stream: true,
+      ...(systemPrompt && { system: systemPrompt }),
+    });
 
     let idx_counter: number = 0;
     for await (const part of stream) {
-      //TODO: LLM Stream Callback, pending re-work.
+      const content =
+        part.type === "content_block_delta" ? part.delta.text : null;
+
+      if (typeof content !== "string") continue;
 
       idx_counter++;
-      yield { delta: part.completion };
+      yield { delta: content };
     }
     return;
   }

packages/core/src/retriever/index.ts

@@ -1 +0,0 @@
-export * from "./sql/index.js";

packages/core/src/retriever/sql/NLSQLRetriever.ts

@@ -1,259 +0,0 @@
-import { serviceContextFromDefaults } from "../../ServiceContext.js";
-import {
-  TextNode,
-  type BaseRetriever,
-  type CallbackManager,
-  type LLM,
-  type NodeWithScore,
-  type ObjectRetriever,
-  type SQLDatabase,
-  type ServiceContext,
-} from "../../index.js";
-import { QueryBundle } from "../../types.js";
-import { defaultTextToSQLPrompt, type TextToSQLPrompt } from "./prompts.js";
-import {
-  DefaultSQLParser,
-  SQLParserMode,
-  SQLRetriever,
-  type SQLTableSchema,
-} from "./types.js";
-
-export class NLSQLRetriever extends SQLRetriever implements BaseRetriever {
-  sqlDatabase: SQLDatabase;
-  sqlRetriever: SQLRetriever;
-  sqlParser: DefaultSQLParser;
-  textToSQLPrompt: TextToSQLPrompt;
-  contextQueryKwargs: Record<string, any> | undefined;
-  tables: any[] | string[] | undefined;
-  tableRetriever: ObjectRetriever | undefined;
-  contextStrPrefix: string | undefined;
-  sqlParserMode: SQLParserMode;
-  llm: LLM;
-  serviceContext: ServiceContext;
-  returnRaw: boolean;
-  handleSQLErrors: boolean;
-  sqlOnly: boolean;
-  callbackManager: CallbackManager | undefined;
-  verbose: boolean;
-  getTables: any;
-
-  constructor({
-    sqlDatabase,
-    textToSQLPrompt,
-    contextQueryKwargs,
-    tables,
-    tableRetriever,
-    contextStrPrefix,
-    sqlParserMode,
-    llm,
-    serviceContext,
-    returnRaw,
-    handleSQLErrors,
-    sqlOnly,
-    callbackManager,
-    verbose,
-  }: {
-    sqlDatabase: SQLDatabase;
-    textToSQLPrompt?: TextToSQLPrompt;
-    contextQueryKwargs?: Record<string, any>;
-    tables?: any[] | string[];
-    tableRetriever?: ObjectRetriever;
-    contextStrPrefix?: string;
-    sqlParserMode?: SQLParserMode;
-    llm?: LLM;
-    serviceContext?: ServiceContext;
-    returnRaw?: boolean;
-    handleSQLErrors?: boolean;
-    sqlOnly?: boolean;
-    callbackManager?: CallbackManager;
-    verbose?: boolean;
-  }) {
-    super(sqlDatabase, returnRaw, callbackManager);
-
-    this.sqlRetriever = new SQLRetriever(sqlDatabase, returnRaw);
-    this.sqlDatabase = sqlDatabase;
-    this.getTables = this.loadGetTablesFn(
-      sqlDatabase,
-      tables,
-      contextQueryKwargs,
-      tableRetriever,
-    );
-    this.contextStrPrefix = contextStrPrefix;
-    this.serviceContext = serviceContext ?? serviceContextFromDefaults();
-    this.textToSQLPrompt = textToSQLPrompt ?? defaultTextToSQLPrompt;
-    this.sqlParserMode = sqlParserMode ?? SQLParserMode.DEFAULT;
-    this.sqlParser = this.loadSQLParser(
-      this.sqlParserMode,
-      this.serviceContext,
-    );
-    this.handleSQLErrors = handleSQLErrors ?? true;
-    this.sqlOnly = sqlOnly ?? false;
-    this.verbose = verbose ?? false;
-    this.returnRaw = returnRaw ?? false;
-    this.llm = llm ?? this.serviceContext.llm;
-  }
-
-  _getPrompts() {
-    return {
-      textToSQLPrompt: this.textToSQLPrompt,
-    };
-  }
-
-  _updatePrompts(prompts: Record<string, any>) {
-    if ("textToSQLPrompt" in prompts) {
-      this.textToSQLPrompt = prompts.textToSQLPrompt;
-    }
-  }
-
-  _getPromptModules() {
-    return {};
-  }
-
-  getServiceContext(): ServiceContext {
-    return this.serviceContext;
-  }
-
-  loadSQLParser(sqlParserMode: SQLParserMode, serviceContext: ServiceContext) {
-    if (sqlParserMode === SQLParserMode.DEFAULT) {
-      return new DefaultSQLParser();
-    } else {
-      throw new Error(`Unknown SQL parser mode: ${sqlParserMode}`);
-    }
-  }
-
-  loadGetTablesFn(
-    sqlDatabase: SQLDatabase,
-    tables: any[] | string[] | undefined,
-    contextQueryKwargs: Record<string, any> | undefined,
-    tableRetriever: ObjectRetriever | undefined,
-  ) {
-    contextQueryKwargs = contextQueryKwargs || {};
-
-    if (tableRetriever) {
-      return async (queryStr: string) =>
-        await tableRetriever.retrieve(queryStr);
-    } else {
-      let tableNames: SQLTableSchema[] | string[];
-
-      if (tables) {
-        tableNames = tables.map((t) => t);
-      } else {
-        tableNames = Array.from(sqlDatabase.usableTableNames);
-      }
-
-      const contextStrs: string[] = [];
-
-      const tableSchemas = tableNames.map((t, i) => {
-        if (typeof t === "string") {
-          return {
-            tableName: t,
-            ...(contextQueryKwargs
-              ? { contextStr: contextQueryKwargs[t] }
-              : {}),
-          };
-        }
-
-        return {
-          tableName: t.tableName,
-          ...(contextQueryKwargs
-            ? { contextStr: contextQueryKwargs[t.tableName] }
-            : {}),
-        };
-      });
-
-      return () => tableSchemas;
-    }
-  }
-
-  async retrieveWithMetadata(strOrQueryBundle: string | QueryBundle): Promise<
-    [
-      NodeWithScore[],
-      {
-        sqlQuery: string;
-      },
-    ]
-  > {
-    const queryBundle =
-      typeof strOrQueryBundle === "string"
-        ? { queryStr: strOrQueryBundle }
-        : strOrQueryBundle;
-
-    const tableDescStr = await this.getTableContext(queryBundle);
-
-    if (this.verbose) {
-      console.log(`> Table desc str: ${tableDescStr}`);
-    }
-
-    const response = await this.serviceContext?.llm?.complete({
-      prompt: this.textToSQLPrompt({
-        dialect: "sql",
-        schema: tableDescStr,
-        queryStr: queryBundle.queryStr,
-      }),
-    });
-
-    if (!response) {
-      throw new Error("No response from LLM");
-    }
-
-    const sqlQueryStr = this.sqlParser.parseResponseToSQL(
-      response?.text,
-      queryBundle,
-    );
-
-    if (this.verbose) {
-      console.log(`> Predicted SQL query: ${sqlQueryStr}`);
-    }
-
-    let retrievedNodes: NodeWithScore[];
-    let metadata: Record<string, unknown> = {};
-
-    if (this.sqlOnly) {
-      const sqlOnlyNode = new TextNode({ text: sqlQueryStr });
-      retrievedNodes = [{ node: sqlOnlyNode }];
-      metadata = {};
-    } else {
-      try {
-        const retrieverResponse = await this.sqlRetriever.retrieveWithMetadata({
-          queryStr: sqlQueryStr,
-        });
-
-        retrievedNodes = retrieverResponse[0];
-        metadata = retrieverResponse[1];
-      } catch (e) {
-        if (this.handleSQLErrors) {
-          const errNode = new TextNode({ text: `Error: ${e}` });
-          retrievedNodes = [{ node: errNode }];
-          metadata = {};
-        } else {
-          throw e;
-        }
-      }
-    }
-    return [retrievedNodes, { sqlQuery: sqlQueryStr, ...metadata }];
-  }
-
-  async retrieve(query: string): Promise<NodeWithScore[]> {
-    const [retrievedNodes] = await this.retrieveWithMetadata(query);
-    return retrievedNodes;
-  }
-
-  async getTableContext(queryBundle: QueryBundle) {
-    const tableSchemaObjs = this.getTables(queryBundle.queryStr);
-    const contextStrs = [];
-    if (this.contextStrPrefix) {
-      contextStrs.push(this.contextStrPrefix);
-    }
-    for (const tableSchemaObj of tableSchemaObjs) {
-      let tableInfo = await this.sqlDatabase.getSingleTableInfo(
-        tableSchemaObj.tableName,
-      );
-      if (tableSchemaObj.contextStr) {
-        const tableOptContext = `The table description is: ${tableSchemaObj.contextStr}`;
-        tableInfo += tableOptContext;
-      }
-      contextStrs.push(tableInfo);
-    }
-    return contextStrs.join("\n\n");
-  }
-}

packages/core/src/retriever/sql/index.ts

@@ -1 +0,0 @@
-export * from "./NLSQLRetriever.js";

packages/core/src/retriever/sql/prompts.ts

@@ -1,31 +0,0 @@
-export const defaultTextToSQLPrompt = ({
-  dialect,
-  schema,
-  queryStr,
-}: {
-  dialect: string;
-  schema: string;
-  queryStr: string;
-}) => `Given an input question, first create a syntactically correct ${dialect}
-query to run, then look at the results of the query and return the answer.
-You can order the results by a relevant column to return the most
-interesting examples in the database.
-Never query for all the columns from a specific table, only ask for a
-few relevant columns given the question.
-Pay attention to use only the column names that you can see in the schema
-description.
-Be careful to not query for columns that do not exist.
-Pay attention to which column is in which table.
-Also, qualify column names with the table name when needed.
-You are required to use the following format, each taking one line:
-Question: Question here
-SQLQuery: SQL Query to run
-SQLResult: Result of the SQLQuery
-Answer: Final answer here
-Only use tables listed below.
-${schema}
-Question: ${queryStr}
-SQLQuery:
-`;
-
-export type TextToSQLPrompt = typeof defaultTextToSQLPrompt;

packages/core/src/retriever/sql/types.ts

@@ -1,105 +0,0 @@
-import type { BaseRetriever } from "../../Retriever.js";
-import {
-  TextNode,
-  type CallbackManager,
-  type Event,
-  type NodeWithScore,
-  type SQLDatabase,
-  type ServiceContext,
-} from "../../index.js";
-import type { QueryBundle } from "../../types.js";
-
-export interface SQLTableSchema {
-  tableName: string;
-  contextStr: string;
-}
-
-export enum SQLParserMode {
-  DEFAULT = "default",
-  PGVECTOR = "pgvector",
-}
-
-// export type SQLParserMode = "default" | "pgvector";
-
-export interface BaseSQLParser {
-  parseResponseToSQL(response: string, queryBundle: QueryBundle): string;
-}
-
-export class DefaultSQLParser implements BaseSQLParser {
-  parseResponseToSQL(response: string, queryBundle: QueryBundle): string {
-    const sqlQueryStart = response.indexOf("SQLQuery:");
-    if (sqlQueryStart !== -1) {
-      response = response.slice(sqlQueryStart);
-      if (response.startsWith("SQLQuery:")) {
-        response = response.slice("SQLQuery:".length);
-      }
-    }
-    const sqlResultStart = response.indexOf("SQLResult:");
-    if (sqlResultStart !== -1) {
-      response = response.slice(0, sqlResultStart);
-    }
-    return response.trim().replace("```", "").trim();
-  }
-}
-
-export class SQLRetriever implements BaseRetriever {
-  sqlDatabase: SQLDatabase;
-  returnRaw: boolean;
-
-  constructor(
-    sqlDatabase: SQLDatabase,
-    returnRaw: boolean = true,
-    callbackManager: CallbackManager | null = null,
-    kwargs: any = {},
-  ) {
-    this.sqlDatabase = sqlDatabase;
-    this.returnRaw = returnRaw;
-  }
-
-  getServiceContext(): ServiceContext {
-    throw new Error("Method not implemented.");
-  }
-
-  _formatNodeResults(results: any[][], colKeys: string[]): NodeWithScore[] {
-    const nodes: NodeWithScore[] = [];
-    for (const result of results) {
-      const metadata = Object.fromEntries(
-        colKeys.map((key, i) => [key, result[i]]),
-      );
-      const textNode = new TextNode({
-        text: "",
-        metadata,
-      });
-      nodes.push({ node: textNode });
-    }
-    return nodes;
-  }
-
-  async retrieveWithMetadata(
-    strOrQueryBundle: QueryBundle,
-  ): Promise<[NodeWithScore[], any]> {
-    const [rawResponseStr, metadata] = await this.sqlDatabase.runSQL(
-      strOrQueryBundle.queryStr,
-    );
-
-    if (this.returnRaw) {
-      return [[{ node: new TextNode({ text: rawResponseStr }) }], metadata];
-    } else {
-      const results = metadata.result;
-      const colKeys = metadata.colKeys;
-      return [this._formatNodeResults(results, colKeys), metadata];
-    }
-  }
-
-  async retrieve(
-    query: string,
-    parentEvent: Event | undefined,
-    preFilters: unknown,
-  ): Promise<NodeWithScore[]> {
-    const retrievedNodes = await this.retrieveWithMetadata({
-      queryStr: query,
-    });
-
-    return retrievedNodes;
-  }
-}

packages/core/src/utilities/SQLWrapper.ts

@@ -1,126 +0,0 @@
-import knex from "knex";
-
-type SQLDatabaseParams = {
-  engine: knex.Knex;
-  schema: string | undefined;
-  metadata: any;
-  ignoreTables: string[] | undefined;
-  includeTables: string[] | undefined;
-  sampleRowsInTableInfo: number;
-  indexesInTableInfo: boolean;
-  customTableInfo: Record<string, any> | undefined;
-  maxStringLength: number;
-};
-
-export class SQLDatabase {
-  engine: knex.Knex;
-  schema: string | undefined;
-  metadata: any;
-  inspector: knex.Knex;
-  allTables: Set<string>;
-  includeTables: Set<string>;
-  ignoreTables: Set<string>;
-  usableTables: Set<string>;
-  sampleRowsInTableInfo: number;
-  indexesInTableInfo: boolean;
-  customTableInfo: Record<string, any> | undefined;
-  maxStringLength: number;
-
-  constructor({
-    engine,
-    schema,
-    metadata,
-    ignoreTables,
-    includeTables,
-    sampleRowsInTableInfo,
-    indexesInTableInfo,
-    customTableInfo,
-    maxStringLength,
-  }: SQLDatabaseParams) {
-    this.engine = engine;
-    this.schema = schema;
-    this.metadata = metadata;
-    this.inspector = engine;
-    this.allTables = new Set(["test_table_1"]);
-    this.includeTables = new Set(includeTables || []);
-    this.ignoreTables = new Set(ignoreTables || []);
-    this.usableTables = new Set();
-    this.sampleRowsInTableInfo = sampleRowsInTableInfo;
-    this.indexesInTableInfo = indexesInTableInfo;
-    this.customTableInfo = customTableInfo;
-    this.maxStringLength = maxStringLength;
-  }
-
-  get usableTableNames(): string[] {
-    if (this.includeTables.size > 0) {
-      return Array.from(this.includeTables);
-    }
-    return Array.from(this.allTables);
-  }
-
-  async getTableColumns(tableName: string) {
-    return await this.inspector(tableName).columnInfo();
-  }
-
-  async getSingleTableInfo(tableName: string) {
-    const columns = await this.getTableColumns(tableName);
-
-    const columnStr = Object.keys(columns)
-      .map((column) => {
-        return `${column} (${columns[column].type})`;
-      })
-      .join(", ");
-
-    return `Table '${tableName}' has columns: ${columnStr}.`;
-  }
-
-  insertIntoTable(tableName: string, data: Record<string, any>): Promise<void> {
-    return this.engine(tableName).insert(data);
-  }
-
-  truncateWord(content: any, length: number, suffix = "..."): string {
-    if (typeof content !== "string" || length <= 0) {
-      return content;
-    }
-
-    if (content.length <= length) {
-      return content;
-    }
-
-    return (
-      content
-        .slice(0, length - suffix.length - 1)
-        .split(" ")
-        .slice(0, -1)
-        .join(" ") + suffix
-    );
-  }
-
-  async runSQL(
-    command: string,
-  ): Promise<[string, { result: any[]; colKeys: string[] }]> {
-    return this.engine.raw(command).then((result: any) => {
-      if (result.length > 0) {
-        const truncatedResults = result.map((row: any) =>
-          this.truncateWord(row, this.maxStringLength),
-        );
-        return [
-          JSON.stringify(truncatedResults),
-          { result: truncatedResults, colKeys: Object.keys(result[0]) },
-        ];
-      }
-      return ["", { result: [], colKeys: [] }];
-    });
-  }
-
-  async getTableInfo(tableName: string): Promise<string> {
-    const columns = await this.getTableColumns(tableName);
-    const columnStr = Object.keys(columns)
-      .map((column: any) => {
-        const comment = column.COMMENT ? `'${column.COMMENT}'` : "";
-        return `${column.COLUMN_NAME} (${column.DATA_TYPE}): ${comment}`;
-      })
-      .join(", ");
-    return `Table '${tableName}' has columns: ${columnStr}.`;
-  }
-}

packages/core/src/utilities/index.ts

@@ -1 +0,0 @@
-export * from "./SQLWrapper.js";

packages/core/tsconfig.json

@@ -3,7 +3,7 @@
   "compilerOptions": {
     "rootDir": "./src",
     "outDir": "./dist/type",
-    "tsBuildInfoFile": ".tsbuildinfo",
+    "tsBuildInfoFile": "./dist/.tsbuildinfo",
     "emitDeclarationOnly": true,
     "module": "node16",
     "moduleResolution": "node16",

packages/env/tsconfig.json

@@ -3,7 +3,7 @@
   "compilerOptions": {
     "rootDir": "./src",
     "outDir": "./dist/type",
-    "tsBuildInfoFile": ".tsbuildinfo",
+    "tsBuildInfoFile": "./dist/.tsbuildinfo",
     "emitDeclarationOnly": true,
     "module": "node16",
     "moduleResolution": "node16",