Accept dimensions & user for embedding calls

edgedb · Dec 10, 2024 · dfd28ca · dfd28ca
1 parent 48d11e8
commit dfd28ca
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 24 deletions.
diff --git a/packages/ai/src/core.ts b/packages/ai/src/core.ts
@@ -11,6 +11,7 @@ import {
   type QueryContext,
   type StreamingMessage,
   type RagRequest,
+  type EmbeddingRequest,
   isPromptRequest,
 } from "./types.js";
 import { getHTTPSCRAMAuth } from "edgedb/dist/httpScram.js";
@@ -209,7 +210,7 @@ export class EdgeDBAI {
     };
   }
 
-  async generateEmbeddings(inputs: string[], model: string): Promise<number[]> {
+  async generateEmbeddings(request: EmbeddingRequest): Promise<number[]> {
     const response = await (
       await this.authenticatedFetch
     )("embeddings", {
@@ -218,8 +219,8 @@ export class EdgeDBAI {
         "Content-Type": "application/json",
       },
       body: JSON.stringify({
-        model,
-        input: inputs,
+        ...request,
+        input: request.inputs,
       }),
     });
 

diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts
@@ -49,18 +49,15 @@ export interface QueryContext {
   max_object_count?: number;
 }
 
-interface RagRequestBase {
-  stream?: boolean;
+export interface RagRequestPrompt {
+  prompt: string;
   [key: string]: unknown;
 }
 
-export type RagRequestPrompt = RagRequestBase & {
-  prompt: string;
-};
-
-export type RagRequestMessages = RagRequestBase & {
+export interface RagRequestMessages {
   messages: EdgeDBMessage[];
-};
+  [key: string]: unknown;
+}
 
 export type RagRequest = RagRequestPrompt | RagRequestMessages;
 
@@ -153,3 +150,10 @@ export type StreamingMessage =
   | MessageDelta
   | MessageStop
   | MessageError;
+
+export interface EmbeddingRequest {
+  inputs: string[];
+  model: string;
+  dimensions?: number;
+  user?: string;
+}
diff --git a/packages/vercel-ai-provider/src/edgedb-chat-language-model.ts b/packages/vercel-ai-provider/src/edgedb-chat-language-model.ts
@@ -9,14 +9,14 @@ import type {
 } from "@ai-sdk/provider";
 import {
   type ParseResult,
+  type FetchFunction,
   createEventSourceResponseHandler,
   createJsonResponseHandler,
   postJsonToApi,
   generateId,
   combineHeaders,
 } from "@ai-sdk/provider-utils";
 import {
-  type EdgeDBChatConfig,
   type EdgeDBChatModelId,
   type EdgeDBChatSettings,
   type EdgeDBMessage,
@@ -31,6 +31,13 @@ import {
 import { convertToEdgeDBMessages } from "./convert-to-edgedb-messages";
 import { prepareTools } from "./edgedb-prepare-tools";
 
+export interface EdgeDBChatConfig {
+  provider: string;
+  fetch: FetchFunction;
+  baseURL: string | null;
+  headers: () => Record<string, string | undefined>;
+}
+
 export class EdgeDBChatLanguageModel implements LanguageModelV1 {
   readonly specificationVersion = "v1";
   readonly defaultObjectGenerationMode = "json";

diff --git a/packages/vercel-ai-provider/src/edgedb-chat-settings.ts b/packages/vercel-ai-provider/src/edgedb-chat-settings.ts
@@ -1,5 +1,3 @@
-import type { FetchFunction } from "@ai-sdk/provider-utils";
-
 export type OpenAIModelId =
   | "gpt-4o"
   | "gpt-4o-mini"
@@ -67,13 +65,6 @@ export interface QueryContext {
   max_object_count?: number;
 }
 
-export interface EdgeDBChatConfig {
-  provider: string;
-  fetch: FetchFunction;
-  baseURL: string;
-  headers: () => Record<string, string | undefined>;
-}
-
 export interface EdgeDBChatSettings {
   context?: QueryContext;
   prompt?: Prompt;

diff --git a/packages/vercel-ai-provider/src/edgedb-embedding-model.ts b/packages/vercel-ai-provider/src/edgedb-embedding-model.ts
@@ -19,7 +19,7 @@ import { edgedbFailedResponseHandler } from "./edgedb-error";
 interface EdgeDBEmbeddingConfig {
   provider: string;
   fetch?: FetchFunction;
-  baseURL: string;
+  baseURL: string | null;
   headers: () => Record<string, string | undefined>;
 }
 
@@ -82,6 +82,7 @@ export class EdgeDBEmbeddingModel implements EmbeddingModelV1<string> {
         model: this.modelId,
         input: values,
         encoding_format: "float",
+        // OpenAI props
         dimensions: this.settings.dimensions,
         user: this.settings.user,
       },

diff --git a/packages/vercel-ai-provider/src/edgedb-provider.ts b/packages/vercel-ai-provider/src/edgedb-provider.ts
@@ -41,7 +41,6 @@ export interface EdgeDBProvider extends ProviderV1 {
 export interface EdgeDBProviderSettings {
   /**
 Use a different URL prefix for API calls, e.g. to use proxy servers.
-The default prefix is `https://api.mistral.ai/v1`.
    */
   baseURL?: string;
 
@@ -56,8 +55,9 @@ export async function createEdgeDB(
   options: EdgeDBProviderSettings = {},
 ): Promise<EdgeDBProvider> {
   const connectConfig = await client.resolveConnectionParams();
-  const baseURL = withoutTrailingSlash(options.baseURL) ?? "";
+  const baseURL = withoutTrailingSlash(options.baseURL) ?? null;
 
+  // In case we want to add more things to this in the future
   const getHeaders = () => ({
     ...options.headers,
   });