OpenAI Adapters (continuedev#1859)

* embeddings support in openai-adapters

* rerank support in openai-adapters

* cohere embed/rerank

* continue-proxy reranker and embeddings providers

* test for openai-adapters

* embeddings provider and reranker for continue proxy

Author: sestinj

.vscode/launch.json

@@ -68,7 +68,10 @@
         "${workspaceFolder}/extensions/vscode/out/extension.js",
         "/Users/natesesti/.continue/config.ts"
       ],
-      "preLaunchTask": "vscode-extension:build"
+      "preLaunchTask": "vscode-extension:build",
+      "env": {
+        "CONTROL_PLANE_ENV": "local"
+      }
     },
     // Has to be run after starting the server (separately or using the compound configuration)
     {
@@ -112,5 +115,24 @@
       "console": "integratedTerminal",
       "internalConsoleOptions": "neverOpen"
     }
+    // {
+    //   "name": "[openai-adapters] Jest Test Debugger, Current Open File",
+    //   "type": "node",
+    //   "request": "launch",
+    //   "runtimeArgs": [
+    //     "--inspect-brk",
+    //     "${workspaceRoot}/packages/openai-adapters/node_modules/jest/bin/jest.js",
+    //     "--runInBand",
+    //     "--config",
+    //     "${workspaceRoot}/packages/openai-adapters/jest.config.mjs",
+    //     "${relativeFile}"
+    //   ],
+    //   "cwd": "${workspaceRoot}/packages/openai-adapters",
+    //   "console": "integratedTerminal",
+    //   "internalConsoleOptions": "neverOpen",
+    //   "env": {
+    //     "NODE_OPTIONS": "--experimental-vm-modules"
+    //   }
+    // }
   ]
 }

core/config/profile/doLoadConfig.ts

@@ -4,8 +4,10 @@ import {
   IdeSettings,
   SerializedContinueConfig,
 } from "../..";
+import { ContinueProxyReranker } from "../../context/rerankers/ContinueProxyReranker";
 import { ControlPlaneClient } from "../../control-plane/client";
 import { TeamAnalytics } from "../../control-plane/TeamAnalytics";
+import ContinueProxyEmbeddingsProvider from "../../indexing/embeddings/ContinueProxyEmbeddingsProvider";
 import ContinueProxy from "../../llm/llms/stubs/ContinueProxy";
 import { Telemetry } from "../../util/posthog";
 import { loadFullConfigNode } from "../load";
@@ -65,5 +67,16 @@ export default async function doLoadConfig(
     },
   );
 
+  if (newConfig.embeddingsProvider?.providerName === "continue-proxy") {
+    (
+      newConfig.embeddingsProvider as ContinueProxyEmbeddingsProvider
+    ).workOsAccessToken = workOsAccessToken;
+  }
+
+  if (newConfig.reranker?.name === "continue-proxy") {
+    (newConfig.reranker as ContinueProxyReranker).workOsAccessToken =
+      workOsAccessToken;
+  }
+
   return newConfig;
 }

core/context/rerankers/ContinueProxyReranker.ts

@@ -0,0 +1,46 @@
+import fetch from "node-fetch";
+import { CONTROL_PLANE_URL } from "../../control-plane/client.js";
+import { Chunk, Reranker } from "../../index.js";
+
+export class ContinueProxyReranker implements Reranker {
+  name = "continue-proxy";
+
+  private _workOsAccessToken: string | undefined = undefined;
+
+  get workOsAccessToken(): string | undefined {
+    return this._workOsAccessToken;
+  }
+
+  set workOsAccessToken(value: string | undefined) {
+    if (this._workOsAccessToken !== value) {
+      this._workOsAccessToken = value;
+      this.params.apiKey = value!;
+    }
+  }
+
+  constructor(
+    private readonly params: {
+      apiKey: string;
+      model?: string;
+    },
+  ) {}
+
+  async rerank(query: string, chunks: Chunk[]): Promise<number[]> {
+    const url = new URL("/model-proxy/v1/rerank", CONTROL_PLANE_URL);
+    const resp = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${this.params.apiKey}`,
+      },
+      body: JSON.stringify({
+        query,
+        documents: chunks.map((chunk) => chunk.content),
+        model: this.params.model,
+      }),
+    });
+    const data: any = await resp.json();
+    const results = data.data.sort((a: any, b: any) => a.index - b.index);
+    return results.map((result: any) => result.relevance_score);
+  }
+}

core/context/rerankers/index.ts

@@ -1,14 +1,16 @@
 import { RerankerName } from "../../index.js";
 import { CohereReranker } from "./cohere.js";
+import { ContinueProxyReranker } from "./ContinueProxyReranker.js";
 import { FreeTrialReranker } from "./freeTrial.js";
 import { LLMReranker } from "./llm.js";
+import { HuggingFaceTEIReranker } from "./tei.js";
 import { VoyageReranker } from "./voyage.js";
-import {HuggingFaceTEIReranker} from "./tei.js"
 
 export const AllRerankers: { [key in RerankerName]: any } = {
   cohere: CohereReranker,
   llm: LLMReranker,
   voyage: VoyageReranker,
   "free-trial": FreeTrialReranker,
-  "huggingface-tei": HuggingFaceTEIReranker
+  "huggingface-tei": HuggingFaceTEIReranker,
+  "continue-proxy": ContinueProxyReranker,
 };

core/control-plane/client.ts

@@ -18,9 +18,10 @@ export interface ControlPlaneWorkspace {
 
 export interface ControlPlaneModelDescription extends ModelDescription {}
 
-// export const CONTROL_PLANE_URL = "http://localhost:3001";
 export const CONTROL_PLANE_URL =
-  "https://control-plane-api-service-i3dqylpbqa-uc.a.run.app";
+  process.env.CONTROL_PLANE_ENV === "local"
+    ? "http://localhost:3001"
+    : "https://control-plane-api-service-i3dqylpbqa-uc.a.run.app";
 
 export class ControlPlaneClient {
   private static URL = CONTROL_PLANE_URL;

core/index.d.ts

@@ -739,7 +739,9 @@ export type EmbeddingsProviderName =
   | "openai"
   | "cohere"
   | "free-trial"
-  | "gemini";
+  | "gemini"
+  | "continue-proxy"
+  | "deepinfra";
 
 export interface EmbedOptions {
   apiBase?: string;
@@ -758,6 +760,7 @@ export interface EmbeddingsProviderDescription extends EmbedOptions {
 
 export interface EmbeddingsProvider {
   id: string;
+  providerName: EmbeddingsProviderName;
   maxChunkSize: number;
   embed(chunks: string[]): Promise<number[][]>;
 }
@@ -767,7 +770,8 @@ export type RerankerName =
   | "voyage"
   | "llm"
   | "free-trial"
-  | "huggingface-tei";
+  | "huggingface-tei"
+  | "continue-proxy";
 
 export interface RerankerDescription {
   name: RerankerName;

core/indexing/embeddings/BaseEmbeddingsProvider.ts

@@ -1,6 +1,7 @@
 import {
   EmbedOptions,
   EmbeddingsProvider,
+  EmbeddingsProviderName,
   FetchFunction,
 } from "../../index.js";
 
@@ -17,6 +18,11 @@ abstract class BaseEmbeddingsProvider implements IBaseEmbeddingsProvider {
   static maxBatchSize: IBaseEmbeddingsProvider["maxBatchSize"];
   static defaultOptions: IBaseEmbeddingsProvider["defaultOptions"];
 
+  static providerName: EmbeddingsProviderName;
+  get providerName(): EmbeddingsProviderName {
+    return (this.constructor as typeof BaseEmbeddingsProvider).providerName;
+  }
+
   options: IBaseEmbeddingsProvider["options"];
   fetch: IBaseEmbeddingsProvider["fetch"];
   id: IBaseEmbeddingsProvider["id"];
@@ -38,6 +44,8 @@ abstract class BaseEmbeddingsProvider implements IBaseEmbeddingsProvider {
       this.id = `${this.constructor.name}::${this.options.model}`;
     }
   }
+  defaultOptions?: EmbedOptions | undefined;
+  maxBatchSize?: number | undefined;
 
   abstract embed(chunks: string[]): Promise<number[][]>;
 

core/indexing/embeddings/CohereEmbeddingsProvider.ts

@@ -1,11 +1,13 @@
 import { Response } from "node-fetch";
-import { EmbedOptions } from "../../index.js";
+import { EmbeddingsProviderName, EmbedOptions } from "../../index.js";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
 
 class CohereEmbeddingsProvider extends BaseEmbeddingsProvider {
   static maxBatchSize = 96;
 
+  static providerName: EmbeddingsProviderName = "cohere";
+
   static defaultOptions: Partial<EmbedOptions> | undefined = {
     apiBase: "https://api.cohere.ai/v1/",
     model: "embed-english-v3.0",

core/indexing/embeddings/ContinueProxyEmbeddingsProvider.ts

@@ -0,0 +1,25 @@
+import { EmbeddingsProviderName, EmbedOptions } from "../..";
+import { CONTROL_PLANE_URL } from "../../control-plane/client";
+import OpenAIEmbeddingsProvider from "./OpenAIEmbeddingsProvider";
+
+class ContinueProxyEmbeddingsProvider extends OpenAIEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "continue-proxy";
+  static defaultOptions: Partial<EmbedOptions> | undefined = {
+    apiBase: new URL("/model-proxy/v1", CONTROL_PLANE_URL).toString(),
+  };
+
+  private _workOsAccessToken: string | undefined = undefined;
+
+  get workOsAccessToken(): string | undefined {
+    return this._workOsAccessToken;
+  }
+
+  set workOsAccessToken(value: string | undefined) {
+    if (this._workOsAccessToken !== value) {
+      this._workOsAccessToken = value;
+      this.options.apiKey = value;
+    }
+  }
+}
+
+export default ContinueProxyEmbeddingsProvider;

core/indexing/embeddings/DeepInfraEmbeddingsProvider.ts

@@ -1,8 +1,9 @@
-import { EmbedOptions } from "../../index.js";
+import { EmbeddingsProviderName, EmbedOptions } from "../../index.js";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
 
 class DeepInfraEmbeddingsProvider extends BaseEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "deepinfra";
   static defaultOptions: Partial<EmbedOptions> | undefined = {
     model: "sentence-transformers/all-MiniLM-L6-v2",
   };

core/indexing/embeddings/FreeTrialEmbeddingsProvider.ts

@@ -1,11 +1,16 @@
 import { Response } from "node-fetch";
 import { getHeaders } from "../../continueServer/stubs/headers.js";
 import { constants } from "../../deploy/constants.js";
-import { EmbedOptions, FetchFunction } from "../../index.js";
+import {
+  EmbeddingsProviderName,
+  EmbedOptions,
+  FetchFunction,
+} from "../../index.js";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
 
 class FreeTrialEmbeddingsProvider extends BaseEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "free-trial";
   static maxBatchSize = 128;
 
   static defaultOptions: Partial<EmbedOptions> | undefined = {

core/indexing/embeddings/GeminiEmbeddingsProvider.ts

@@ -1,17 +1,17 @@
-import { Response } from "node-fetch";
-import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
-import BaseEmbeddingsProvider, {
-  IBaseEmbeddingsProvider,
-} from "./BaseEmbeddingsProvider.js";
 import {
   EmbedContentRequest,
   EmbedContentResponse,
 } from "@google/generative-ai";
+import { Response } from "node-fetch";
+import { EmbeddingsProviderName } from "../../index.js";
+import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
+import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
 
 /**
  * [View the Gemini Text Embedding docs.](https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding-and-embedding)
  */
 class GeminiEmbeddingsProvider extends BaseEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "gemini";
   static maxBatchSize = 2048;
 
   static defaultOptions = {

core/indexing/embeddings/HuggingFaceTEIEmbeddingsProvider.ts

@@ -1,10 +1,11 @@
-import fetch, { Response } from "node-fetch";
-import { EmbedOptions, FetchFunction } from "../..";
+import { Response } from "node-fetch";
+import { EmbeddingsProviderName, EmbedOptions, FetchFunction } from "../..";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider";
 
 class HuggingFaceTEIEmbeddingsProvider extends BaseEmbeddingsProvider {
-  private maxBatchSize = 32;
+  static providerName: EmbeddingsProviderName = "huggingface-tei";
+  maxBatchSize = 32;
 
   static defaultOptions: Partial<EmbedOptions> | undefined = {
     apiBase: "http://localhost:8080",
@@ -17,7 +18,7 @@ class HuggingFaceTEIEmbeddingsProvider extends BaseEmbeddingsProvider {
     if (!this.options.apiBase?.endsWith("/")) {
       this.options.apiBase += "/";
     }
-    this.doInfoRequest().then(response => {
+    this.doInfoRequest().then((response) => {
       this.options.model = response.model_id;
       this.maxBatchSize = response.max_client_batch_size;
     });
@@ -26,7 +27,9 @@ class HuggingFaceTEIEmbeddingsProvider extends BaseEmbeddingsProvider {
   async embed(chunks: string[]) {
     const promises = [];
     for (let i = 0; i < chunks.length; i += this.maxBatchSize) {
-      promises.push(this.doEmbedRequest(chunks.slice(i, i + this.maxBatchSize)));
+      promises.push(
+        this.doEmbedRequest(chunks.slice(i, i + this.maxBatchSize)),
+      );
     }
     const results = await Promise.all(promises);
     return results.flat();
@@ -37,11 +40,11 @@ class HuggingFaceTEIEmbeddingsProvider extends BaseEmbeddingsProvider {
       this.fetch(new URL("embed", this.options.apiBase), {
         method: "POST",
         body: JSON.stringify({
-          inputs: batch
+          inputs: batch,
         }),
         headers: {
           "Content-Type": "application/json",
-        }
+        },
       }),
     );
     if (!resp.ok) {
@@ -75,9 +78,9 @@ class TEIEmbedError extends Error {
 }
 
 type TEIEmbedErrorResponse = {
-  error: string
-  error_type: string
-}
+  error: string;
+  error_type: string;
+};
 
 type TEIInfoResponse = {
   model_id: string;
@@ -86,7 +89,7 @@ type TEIInfoResponse = {
   model_type: {
     embedding: {
       pooling: string;
-    }
+    };
   };
   max_concurrent_requests: number;
   max_input_length: number;

core/indexing/embeddings/OllamaEmbeddingsProvider.ts

@@ -1,4 +1,8 @@
-import { EmbedOptions, FetchFunction } from "../../index.js";
+import {
+  EmbeddingsProviderName,
+  EmbedOptions,
+  FetchFunction,
+} from "../../index.js";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
 import BaseEmbeddingsProvider, {
   IBaseEmbeddingsProvider,
@@ -41,6 +45,7 @@ async function embedOne(
 }
 
 class OllamaEmbeddingsProvider extends BaseEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "ollama";
   static defaultOptions: IBaseEmbeddingsProvider["defaultOptions"] = {
     apiBase: "http://localhost:11434/",
     model: "nomic-embed-text",

core/indexing/embeddings/OpenAIEmbeddingsProvider.ts

@@ -1,9 +1,10 @@
 import { Response } from "node-fetch";
-import { EmbedOptions } from "../../index.js";
+import { EmbeddingsProviderName, EmbedOptions } from "../../index.js";
 import { withExponentialBackoff } from "../../util/withExponentialBackoff.js";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
 
 class OpenAIEmbeddingsProvider extends BaseEmbeddingsProvider {
+  static providerName: EmbeddingsProviderName = "openai";
   // https://platform.openai.com/docs/api-reference/embeddings/create is 2048
   // but Voyage is 128
   static maxBatchSize = 128;