Merge pull request #62 from donvito/feature/baseten-provider

Feature/baseten provider

Author: donvito

README.md

@@ -50,6 +50,7 @@ More to come...check swagger docs for updated endpoints.
 | [Vercel AI Gateway](https://vercel.com/ai-gateway) | Open source and private models | Available |
 | [LlamaCpp](https://github.com/ggml-org/llama.cpp) | Local models via llama.cpp server (self-hosted) | Available |
 | [Google Gemini](https://ai.google.dev/) | Gemini models via OpenAI-compatible interface | Available |
+| [Baseten](https://baseten.co/) | Cloud-hosted ML models with OpenAI-compatible API | Available |
 
 
 ## Run the project
@@ -169,6 +170,10 @@ LMSTUDIO_BASE_URL=http://localhost:1234
 
 # OpenRouter Configuration 
 OPENROUTER_API_KEY=your-openrouter-api-key
+
+# Baseten Configuration
+BASETEN_API_KEY=your-baseten-api-key
+BASETEN_BASE_URL=https://inference.baseten.co/v1
 ```
 
 ### Google Gemini Setup

src/config/models.json

@@ -79,6 +79,12 @@
       "models": [
         { "name": "default", "capabilities": ["summarize", "pdf-summarizer", "rewrite", "compose", "planning", "keywords", "sentiment", "askText", "emailReply", "translate", "outline"], "notes": "Default model used is what you configued when running llama-server" }
       ]
+    },
+    "baseten": {
+      "enabled": true,
+      "models": [
+        { "name": "openai/gpt-oss-120b", "capabilities": ["summarize", "pdf-summarizer", "pdf-translate", "rewrite", "compose", "planning", "keywords", "sentiment", "askText", "emailReply", "translate", "meetingNotes", "outline"], "notes": "Baseten hosted OpenAI GPT OSS 120B model with comprehensive capabilities." }
+      ]
     }
   }
 }

src/config/services.ts

@@ -56,6 +56,14 @@ export interface GoogleConfig extends ServiceConfig {
   model: string;
 }
 
+export interface BasetenConfig extends ServiceConfig {
+  apiKey: string;
+  baseURL: string;
+  model: string;
+  chatModel: string;
+  timeout?: number;
+}
+
 // OpenAI Configuration
 export const openaiConfig: OpenAIConfig = {
   name: 'OpenAI',
@@ -137,8 +145,20 @@ export const googleConfig: GoogleConfig = {
   model: process.env.GEMINI_MODEL || 'gemini-2.5-flash-lite',
 };
 
+// Baseten Configuration
+export const basetenConfig: BasetenConfig = {
+  name: 'Baseten',
+  enabled: !!process.env.BASETEN_API_KEY,
+  priority: 9,
+  apiKey: process.env.BASETEN_API_KEY || '',
+  baseURL: process.env.BASETEN_BASE_URL || 'https://inference.baseten.co/v1',
+  model: process.env.BASETEN_MODEL || 'default',
+  chatModel: process.env.BASETEN_CHAT_MODEL || process.env.BASETEN_MODEL || 'openai/gpt-oss-120b',
+  timeout: parseInt(process.env.BASETEN_TIMEOUT || '30000'),
+};
+
 // Available services
-export const availableServices = [openaiConfig, anthropicConfig, ollamaConfig, openrouterConfig, lmstudioConfig, aigatewayConfig, llamacppConfig, googleConfig];
+export const availableServices = [openaiConfig, anthropicConfig, ollamaConfig, openrouterConfig, lmstudioConfig, aigatewayConfig, llamacppConfig, googleConfig, basetenConfig];
 
 // Get the primary service (highest priority enabled service)
 export function getPrimaryService(): ServiceConfig | null {

src/services/baseten.ts

@@ -0,0 +1,103 @@
+import { z } from 'zod';
+import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
+import { generateText, streamText, generateObject } from 'ai';
+import type { AIProvider } from './interfaces';
+import { basetenConfig } from '../config/services';
+
+const normalizedBase = (basetenConfig.baseURL || 'https://inference.baseten.co/v1').replace(/\/$/, '');
+const BASETEN_BASE_URL = normalizedBase;
+
+const baseten = createOpenAICompatible({
+  name: 'baseten',
+  baseURL: `${BASETEN_BASE_URL}`,
+  headers: {
+    'Authorization': `Api-Key ${basetenConfig.apiKey}`,
+  },
+});
+
+class BasetenProvider implements AIProvider {
+  name = 'baseten' as const;
+
+  async generateChatStructuredResponse(
+    prompt: string,
+    schema: z.ZodType,
+    model: string = basetenConfig.chatModel,
+    temperature: number = 0
+  ): Promise<any> {
+    try {
+      const modelToUse = model || basetenConfig.chatModel;
+      
+      const result = await generateObject({
+        model: baseten(modelToUse),
+        schema,
+        prompt,
+        temperature,
+      });
+
+      return {
+        object: result.object,
+        finishReason: result.finishReason,
+        usage: {
+          promptTokens: result.usage?.promptTokens || 0,
+          completionTokens: result.usage?.completionTokens || 0,
+          totalTokens: result.usage?.totalTokens || 0,
+        },
+        warnings: result.warnings,
+      };
+    } catch (error) {
+      throw new Error(`Baseten structured response error: ${error}`);
+    }
+  }
+  
+  
+  async generateChatTextResponse(
+    prompt: string,
+    model?: string,
+    temperature: number = 0
+  ): Promise<any> {
+    try {
+    const modelToUse = baseten(model || basetenConfig.model);
+
+    const result = await generateText({
+      model: modelToUse,
+      prompt,
+      temperature,
+    });
+
+    return result;
+    } catch (error) {
+      console.error('Baseten text response error: ', error);
+      throw new Error(`Baseten text response error: ${error}`);
+    }
+  }
+
+  async generateChatTextStreamResponse(
+    prompt: string,
+    model?: string,
+    temperature: number = 0
+  ): Promise<any> {
+    try {
+    const modelToUse = baseten(model || basetenConfig.model);
+
+    const result = await streamText({
+      model: modelToUse,
+      prompt,
+      temperature,
+    });
+
+      return result;
+    } catch (error) {
+      console.error('Baseten streaming response error: ', error);
+      throw new Error(`Baseten streaming response error: ${error}`);
+    }
+  }
+
+  async getAvailableModels(): Promise<string[]> {
+    return ['openai/gpt-oss-120b'];
+  }
+}
+
+const provider = new BasetenProvider();
+
+export default provider;
+export { BASETEN_BASE_URL };

src/services/interfaces.ts

@@ -1,6 +1,6 @@
 import { z } from 'zod';
 
-export type ProviderName = 'openai' | 'anthropic' | 'ollama' | 'openrouter' | 'lmstudio' | 'aigateway' | 'llamacpp' | 'google';
+export type ProviderName = 'openai' | 'anthropic' | 'ollama' | 'openrouter' | 'lmstudio' | 'aigateway' | 'llamacpp' | 'google' | 'baseten';
 
 export interface AIProvider {
   name: ProviderName;

src/services/providers.ts

@@ -1,3 +1,3 @@
 import { z } from 'zod';
 
-export const providersSupported = z.enum(['ollama', 'openai', 'anthropic', 'openrouter', 'lmstudio', 'aigateway', 'llamacpp', 'google']);
+export const providersSupported = z.enum(['ollama', 'openai', 'anthropic', 'openrouter', 'lmstudio', 'aigateway', 'llamacpp', 'google', 'baseten']);

src/services/registry.ts

@@ -8,6 +8,7 @@ import lmstudioProvider from './lmstudio';
 import aigatewayProvider from './aigateway';
 import llamacppProvider from './llamacpp';
 import geminiProvider from './google';
+import basetenProvider from './baseten';
 
 export class ServiceRegistry {
   private providers = new Map<ProviderName, AIProvider>();
@@ -44,6 +45,7 @@ serviceRegistry.register(lmstudioProvider);
 serviceRegistry.register(aigatewayProvider);
 serviceRegistry.register(llamacppProvider);
 serviceRegistry.register(geminiProvider);
+serviceRegistry.register(basetenProvider);
 
 // Helper for tests to replace the registry content
 export function replaceRegistryForTests(registry: ServiceRegistry) {