feat: move vLLM rerank response in the VLLM.ts and remove unused types

Author: Iamshankhadeep

core/llm/index.ts

@@ -4,8 +4,7 @@ import { findLlmInfo } from "@continuedev/llm-info";
 import {
   BaseLlmApi,
   ChatCompletionCreateParams,
-  constructLlmApi,
-  VllmRerankResponse
+  constructLlmApi
 } from "@continuedev/openai-adapters";
 import Handlebars from "handlebars";
 
@@ -1053,23 +1052,17 @@ export abstract class BaseLLM implements ILLM {
         documents: chunks.map((chunk) => chunk.content),
       });
 
-      // Handle different response formats: OpenAI (data), vLLM (results)
-      let dataArray: Array<{ index: number; relevance_score: number }>;
-      
+      // Standard OpenAI format
       if (results.data && Array.isArray(results.data)) {
-        dataArray = results.data;
-      } else if ((results as VllmRerankResponse).results && Array.isArray((results as VllmRerankResponse).results)) {
-        dataArray = (results as VllmRerankResponse).results;
-      }else {
-        throw new Error(
-          `Unexpected rerank response format from ${this.providerName}. ` +
-          `Expected 'data' or 'results' array but got: ${JSON.stringify(results)}`
-        );
+        return results.data
+          .sort((a, b) => a.index - b.index)
+          .map((result) => result.relevance_score);
       }
 
-      return dataArray
-        .sort((a, b) => a.index - b.index)
-        .map((result) => result.relevance_score);
+      throw new Error(
+        `Unexpected rerank response format from ${this.providerName}. ` +
+        `Expected 'data' array but got: ${JSON.stringify(Object.keys(results))}`
+      );
     }
 
     throw new Error(

core/llm/llms/Vllm.ts

@@ -1,7 +1,25 @@
-import { LLMOptions } from "../../index.js";
+import { Chunk, LLMOptions } from "../../index.js";
 
 import OpenAI from "./OpenAI.js";
 
+// vLLM-specific rerank response types
+interface VllmRerankItem {
+  index: number;
+  document: {
+    text: string;
+  };
+  relevance_score: number;
+}
+
+interface VllmRerankResponse {
+  id: string;
+  model: string;
+  usage: {
+    total_tokens: number;
+  };
+  results: VllmRerankItem[];
+}
+
 class Vllm extends OpenAI {
   static providerName = "vllm";
   constructor(options: LLMOptions) {
@@ -12,6 +30,28 @@ class Vllm extends OpenAI {
     }
   }
 
+  async rerank(query: string, chunks: Chunk[]): Promise<number[]> {
+    if (this.useOpenAIAdapterFor.includes("rerank") && this.openaiAdapter) {
+      const results = await this.openaiAdapter.rerank({
+        model: this.model,
+        query,
+        documents: chunks.map((chunk) => chunk.content),
+      }) as unknown as VllmRerankResponse;
+
+      // vLLM uses 'results' array instead of 'data'
+      if (results.results && Array.isArray(results.results)) {
+        const sortedResults = results.results.sort((a, b) => a.index - b.index);
+        return sortedResults.map((result) => result.index);
+      }
+
+      throw new Error(
+        `vLLM rerank response missing 'results' array. Got: ${JSON.stringify(Object.keys(results))}`
+      );
+    }
+
+    throw new Error("vLLM rerank requires OpenAI adapter");
+  }
+
   private _setupCompletionOptions() {
     this.fetch(this._getEndpoint("models"), {
       method: "GET",

packages/openai-adapters/src/apis/base.ts

@@ -1,14 +1,14 @@
 import {
-  ChatCompletion,
-  ChatCompletionChunk,
-  ChatCompletionCreateParamsNonStreaming,
-  ChatCompletionCreateParamsStreaming,
-  Completion,
-  CompletionCreateParamsNonStreaming,
-  CompletionCreateParamsStreaming,
-  CreateEmbeddingResponse,
-  EmbeddingCreateParams,
-  Model,
+    ChatCompletion,
+    ChatCompletionChunk,
+    ChatCompletionCreateParamsNonStreaming,
+    ChatCompletionCreateParamsStreaming,
+    Completion,
+    CompletionCreateParamsNonStreaming,
+    CompletionCreateParamsStreaming,
+    CreateEmbeddingResponse,
+    EmbeddingCreateParams,
+    Model,
 } from "openai/resources/index";
 
 export interface FimCreateParamsStreaming
@@ -37,23 +37,7 @@ export interface CreateRerankResponse {
   };
 }
 
-// vLLM-specific rerank response format
-export interface VllmRerankItem {
-  index: number;
-  document: {
-    text: string;
-  };
-  relevance_score: number;
-}
 
-export interface VllmRerankResponse {
-  id: string;
-  model: string;
-  usage: {
-    total_tokens: number;
-  };
-  results: VllmRerankItem[];
-}
 
 export interface BaseLlmApi {
   // Chat, no stream

packages/openai-adapters/src/index.ts

@@ -114,13 +114,10 @@ export {
   type Completion,
   type CompletionCreateParams,
   type CompletionCreateParamsNonStreaming,
-  type CompletionCreateParamsStreaming,
+  type CompletionCreateParamsStreaming
 } from "openai/resources/index";
 
 // export
-export type {
-  BaseLlmApi,
-  VllmRerankItem,
-  VllmRerankResponse,
-} from "./apis/base.js";
+export type { BaseLlmApi } from "./apis/base.js";
 export type { LLMConfig } from "./types.js";
+