scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: migrate LiteLLM provider to AI SDK (@ai-sdk/openai-compatible) (#11348)

Browse source 
* feat: migrate LiteLLM provider to AI SDK (@ai-sdk/openai-compatible)

- Replace raw OpenAI SDK (RouterProvider) with Vercel AI SDK's
  createOpenAICompatible via OpenAICompatibleHandler base class
- Retain dynamic model fetching from LiteLLM server via /v1/model/info
- Use centralized getModelMaxOutputTokens() to cap output tokens at 20%
  of context window, preventing overflow errors
- Remove LiteLLM-specific workarounds (Gemini thought signature injection,
  prompt cache control headers) now handled by the proxy or AI SDK
- Rewrite tests to mock AI SDK (streamText, generateText) instead of
  raw OpenAI SDK

* fix: call fetchModel() in completePrompt() before execution

Addresses review feedback - completePrompt() now fetches models
before executing to ensure correct model info for token limits,
matching the behavior of createMessage().
This commit is contained in:
Daniel 2026-02-09 17:20:42 -05:00 • committed by GitHub
parent 70775f0ec1
commit c74acf36ef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 503 additions and 1105 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1225
src/api/providers/__tests__/lite-llm.spec.ts
View file

File diff suppressed because it is too large Load diff

383
src/api/providers/lite-llm.ts
View file

@ -1,331 +1,112 @@
import OpenAI from "openai"
import { Anthropic } from "@anthropic-ai/sdk" // Keep for type usage only
/**
* LiteLLM provider handler using Vercel AI SDK.
*
* This handler uses @ai-sdk/openai-compatible to communicate with LiteLLM proxy servers.
* LiteLLM follows the OpenAI API format, making it compatible with the OpenAI-compatible provider.
* Models are dynamically fetched from the LiteLLM server via /v1/model/info.
*/
import { litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types"
import { Anthropic } from "@anthropic-ai/sdk"
import { LanguageModel } from "ai"
import { calculateApiCostOpenAI } from "../../shared/cost"
import { litellmDefaultModelId, litellmDefaultModelInfo, type ModelInfo, type ModelRecord } from "@roo-code/types"
import { ApiHandlerOptions } from "../../shared/api"
import { type ApiHandlerOptions, getModelMaxOutputTokens } from "../../shared/api"
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { sanitizeOpenAiCallId } from "../../utils/tool-id"
import { OpenAICompatibleHandler } from "./openai-compatible"
import { getModels, getModelsFromCache } from "./fetchers/modelCache"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
import { RouterProvider } from "./router-provider"
/**
* LiteLLM provider handler
*
* This handler uses the LiteLLM API to proxy requests to various LLM providers.
* It follows the OpenAI API format for compatibility.
*/
export class LiteLLMHandler extends RouterProvider implements SingleCompletionHandler {
export class LiteLLMHandler extends OpenAICompatibleHandler implements SingleCompletionHandler {
private models: ModelRecord = {}
constructor(options: ApiHandlerOptions) {
super({
options,
name: "litellm",
baseURL: `${options.litellmBaseUrl || "http://localhost:4000"}`,
const modelId = options.litellmModelId || litellmDefaultModelId
super(options, {
providerName: "LiteLLM",
baseURL: options.litellmBaseUrl || "http://localhost:4000",
apiKey: options.litellmApiKey || "dummy-key",
modelId: options.litellmModelId,
defaultModelId: litellmDefaultModelId,
defaultModelInfo: litellmDefaultModelInfo,
modelId,
modelInfo: litellmDefaultModelInfo,
temperature: options.modelTemperature ?? 0,
modelMaxTokens: options.modelMaxTokens,
})
}
private isGpt5(modelId: string): boolean {
// Match gpt-5, gpt5, and variants like gpt-5o, gpt-5-turbo, gpt5-preview, gpt-5.1
// Avoid matching gpt-50, gpt-500, etc.
return /\bgpt-?5(?!\d)/i.test(modelId)
private async fetchModel() {
this.models = await getModels({
provider: "litellm",
apiKey: this.config.apiKey,
baseUrl: this.config.baseURL,
})
return this.getModel()
}
/**
* Detect if the model is a Gemini model that requires thought signature handling.
* Gemini 3 models validate thought signatures for tool/function calling steps.
*/
private isGeminiModel(modelId: string): boolean {
// Match various Gemini model patterns:
// - gemini-3-pro, gemini-3-flash, gemini-3-*
// - gemini 3 pro, Gemini 3 Pro (space-separated, case-insensitive)
// - gemini/gemini-3-*, google/gemini-3-*
// - vertex_ai/gemini-3-*, vertex/gemini-3-*
// Also match Gemini 2.5+ models which use similar validation
const lowerModelId = modelId.toLowerCase()
override getModel(): { id: string; info: ModelInfo } {
const id = this.config.modelId || litellmDefaultModelId
if (this.models[id]) {
return { id, info: this.models[id] }
}
const cachedModels = getModelsFromCache("litellm")
if (cachedModels?.[id]) {
this.models = cachedModels
return { id, info: cachedModels[id] }
}
return { id: this.config.modelId || litellmDefaultModelId, info: litellmDefaultModelInfo }
}
protected override getLanguageModel(): LanguageModel {
const { id } = this.getModel()
return this.provider(id)
}
protected override getMaxOutputTokens(): number | undefined {
const { id, info } = this.getModel()
return (
// Match hyphenated versions: gemini-3, gemini-2.5
lowerModelId.includes("gemini-3") ||
lowerModelId.includes("gemini-2.5") ||
// Match space-separated versions: "gemini 3", "gemini 2.5"
// This handles model names like "Gemini 3 Pro" from LiteLLM model groups
lowerModelId.includes("gemini 3") ||
lowerModelId.includes("gemini 2.5") ||
// Also match provider-prefixed versions
/\b(gemini|google|vertex_ai|vertex)\/gemini[-\s](3|2\.5)/i.test(modelId)
getModelMaxOutputTokens({
modelId: id,
model: info,
settings: this.options,
format: "openai",
}) ?? undefined
)
}
/**
* Inject thought signatures for Gemini models via provider_specific_fields.
* This is required when switching from other models to Gemini to satisfy API validation
* for function calls that weren't generated by Gemini (and thus lack thought signatures).
*
* Per LiteLLM documentation:
* - Thought signatures are stored in provider_specific_fields.thought_signature of tool calls
* - The dummy signature base64("skip_thought_signature_validator") bypasses validation
*
* We inject the dummy signature on EVERY tool call unconditionally to ensure Gemini
* doesn't complain about missing/corrupted signatures when conversation history
* contains tool calls from other models (like Claude).
*/
private injectThoughtSignatureForGemini(
openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[],
): OpenAI.Chat.ChatCompletionMessageParam[] {
// Base64 encoded "skip_thought_signature_validator" as per LiteLLM docs
const dummySignature = Buffer.from("skip_thought_signature_validator").toString("base64")
return openAiMessages.map((msg) => {
if (msg.role === "assistant") {
const toolCalls = (msg as any).tool_calls as any[] | undefined
// Only process if there are tool calls
if (toolCalls && toolCalls.length > 0) {
// Inject dummy signature into ALL tool calls' provider_specific_fields
// This ensures Gemini doesn't reject tool calls from other models
const updatedToolCalls = toolCalls.map((tc) => ({
...tc,
provider_specific_fields: {
...(tc.provider_specific_fields || {}),
thought_signature: dummySignature,
},
}))
return {
...msg,
tool_calls: updatedToolCalls,
}
}
}
return msg
})
}
override async *createMessage(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
metadata?: ApiHandlerCreateMessageMetadata,
): ApiStream {
const { id: modelId, info } = await this.fetchModel()
const openAiMessages = convertToOpenAiMessages(messages, {
normalizeToolCallId: sanitizeOpenAiCallId,
})
// Prepare messages with cache control if enabled and supported
let systemMessage: OpenAI.Chat.ChatCompletionMessageParam
let enhancedMessages: OpenAI.Chat.ChatCompletionMessageParam[]
if (this.options.litellmUsePromptCache && info.supportsPromptCache) {
// Create system message with cache control in the proper format
systemMessage = {
role: "system",
content: [
{
type: "text",
text: systemPrompt,
cache_control: { type: "ephemeral" },
} as any,
],
}
// Find the last two user messages to apply caching
const userMsgIndices = openAiMessages.reduce(
(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
[] as number[],
)
const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
const secondLastUserMsgIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
// Apply cache_control to the last two user messages
enhancedMessages = openAiMessages.map((message, index) => {
if ((index === lastUserMsgIndex || index === secondLastUserMsgIndex) && message.role === "user") {
// Handle both string and array content types
if (typeof message.content === "string") {
return {
...message,
content: [
{
type: "text",
text: message.content,
cache_control: { type: "ephemeral" },
} as any,
],
}
} else if (Array.isArray(message.content)) {
// Apply cache control to the last content item in the array
return {
...message,
content: message.content.map((content, contentIndex) =>
contentIndex === message.content.length - 1
? ({
...content,
cache_control: { type: "ephemeral" },
} as any)
: content,
),
}
}
}
return message
})
} else {
// No cache control - use simple format
systemMessage = { role: "system", content: systemPrompt }
enhancedMessages = openAiMessages
}
// Required by some providers; others default to max tokens allowed
let maxTokens: number | undefined = info.maxTokens ?? undefined
// Check if this is a GPT-5 model that requires max_completion_tokens instead of max_tokens
const isGPT5Model = this.isGpt5(modelId)
// For Gemini models with native protocol: inject fake reasoning.encrypted block for tool calls
// This is required when switching from other models to Gemini to satisfy API validation.
// Gemini 3 models validate thought signatures for function calls, and when conversation
// history contains tool calls from other models (like Claude), they lack the required
// signatures. The "skip_thought_signature_validator" value bypasses this validation.
const isGemini = this.isGeminiModel(modelId)
let processedMessages = enhancedMessages
if (isGemini) {
processedMessages = this.injectThoughtSignatureForGemini(enhancedMessages)
}
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
model: modelId,
messages: [systemMessage, ...processedMessages],
stream: true,
stream_options: {
include_usage: true,
},
tools: this.convertToolsForOpenAI(metadata?.tools),
tool_choice: metadata?.tool_choice,
}
// GPT-5 models require max_completion_tokens instead of the deprecated max_tokens parameter
if (isGPT5Model && maxTokens) {
requestOptions.max_completion_tokens = maxTokens
} else if (maxTokens) {
requestOptions.max_tokens = maxTokens
}
if (this.supportsTemperature(modelId)) {
requestOptions.temperature = this.options.modelTemperature ?? 0
}
try {
const { data: completion } = await this.client.chat.completions.create(requestOptions).withResponse()
let lastUsage
for await (const chunk of completion) {
const delta = chunk.choices[0]?.delta
const usage = chunk.usage as LiteLLMUsage
if (delta?.content) {
yield { type: "text", text: delta.content }
}
// Handle tool calls in stream - emit partial chunks for NativeToolCallParser
if (delta?.tool_calls) {
for (const toolCall of delta.tool_calls) {
yield {
type: "tool_call_partial",
index: toolCall.index,
id: toolCall.id,
name: toolCall.function?.name,
arguments: toolCall.function?.arguments,
}
}
}
if (usage) {
lastUsage = usage
}
}
if (lastUsage) {
// Extract cache-related information if available
// LiteLLM may use different field names for cache tokens
const cacheWriteTokens =
lastUsage.cache_creation_input_tokens || (lastUsage as any).prompt_cache_miss_tokens || 0
const cacheReadTokens =
lastUsage.prompt_tokens_details?.cached_tokens ||
(lastUsage as any).cache_read_input_tokens ||
(lastUsage as any).prompt_cache_hit_tokens ||
0
const { totalCost } = calculateApiCostOpenAI(
info,
lastUsage.prompt_tokens || 0,
lastUsage.completion_tokens || 0,
cacheWriteTokens,
cacheReadTokens,
)
const usageData: ApiStreamUsageChunk = {
type: "usage",
inputTokens: lastUsage.prompt_tokens || 0,
outputTokens: lastUsage.completion_tokens || 0,
cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined,
cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined,
totalCost,
}
yield usageData
}
} catch (error) {
if (error instanceof Error) {
throw new Error(`LiteLLM streaming error: ${error.message}`)
}
throw error
}
await this.fetchModel()
yield* super.createMessage(systemPrompt, messages, metadata)
}
async completePrompt(prompt: string): Promise<string> {
const { id: modelId, info } = await this.fetchModel()
override async completePrompt(prompt: string): Promise<string> {
await this.fetchModel()
return super.completePrompt(prompt)
}
// Check if this is a GPT-5 model that requires max_completion_tokens instead of max_tokens
const isGPT5Model = this.isGpt5(modelId)
try {
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
model: modelId,
messages: [{ role: "user", content: prompt }],
}
if (this.supportsTemperature(modelId)) {
requestOptions.temperature = this.options.modelTemperature ?? 0
}
// GPT-5 models require max_completion_tokens instead of the deprecated max_tokens parameter
if (isGPT5Model && info.maxTokens) {
requestOptions.max_completion_tokens = info.maxTokens
} else if (info.maxTokens) {
requestOptions.max_tokens = info.maxTokens
}
const response = await this.client.chat.completions.create(requestOptions)
return response.choices[0]?.message.content || ""
} catch (error) {
if (error instanceof Error) {
throw new Error(`LiteLLM completion error: ${error.message}`)
}
throw error
protected override processUsageMetrics(usage: {
inputTokens?: number
outputTokens?: number
details?: {
cachedInputTokens?: number
reasoningTokens?: number
}
raw?: Record<string, unknown>
}): ApiStreamUsageChunk {
return {
type: "usage",
inputTokens: usage.inputTokens || 0,
outputTokens: usage.outputTokens || 0,
cacheReadTokens: usage.details?.cachedInputTokens,
reasoningTokens: usage.details?.reasoningTokens,
}
}
}
// LiteLLM usage may include an extra field for Anthropic use cases.
interface LiteLLMUsage extends OpenAI.CompletionUsage {
cache_creation_input_tokens?: number
}