Merge pull request #1 from singlestore-labs/improve_query_builder

Improve query builder

examples/estore/index.ts

@@ -101,29 +101,19 @@ async function main() {
     await Promise.all([db.table("users").insert(dataset.users), db.table("products").insert(dataset.products)]);
     console.log("Dataset inserted.");
 
-    console.log('Selecting user with name "Alice"...');
-    console.log(await db.table("users").select({ name: "Alice" }));
+    console.log('Finding user with name "Alice"...');
+    console.log(await db.table("users").find({ where: { name: "Alice" } }));
 
-    console.log("Selecting products priced under 300...");
+    console.log("Finding products priced under 300...");
     console.log(
-      await db.table("products").select({ price: { lte: 300 } }, { columns: ["id", "name", "description", "price"] }),
+      await db.table("products").find({ select: ["id", "name", "description", "price"], where: { price: { lte: 300 } } }),
     );
 
-    console.log('Selecting products with name "Smartphone" priced at 800...');
+    console.log('Finding products with name "Smartphone" priced at 800...');
     console.log(
       await db
         .table("products")
-        .select({ name: "Smartphone", price: { lte: 800 } }, { columns: ["id", "name", "description", "price"] }),
-    );
-
-    console.log('Selecting products named "Smartphone" or "Laptop" priced above 500...');
-    console.log(
-      await db
-        .table("products")
-        .select(
-          { and: [{ or: [{ name: "Smartphone" }, { name: "Laptop" }] }, { price: { gt: 500 } }] },
-          { columns: ["id", "name", "description", "price"] },
-        ),
+        .find({ select: ["id", "name", "description", "price"], where: { name: "Smartphone", price: { lte: 800 } } }),
     );
 
     console.log("Executing custom query to select all users...");
@@ -165,14 +155,16 @@ async function main() {
 
     console.log("Executing database methods...");
     console.log('Creating "users" table...');
-    await db.createTable<Database["tables"]["users"]>({
+    const newUsersTable = await db.createTable<Database["tables"]["users"]>({
       name: "users",
       columns: {
         id: { type: "bigint", autoIncrement: true, primaryKey: true },
         name: { type: "varchar(32)" },
       },
     });
 
+    await newUsersTable.insert(dataset.users);
+
     console.log('Using the "users" table...');
     db.table("users");
 
@@ -182,7 +174,7 @@ async function main() {
     console.log(
       await db
         .table("products")
-        .vectorSearch({ prompt: prompt_19_1, vectorColumn: "description_v" }, { columns: ["name", "description"], limit: 1 }),
+        .vectorSearch({ prompt: prompt_19_1, vectorColumn: "description_v" }, { select: ["name", "description"], limit: 1 }),
     );
 
     const prompt_19_2 = "What monitor do I have in my store?";
@@ -192,7 +184,7 @@ async function main() {
         .table("products")
         .createChatCompletion(
           { prompt: prompt_19_2, vectorColumn: "description_v" },
-          { columns: ["name", "description", "price"], limit: 1 },
+          { select: ["name", "description", "price"], limit: 1 },
         ),
     );
 

packages/ai/src/chat-completions/index.ts

@@ -36,40 +36,43 @@ export interface ChatCompletionMessage {
 /**
  * Parameters for creating a chat completion.
  *
- * @typeParam T - Indicates whether the completion should be streamed (boolean).
- * @typeParam U - An array of tools that can be used during the chat completion.
+ * @typeParam TStream - Indicates whether the completion should be streamed (boolean).
+ * @typeParam TChatCompletionTool - An array of tools that can be used during the chat completion.
  *
  * @property {string} [prompt] - The initial prompt to generate the chat completion.
  * @property {string} [model] - The model to use for generating the completion.
  * @property {string} [systemRole] - The role of the system in the conversation.
- * @property {T} [stream] - Whether the completion should be streamed.
+ * @property {TStream} [stream] - Whether the completion should be streamed.
  * @property {ChatCompletionMessage[]} [messages] - An array of previous messages in the conversation.
- * @property {U} [tools] - An array of tools that can be used during the completion.
+ * @property {TChatCompletionTool} [tools] - An array of tools that can be used during the completion.
  *
  * @property {Record<string, (tool: AnyChatCompletionTool, params: any) => Promise<void>>} [toolCallHandlers] - Optional handlers for when a tool is called during the completion.
  *
  * @property {Record<string, (tool: AnyChatCompletionTool, result: any, params: any) => Promise<void>>} [toolCallResultHandlers] - Optional handlers for when a tool returns a result during the completion.
  */
-export interface CreateChatCompletionParams<T extends boolean | undefined, U extends AnyChatCompletionTool[] | undefined> {
+export interface CreateChatCompletionParams<
+  TStream extends boolean | undefined,
+  TChatCompletionTool extends AnyChatCompletionTool[] | undefined,
+> {
   prompt?: string;
   model?: string;
   systemRole?: string;
-  stream?: T;
+  stream?: TStream;
   messages?: ChatCompletionMessage[];
-  tools?: U;
+  tools?: TChatCompletionTool;
 
-  toolCallHandlers?: U extends AnyChatCompletionTool[]
+  toolCallHandlers?: TChatCompletionTool extends AnyChatCompletionTool[]
     ? {
-        [K in U[number] as K["name"]]?: (
+        [K in TChatCompletionTool[number] as K["name"]]?: (
           tool: K,
           params: K["params"] extends z.AnyZodObject ? z.infer<K["params"]> : undefined,
         ) => Promise<void>;
       }
     : undefined;
 
-  toolCallResultHandlers?: U extends AnyChatCompletionTool[]
+  toolCallResultHandlers?: TChatCompletionTool extends AnyChatCompletionTool[]
     ? {
-        [K in U[number] as K["name"]]?: (
+        [K in TChatCompletionTool[number] as K["name"]]?: (
           tool: K,
           result: Awaited<ReturnType<K["call"]>>,
           params: K["params"] extends z.AnyZodObject ? z.infer<K["params"]> : undefined,
@@ -81,29 +84,31 @@ export interface CreateChatCompletionParams<T extends boolean | undefined, U ext
 /**
  * The result of creating a chat completion.
  *
- * @typeParam T - Indicates whether the result is a stream or a single completion.
+ * @typeParam TStream - Indicates whether the result is a stream or a single completion.
  *
  * @returns If T is true, returns a `ChatCompletionStream`, otherwise returns a single `ChatCompletion`.
  */
-export type CreateChatCompletionResult<T extends boolean | undefined> = T extends true ? ChatCompletionStream : ChatCompletion;
+export type CreateChatCompletionResult<TStream extends boolean | undefined> = TStream extends true
+  ? ChatCompletionStream
+  : ChatCompletion;
 
 /**
  * Abstract class representing a chat completions generator, capable of handling various tools.
  *
- * @typeParam T - An array of tools that can be used during the chat completion.
+ * @typeParam TChatCompletionTool - An array of tools that can be used during the chat completion.
  */
-export abstract class ChatCompletions<T extends AnyChatCompletionTool[] | undefined> {
+export abstract class ChatCompletions<TChatCompletionTool extends AnyChatCompletionTool[] | undefined> {
   /**
    * The tools that can be used during chat completion. Initialized to undefined.
    */
-  tools = undefined as T;
+  tools = undefined as TChatCompletionTool;
 
   /**
    * Initializes the tools to be used in chat completion.
    *
    * @param tools - An array of tools to be used in the chat completion.
    */
-  initTools(tools: T) {
+  initTools(tools: TChatCompletionTool) {
     this.tools = tools;
   }
 

packages/ai/src/chat-completions/openai.ts

@@ -29,23 +29,27 @@ interface _OpenAICreateChatCompletionParams
 /**
  * Parameters for creating an OpenAI chat completion.
  *
- * @typeParam T - Indicates whether the completion should be streamed (boolean).
- * @typeParam U - An array of tools that can be used during the chat completion.
+ * @typeParam TStream - Indicates whether the completion should be streamed (boolean).
+ * @typeParam TChatCompletionTool - An array of tools that can be used during the chat completion.
  *
  * @property {OpenAIChatCompletionModel} [model] - The model to use for generating the completion.
  */
-export interface OpenAICreateChatCompletionParams<T extends boolean | undefined, U extends AnyChatCompletionTool[] | undefined>
-  extends CreateChatCompletionParams<T, U>,
+export interface OpenAICreateChatCompletionParams<
+  TStream extends boolean | undefined,
+  TChatCompletionTool extends AnyChatCompletionTool[] | undefined,
+> extends CreateChatCompletionParams<TStream, TChatCompletionTool>,
     _OpenAICreateChatCompletionParams {
   model?: OpenAIChatCompletionModel;
 }
 
 /**
  * Class representing chat completions using OpenAI's API, with support for tools and streaming.
  *
- * @typeParam T - An array of tools that can be used during the chat completion.
+ * @typeParam TChatCompletionTool - An array of tools that can be used during the chat completion.
  */
-export class OpenAIChatCompletions<T extends AnyChatCompletionTool[] | undefined> extends ChatCompletions<T> {
+export class OpenAIChatCompletions<
+  TChatCompletionTool extends AnyChatCompletionTool[] | undefined,
+> extends ChatCompletions<TChatCompletionTool> {
   constructor(private _openai: OpenAI) {
     super();
   }
@@ -80,22 +84,27 @@ export class OpenAIChatCompletions<T extends AnyChatCompletionTool[] | undefined
   /**
    * Creates a chat completion using OpenAI's API, with optional support for streaming and tool integration.
    *
-   * @typeParam U - Indicates whether the completion should be streamed (boolean).
-   * @typeParam K - An array of tools that can be used during the chat completion.
+   * @typeParam TStream - Indicates whether the completion should be streamed (boolean).
+   * @typeParam TChatCompletionTool - An array of tools that can be used during the chat completion.
    *
-   * @param {OpenAICreateChatCompletionParams<U, MergeChatCompletionTools<T, K>>} params - Parameters for creating the chat completion, including prompt, system role, messages, tools, and handlers.
+   * @param {OpenAICreateChatCompletionParams<TStream, MergeChatCompletionTools<TChatCompletionTool, TChatCompletionTool>>} params - Parameters for creating the chat completion, including prompt, system role, messages, tools, and handlers.
    *
    * @returns A promise that resolves to either a single chat completion or a stream of chat completions, depending on the `stream` parameter.
    */
-  async create<U extends boolean | undefined = false, K extends AnyChatCompletionTool[] | undefined = undefined>({
+  async create<
+    TStream extends boolean | undefined = false,
+    TChatCompletionTool extends AnyChatCompletionTool[] | undefined = undefined,
+  >({
     prompt,
     systemRole,
     messages,
     tools,
     toolCallHandlers,
     toolCallResultHandlers,
     ...params
-  }: OpenAICreateChatCompletionParams<U, MergeChatCompletionTools<T, K>>): Promise<CreateChatCompletionResult<U>> {
+  }: OpenAICreateChatCompletionParams<TStream, MergeChatCompletionTools<TChatCompletionTool, TChatCompletionTool>>): Promise<
+    CreateChatCompletionResult<TStream>
+  > {
     let _messages: ChatCompletionMessage[] = [];
     if (systemRole) _messages.push({ role: "system", content: systemRole });
     if (messages?.length) _messages = [..._messages, ...messages];
@@ -197,10 +206,10 @@ export class OpenAIChatCompletions<T extends AnyChatCompletionTool[] | undefined
 
       if (message && "tool_calls" in message && message.tool_calls?.length) {
         const toolCallResults = await handleToolCalls(message.tool_calls);
-        return (await handleToolCallResults(toolCallResults, message)) as CreateChatCompletionResult<U>;
+        return (await handleToolCallResults(toolCallResults, message)) as CreateChatCompletionResult<TStream>;
       }
 
-      return { content: message?.content || "" } as CreateChatCompletionResult<U>;
+      return { content: message?.content || "" } as CreateChatCompletionResult<TStream>;
     }
 
     /**
@@ -243,6 +252,6 @@ export class OpenAIChatCompletions<T extends AnyChatCompletionTool[] | undefined
       }
     }
 
-    return handleStream(response) as CreateChatCompletionResult<U>;
+    return handleStream(response) as CreateChatCompletionResult<TStream>;
   }
 }

packages/ai/src/embeddings/openai.ts

@@ -49,6 +49,7 @@ export class OpenAIEmbeddings extends Embeddings {
    * @returns {OpenAIEmbeddingModel[]} - An array of model names supported by OpenAI.
    */
   getModels(): OpenAIEmbeddingModel[] {
+    // TODO: Replace with dynamic values
     return ["text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002"];
   }
 

packages/ai/src/index.ts

@@ -12,20 +12,24 @@ export { Embeddings, ChatCompletions, ChatCompletionTool };
 /**
  * Configuration object for initializing the `AI` class.
  *
- * @typeParam T - The type of `Embeddings` to be used, defaulting to `OpenAIEmbeddings`.
- * @typeParam U - An array of tools that can be used during chat completions, or undefined.
- * @typeParam K - The type of `ChatCompletions` to be used, defaulting to `OpenAIChatCompletions`.
+ * @typeParam TEmbeddings - The type of `Embeddings` to be used, defaulting to `OpenAIEmbeddings`.
+ * @typeParam TChatCompletionTool - An array of tools that can be used during chat completions, or undefined.
+ * @typeParam TChatCompletion - The type of `ChatCompletions` to be used, defaulting to `OpenAIChatCompletions`.
  *
  * @property {string} [openAIApiKey] - The API key for authenticating with the OpenAI API.
- * @property {T} [embeddings] - An instance of `Embeddings` used for generating embeddings. Defaults to `OpenAIEmbeddings`.
- * @property {K} [chatCompletions] - An instance of `ChatCompletions` used for generating chat completions. Defaults to `OpenAIChatCompletions`.
- * @property {U} [chatCompletionTools] - An optional array of tools that can be used during chat completions.
+ * @property {TEmbeddings} [embeddings] - An instance of `Embeddings` used for generating embeddings. Defaults to `OpenAIEmbeddings`.
+ * @property {TChatCompletion} [chatCompletions] - An instance of `ChatCompletions` used for generating chat completions. Defaults to `OpenAIChatCompletions`.
+ * @property {TChatCompletionTool} [chatCompletionTools] - An optional array of tools that can be used during chat completions.
  */
-export interface AIConfig<T extends Embeddings, U extends AnyChatCompletionTool[] | undefined, K extends ChatCompletions<U>> {
+export interface AIConfig<
+  TEmbeddings extends Embeddings,
+  TChatCompletionTool extends AnyChatCompletionTool[] | undefined,
+  TChatCompletion extends ChatCompletions<TChatCompletionTool>,
+> {
   openAIApiKey?: string;
-  embeddings?: T;
-  chatCompletions?: K;
-  chatCompletionTools?: U;
+  embeddings?: TEmbeddings;
+  chatCompletions?: TChatCompletion;
+  chatCompletionTools?: TChatCompletionTool;
 }
 
 /**
@@ -37,36 +41,32 @@ export type AnyAI = AI<Embeddings, AnyChatCompletionTool[] | undefined, ChatComp
 /**
  * Main class for handling AI operations, including embeddings and chat completions.
  *
- * @typeParam T - The type of `Embeddings` to be used, defaulting to `OpenAIEmbeddings`.
- * @typeParam U - An array of tools that can be used during chat completions, or undefined.
- * @typeParam K - The type of `ChatCompletions` to be used, defaulting to `OpenAIChatCompletions`.
+ * @typeParam TEmbeddings - The type of `Embeddings` to be used, defaulting to `OpenAIEmbeddings`.
+ * @typeParam TChatCompletionTool - An array of tools that can be used during chat completions, or undefined.
+ * @typeParam TChatCompletions - The type of `ChatCompletions` to be used, defaulting to `OpenAIChatCompletions`.
  *
- * @property {T} embeddings - An instance of `Embeddings` used for generating embeddings.
- * @property {K} chatCompletions - An instance of `ChatCompletions` used for generating chat completions.
+ * @property {TEmbeddings} embeddings - An instance of `Embeddings` used for generating embeddings.
+ * @property {TChatCompletions} chatCompletions - An instance of `ChatCompletions` used for generating chat completions.
  */
 export class AI<
-  T extends Embeddings = OpenAIEmbeddings,
-  U extends AnyChatCompletionTool[] | undefined = undefined,
-  K extends ChatCompletions<any> = OpenAIChatCompletions<U>,
+  TEmbeddings extends Embeddings = OpenAIEmbeddings,
+  TChatCompletionTool extends AnyChatCompletionTool[] | undefined = undefined,
+  TChatCompletions extends ChatCompletions<any> = OpenAIChatCompletions<TChatCompletionTool>,
 > {
-  embeddings: T;
-  chatCompletions: K;
+  embeddings: TEmbeddings;
+  chatCompletions: TChatCompletions;
 
   /**
    * Constructs a new `AI` instance.
    *
-   * @param {AIConfig<T, U, K>} config - The configuration object for initializing the `AI` instance.
+   * @param {AIConfig<TEmbeddings, TChatCompletionTool, TChatCompletions>} config - The configuration object for initializing the `AI` instance.
    */
-  constructor(config: AIConfig<T, U, K>) {
+  constructor(config: AIConfig<TEmbeddings, TChatCompletionTool, TChatCompletions>) {
     const openai = new OpenAI({ apiKey: config.openAIApiKey });
 
-    // Initialize embeddings, defaulting to OpenAIEmbeddings if not provided
-    this.embeddings = (config.embeddings ?? new OpenAIEmbeddings(openai)) as T;
-
-    // Initialize chat completions, defaulting to OpenAIChatCompletions if not provided
-    this.chatCompletions = (config.chatCompletions ?? new OpenAIChatCompletions(openai)) as K;
+    this.embeddings = (config.embeddings ?? new OpenAIEmbeddings(openai)) as TEmbeddings;
+    this.chatCompletions = (config.chatCompletions ?? new OpenAIChatCompletions(openai)) as TChatCompletions;
 
-    // If tools are provided, initialize them within the chat completions instance
     if (config.chatCompletionTools?.length) {
       this.chatCompletions.initTools(config.chatCompletionTools);
     }