diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 207c60a524..824b207891 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -36,6 +36,7 @@ export const providerNames = [ "huggingface", "cerebras", "sambanova", + "poe", ] as const export const providerNamesSchema = z.enum(providerNames) @@ -257,6 +258,11 @@ const sambaNovaSchema = apiModelIdProviderModelSchema.extend({ sambaNovaApiKey: z.string().optional(), }) +const poeSchema = apiModelIdProviderModelSchema.extend({ + poeApiKey: z.string().optional(), + poeBaseUrl: z.string().optional(), +}) + const defaultSchema = z.object({ apiProvider: z.undefined(), }) @@ -290,6 +296,7 @@ export const providerSettingsSchemaDiscriminated = z.discriminatedUnion("apiProv litellmSchema.merge(z.object({ apiProvider: z.literal("litellm") })), cerebrasSchema.merge(z.object({ apiProvider: z.literal("cerebras") })), sambaNovaSchema.merge(z.object({ apiProvider: z.literal("sambanova") })), + poeSchema.merge(z.object({ apiProvider: z.literal("poe") })), defaultSchema, ]) @@ -323,6 +330,7 @@ export const providerSettingsSchema = z.object({ ...litellmSchema.shape, ...cerebrasSchema.shape, ...sambaNovaSchema.shape, + ...poeSchema.shape, ...codebaseIndexProviderSchema.shape, }) diff --git a/packages/types/src/providers/index.ts b/packages/types/src/providers/index.ts index d6584e70ec..23500ea055 100644 --- a/packages/types/src/providers/index.ts +++ b/packages/types/src/providers/index.ts @@ -15,6 +15,7 @@ export * from "./moonshot.js" export * from "./ollama.js" export * from "./openai.js" export * from "./openrouter.js" +export * from "./poe.js" export * from "./requesty.js" export * from "./sambanova.js" export * from "./unbound.js" diff --git a/packages/types/src/providers/poe.ts b/packages/types/src/providers/poe.ts new file mode 100644 index 0000000000..0f2940907c --- /dev/null +++ b/packages/types/src/providers/poe.ts @@ -0,0 +1,236 @@ +import type { ModelInfo } from "../model.js" + +// https://creator.poe.com/docs/external-applications/openai-compatible-api +export type PoeModelId = + | "gpt-4o" + | "gpt-4o-mini" + | "o1-preview" + | "o1-mini" + | "claude-3-5-sonnet" + | "claude-3-5-haiku" + | "claude-3-opus" + | "claude-3-sonnet" + | "claude-3-haiku" + | "gemini-1.5-pro" + | "gemini-1.5-flash" + | "gemini-2.0-flash-exp" + | "llama-3.1-405b" + | "llama-3.1-70b" + | "llama-3.1-8b" + | "llama-3.3-70b" + | "mistral-large" + | "grok-2" + | "grok-2-mini" + | "kimi-k1.5" + | "glm-4-plus" + | string // Allow custom bot names + +export const poeDefaultModelId: PoeModelId = "claude-3-5-sonnet" + +export const poeModels = { + // GPT Models + "gpt-4o": { + maxTokens: 16384, + contextWindow: 128000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 2.5, + outputPrice: 10, + description: "OpenAI GPT-4o model via Poe", + }, + "gpt-4o-mini": { + maxTokens: 16384, + contextWindow: 128000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 0.15, + outputPrice: 0.6, + description: "OpenAI GPT-4o Mini model via Poe", + }, + "o1-preview": { + maxTokens: 32768, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 15, + outputPrice: 60, + description: "OpenAI o1-preview reasoning model via Poe", + }, + "o1-mini": { + maxTokens: 65536, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 3, + outputPrice: 12, + description: "OpenAI o1-mini reasoning model via Poe", + }, + // Claude Models + "claude-3-5-sonnet": { + maxTokens: 8192, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 3, + outputPrice: 15, + description: "Anthropic Claude 3.5 Sonnet model via Poe", + }, + "claude-3-5-haiku": { + maxTokens: 8192, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 0.8, + outputPrice: 4, + description: "Anthropic Claude 3.5 Haiku model via Poe", + }, + "claude-3-opus": { + maxTokens: 4096, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 15, + outputPrice: 75, + description: "Anthropic Claude 3 Opus model via Poe", + }, + "claude-3-sonnet": { + maxTokens: 4096, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 3, + outputPrice: 15, + description: "Anthropic Claude 3 Sonnet model via Poe", + }, + "claude-3-haiku": { + maxTokens: 4096, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 0.25, + outputPrice: 1.25, + description: "Anthropic Claude 3 Haiku model via Poe", + }, + // Gemini Models + "gemini-1.5-pro": { + maxTokens: 8192, + contextWindow: 2097152, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 2.5, + outputPrice: 10, + description: "Google Gemini 1.5 Pro model via Poe", + }, + "gemini-1.5-flash": { + maxTokens: 8192, + contextWindow: 1048576, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 0.075, + outputPrice: 0.3, + description: "Google Gemini 1.5 Flash model via Poe", + }, + "gemini-2.0-flash-exp": { + maxTokens: 8192, + contextWindow: 1048576, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 0, + outputPrice: 0, + description: "Google Gemini 2.0 Flash Experimental model via Poe", + }, + // Llama Models + "llama-3.1-405b": { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 2.7, + outputPrice: 2.7, + description: "Meta Llama 3.1 405B model via Poe", + }, + "llama-3.1-70b": { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 0.88, + outputPrice: 0.88, + description: "Meta Llama 3.1 70B model via Poe", + }, + "llama-3.1-8b": { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 0.18, + outputPrice: 0.18, + description: "Meta Llama 3.1 8B model via Poe", + }, + "llama-3.3-70b": { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 0.88, + outputPrice: 0.88, + description: "Meta Llama 3.3 70B model via Poe", + }, + // Other Models + "mistral-large": { + maxTokens: 8192, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 2, + outputPrice: 6, + description: "Mistral Large model via Poe", + }, + "grok-2": { + maxTokens: 4096, + contextWindow: 131072, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 5, + outputPrice: 10, + description: "xAI Grok-2 model via Poe", + }, + "grok-2-mini": { + maxTokens: 4096, + contextWindow: 131072, + supportsImages: true, + supportsPromptCache: false, + inputPrice: 2, + outputPrice: 10, + description: "xAI Grok-2 Mini model via Poe", + }, + "kimi-k1.5": { + maxTokens: 8192, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 14, + outputPrice: 14, + description: "Moonshot AI Kimi K1.5 model via Poe", + }, + "glm-4-plus": { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 10, + outputPrice: 10, + description: "Zhipu AI GLM-4 Plus model via Poe", + }, +} as const satisfies Record + +// Default model info for custom bots +export const poeDefaultModelInfo: ModelInfo = { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: false, + inputPrice: 0, + outputPrice: 0, + description: "Custom Poe bot", +} diff --git a/src/api/index.ts b/src/api/index.ts index 5daa53396f..9502a18aa0 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -33,6 +33,7 @@ import { ClaudeCodeHandler, SambaNovaHandler, DoubaoHandler, + PoeHandler, } from "./providers" export interface SingleCompletionHandler { @@ -124,6 +125,8 @@ export function buildApiHandler(configuration: ProviderSettings): ApiHandler { return new CerebrasHandler(options) case "sambanova": return new SambaNovaHandler(options) + case "poe": + return new PoeHandler(options) default: apiProvider satisfies "gemini-cli" | undefined return new AnthropicHandler(options) diff --git a/src/api/providers/__tests__/poe.spec.ts b/src/api/providers/__tests__/poe.spec.ts new file mode 100644 index 0000000000..b94a49dad9 --- /dev/null +++ b/src/api/providers/__tests__/poe.spec.ts @@ -0,0 +1,160 @@ +// npx vitest run src/api/providers/__tests__/poe.spec.ts + +// Mock vscode first to avoid import errors +vitest.mock("vscode", () => ({})) + +import OpenAI from "openai" +import { Anthropic } from "@anthropic-ai/sdk" + +import { type PoeModelId, poeDefaultModelId, poeModels, poeDefaultModelInfo } from "@roo-code/types" + +import { PoeHandler } from "../poe" + +vitest.mock("openai", () => { + const createMock = vitest.fn() + return { + default: vitest.fn(() => ({ chat: { completions: { create: createMock } } })), + } +}) + +describe("PoeHandler", () => { + let handler: PoeHandler + let mockCreate: any + + beforeEach(() => { + vitest.clearAllMocks() + mockCreate = (OpenAI as unknown as any)().chat.completions.create + handler = new PoeHandler({ poeApiKey: "test-poe-api-key" }) + }) + + it("should use the correct Poe base URL", () => { + new PoeHandler({ poeApiKey: "test-poe-api-key" }) + expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ baseURL: "https://api.poe.com/v1" })) + }) + + it("should use custom base URL when provided", () => { + const customBaseUrl = "https://custom.poe.api/v1" + new PoeHandler({ poeApiKey: "test-poe-api-key", poeBaseUrl: customBaseUrl }) + expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ baseURL: customBaseUrl })) + }) + + it("should use the provided API key", () => { + const poeApiKey = "test-poe-api-key" + new PoeHandler({ poeApiKey }) + expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: poeApiKey })) + }) + + it("should return default model when no model is specified", () => { + const model = handler.getModel() + expect(model.id).toBe(poeDefaultModelId) + expect(model.info).toEqual(poeModels[poeDefaultModelId as keyof typeof poeModels]) + }) + + it("should return specified model when valid model is provided", () => { + const testModelId: PoeModelId = "gpt-4o" + const handlerWithModel = new PoeHandler({ apiModelId: testModelId, poeApiKey: "test-poe-api-key" }) + const model = handlerWithModel.getModel() + expect(model.id).toBe(testModelId) + expect(model.info).toEqual(poeModels[testModelId as keyof typeof poeModels]) + }) + + it("should return custom bot with default info for unknown models", () => { + const customBotName = "MyCustomBot" + const handlerWithCustomBot = new PoeHandler({ apiModelId: customBotName, poeApiKey: "test-poe-api-key" }) + const model = handlerWithCustomBot.getModel() + expect(model.id).toBe(customBotName) + expect(model.info).toEqual(poeDefaultModelInfo) + }) + + it("completePrompt method should return text from Poe API", async () => { + const expectedResponse = "This is a test response from Poe" + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: expectedResponse } }] }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe(expectedResponse) + }) + + it("should handle errors in completePrompt", async () => { + const errorMessage = "Poe API error" + mockCreate.mockRejectedValueOnce(new Error(errorMessage)) + await expect(handler.completePrompt("test prompt")).rejects.toThrow(`Poe completion error: ${errorMessage}`) + }) + + it("createMessage should yield text content from stream", async () => { + const testContent = "This is test content from Poe stream" + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + next: vitest + .fn() + .mockResolvedValueOnce({ + done: false, + value: { choices: [{ delta: { content: testContent } }] }, + }) + .mockResolvedValueOnce({ done: true }), + }), + } + }) + + const stream = handler.createMessage("system prompt", []) + const firstChunk = await stream.next() + + expect(firstChunk.done).toBe(false) + expect(firstChunk.value).toEqual({ type: "text", text: testContent }) + }) + + it("createMessage should yield usage data from stream", async () => { + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + next: vitest + .fn() + .mockResolvedValueOnce({ + done: false, + value: { choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 20 } }, + }) + .mockResolvedValueOnce({ done: true }), + }), + } + }) + + const stream = handler.createMessage("system prompt", []) + const firstChunk = await stream.next() + + expect(firstChunk.done).toBe(false) + expect(firstChunk.value).toEqual({ type: "usage", inputTokens: 10, outputTokens: 20 }) + }) + + it("createMessage should pass correct parameters to Poe client", async () => { + const modelId: PoeModelId = "claude-3-5-sonnet" + const modelInfo = poeModels[modelId as keyof typeof poeModels] + const handlerWithModel = new PoeHandler({ apiModelId: modelId, poeApiKey: "test-poe-api-key" }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Poe" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Poe" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + max_tokens: modelInfo.maxTokens, + temperature: 0.7, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + }) +}) diff --git a/src/api/providers/index.ts b/src/api/providers/index.ts index a1b8f25536..fc2d4727a4 100644 --- a/src/api/providers/index.ts +++ b/src/api/providers/index.ts @@ -20,6 +20,7 @@ export { OllamaHandler } from "./ollama" export { OpenAiNativeHandler } from "./openai-native" export { OpenAiHandler } from "./openai" export { OpenRouterHandler } from "./openrouter" +export { PoeHandler } from "./poe" export { RequestyHandler } from "./requesty" export { SambaNovaHandler } from "./sambanova" export { UnboundHandler } from "./unbound" diff --git a/src/api/providers/poe.ts b/src/api/providers/poe.ts new file mode 100644 index 0000000000..9904ca240d --- /dev/null +++ b/src/api/providers/poe.ts @@ -0,0 +1,31 @@ +import { type PoeModelId, poeDefaultModelId, poeModels, poeDefaultModelInfo } from "@roo-code/types" + +import type { ApiHandlerOptions } from "../../shared/api" + +import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider" + +export class PoeHandler extends BaseOpenAiCompatibleProvider { + constructor(options: ApiHandlerOptions) { + super({ + ...options, + providerName: "Poe", + baseURL: options.poeBaseUrl || "https://api.poe.com/v1", + apiKey: options.poeApiKey, + defaultProviderModelId: poeDefaultModelId, + providerModels: poeModels, + defaultTemperature: 0.7, + }) + } + + override getModel() { + const modelId = this.options.apiModelId || this.defaultProviderModelId + + // Check if it's a known model + if (modelId in this.providerModels) { + return { id: modelId, info: this.providerModels[modelId as PoeModelId] } + } + + // For custom bots, use default model info + return { id: modelId, info: poeDefaultModelInfo } + } +}