diff --git a/README.md b/README.md index e183d364d5..eff650a906 100644 --- a/README.md +++ b/README.md @@ -189,7 +189,7 @@ For additional configurations in HTTP server setup, please read [this](https://g ## 🔍 Overview -Bifrost acts as a bridge between your applications and multiple AI providers (OpenAI, Anthropic, Amazon Bedrock, etc.). It provides a consistent API interface while handling: +Bifrost acts as a bridge between your applications and multiple AI providers (OpenAI, Anthropic, Amazon Bedrock, Mistral, Ollama, etc.). It provides a consistent API while handling: - Authentication and key management - Request routing and load balancing @@ -203,7 +203,7 @@ With Bifrost, you can focus on building your AI-powered applications without wor ## ✨ Features -- **Multi-Provider Support**: Integrate with OpenAI, Anthropic, Amazon Bedrock, and more through a single API +- **Multi-Provider Support**: Integrate with OpenAI, Anthropic, Amazon Bedrock, Mistral, Ollama, and more through a single API - **Fallback Mechanisms**: Automatically retry failed requests with alternative models or providers - **Dynamic Key Management**: Rotate and manage API keys efficiently - **Connection Pooling**: Optimize network resources for better performance @@ -211,7 +211,7 @@ With Bifrost, you can focus on building your AI-powered applications without wor - **Flexible Transports**: Multiple transports for easy integration into your infra - **Plugin First Architecture**: No callback hell, simple addition/creation of custom plugins - **Custom Configuration**: Offers granular control over pool sizes, network retry settings, fallback providers, and network proxy configurations -- **Build in Observability**: Native Prometheus metrics out of the box, no wrappers, no sidecars, just drop it in and scrape +- **Built-in Observability**: Native Prometheus metrics out of the box, no wrappers, no sidecars, just drop it in and scrape --- @@ -269,7 +269,7 @@ client, err := bifrost.Init(schemas.BifrostConfig{ }) ``` -3. **transports**: This package contains transport clients like HTTP to expose your Bifrost client. You can either `go get` this package or directly use the independent Dockerfile to quickly spin up your Bifrost API interface ([Click here](https://github.com/maximhq/bifrost/tree/main/transports/README.md) to read more on this). +1. **transports**: This package contains transport clients like HTTP to expose your Bifrost client. You can either `go get` this package or directly use the independent Dockerfile to quickly spin up your Bifrost API ([Click here](https://github.com/maximhq/bifrost/tree/main/transports/README.md) to read more on this). ### Additional Configurations @@ -318,9 +318,9 @@ Bifrost has been tested under high load conditions to ensure optimal performance | Response Parsing | 11.30 ms | 2.11 ms | | **Bifrost's Overhead** | **`59 µs\*`** | **`11 µs\*`** | -_\*Bifrost's overhead is measured at 59 µs on t3.medium and 11 µs on t3.xlarge, excluding the time taken for JSON marshalling and the HTTP call to the LLM, both of which are required in any custom implementation._ +_\*Bifrost's overhead is measured at 59 µs on t3.medium and 11 µs on t3.xlarge, excluding the time taken for JSON marshalling and the HTTP call to the LLM, both of which are required in any custom implementation._ -**Note**: On the t3.xlarge, we tested with significantly larger response payloads (~10 KB average vs ~1 KB on t3.medium). Even so, response parsing time dropped dramatically thanks to better CPU throughput and Bifrost's optimized memory reuse. +**Note**: On the t3.xlarge, we tested with significantly larger response payloads (~10 KB average vs ~1 KB on t3.medium). Even so, response parsing time dropped dramatically thanks to better CPU throughput and Bifrost's optimized memory reuse. ### Key Performance Highlights @@ -344,13 +344,13 @@ One of Bifrost's key strengths is its flexibility in configuration. You can free - Buffer and Concurrency Settings: Controls the queue size and maximum number of concurrent requests (adjustable per provider). - Retry and Timeout Configurations: Customizable based on your requirements for each provider. -Curious? Run your own benchmarks. The [Bifrost Benchmarking](https://github.com/maximhq/bifrost-benchmarking) repo has everything you need to test it in your own environment. +Curious? Run your own benchmarks. The [Bifrost Benchmarking](https://github.com/maximhq/bifrost-benchmarking) repo has everything you need to test it in your own environment. --- ## 🤝 Contributing -We welcome contributions of all kinds—whether it's bug fixes, features, documentation improvements, or new ideas. Feel free to open an issue, and once its’ assigned, submit a Pull Request. +We welcome contributions of all kinds—whether it's bug fixes, features, documentation improvements, or new ideas. Feel free to open an issue, and once it's assigned, submit a Pull Request. Here's how to get started (after picking up an issue): diff --git a/core/bifrost.go b/core/bifrost.go index 67ddcb09fa..38b00ad5c6 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -140,6 +140,8 @@ func (bifrost *Bifrost) createProviderFromProviderKey(providerKey schemas.ModelP return providers.NewVertexProvider(config, bifrost.logger) case schemas.Mistral: return providers.NewMistralProvider(config, bifrost.logger), nil + case schemas.Ollama: + return providers.NewOllamaProvider(config, bifrost.logger) default: return nil, fmt.Errorf("unsupported provider: %s", providerKey) } @@ -153,8 +155,8 @@ func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, confi return fmt.Errorf("failed to get config for provider: %v", err) } - // Check if the provider has any keys (skip vertex) - if providerKey != schemas.Vertex { + // Check if the provider has any keys (skip keyless providers) + if providerRequiresKey(providerKey) { keys, err := bifrost.account.GetKeysForProvider(providerKey) if err != nil || len(keys) == 0 { return fmt.Errorf("failed to get keys for provider: %v", err) @@ -349,6 +351,12 @@ var retryableStatusCodes = map[int]bool{ 429: true, // Too Many Requests } +// providerRequiresKey returns true if the given provider requires an API key for authentication. +// Some providers like Vertex and Ollama are keyless and don't require API keys. +func providerRequiresKey(providerKey schemas.ModelProvider) bool { + return providerKey != schemas.Vertex && providerKey != schemas.Ollama +} + // calculateBackoff implements exponential backoff with jitter for retry attempts. func (bifrost *Bifrost) calculateBackoff(attempt int, config *schemas.ProviderConfig) time.Duration { // Calculate an exponential backoff: initial * 2^attempt @@ -371,7 +379,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, queue chan Chan var err error key := "" - if provider.GetProviderKey() != schemas.Vertex { + if providerRequiresKey(provider.GetProviderKey()) { key, err = bifrost.selectKeyFromProviderForModel(provider.GetProviderKey(), req.Model) if err != nil { bifrost.logger.Warn(fmt.Sprintf("Error selecting key for model %s: %v", req.Model, err)) diff --git a/core/providers/ollama.go b/core/providers/ollama.go new file mode 100644 index 0000000000..e29ac40c5c --- /dev/null +++ b/core/providers/ollama.go @@ -0,0 +1,182 @@ +// Package providers implements various LLM providers and their utility functions. +// This file contains the Ollama provider implementation. +package providers + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + "github.com/goccy/go-json" + + schemas "github.com/maximhq/bifrost/core/schemas" + "github.com/valyala/fasthttp" +) + +// OllamaResponse represents the response structure from the Ollama API. +type OllamaResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Choices []schemas.BifrostResponseChoice `json:"choices"` + Model string `json:"model"` + Created int `json:"created"` + Usage schemas.LLMUsage `json:"usage"` +} + +// ollamaResponsePool provides a pool for Ollama response objects. +var ollamaResponsePool = sync.Pool{ + New: func() interface{} { + return &OllamaResponse{} + }, +} + +// acquireOllamaResponse gets a Ollama response from the pool and resets it. +func acquireOllamaResponse() *OllamaResponse { + resp := ollamaResponsePool.Get().(*OllamaResponse) + *resp = OllamaResponse{} // Reset the struct + return resp +} + +// releaseOllamaResponse returns a Ollama response to the pool. +func releaseOllamaResponse(resp *OllamaResponse) { + if resp != nil { + ollamaResponsePool.Put(resp) + } +} + +// OllamaProvider implements the Provider interface for Ollama's API. +type OllamaProvider struct { + logger schemas.Logger // Logger for provider operations + client *fasthttp.Client // HTTP client for API requests + baseURL string // Base URL for the provider +} + +// NewOllamaProvider creates a new Ollama provider instance. +// It initializes the HTTP client with the provided configuration and sets up response pools. +// The client is configured with timeouts, concurrency limits, and optional proxy settings. +func NewOllamaProvider(config *schemas.ProviderConfig, logger schemas.Logger) (*OllamaProvider, error) { + config.CheckAndSetDefaults() + + client := &fasthttp.Client{ + ReadTimeout: time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds), + WriteTimeout: time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds), + MaxConnsPerHost: config.ConcurrencyAndBufferSize.BufferSize, + } + + // Pre-warm response pools + for range config.ConcurrencyAndBufferSize.Concurrency { + ollamaResponsePool.Put(&OllamaResponse{}) + bifrostResponsePool.Put(&schemas.BifrostResponse{}) + } + + // Configure proxy if provided + client = configureProxy(client, config.ProxyConfig, logger) + + baseURL := strings.TrimRight(config.NetworkConfig.BaseURL, "/") + if baseURL == "" { + return nil, fmt.Errorf("base_url is required for ollama provider") + } + + return &OllamaProvider{ + logger: logger, + client: client, + baseURL: baseURL, + }, nil +} + +// GetProviderKey returns the provider identifier for Ollama. +func (provider *OllamaProvider) GetProviderKey() schemas.ModelProvider { + return schemas.Ollama +} + +// TextCompletion is not supported by the Ollama provider. +func (provider *OllamaProvider) TextCompletion(ctx context.Context, model, key, text string, params *schemas.ModelParameters) (*schemas.BifrostResponse, *schemas.BifrostError) { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: schemas.ErrorField{ + Message: "text completion is not supported by ollama provider", + }, + } +} + +// ChatCompletion performs a chat completion request to the Ollama API. +func (provider *OllamaProvider) ChatCompletion(ctx context.Context, model, key string, messages []schemas.BifrostMessage, params *schemas.ModelParameters) (*schemas.BifrostResponse, *schemas.BifrostError) { + formattedMessages, preparedParams := prepareOpenAIChatRequest(messages, params) + + requestBody := mergeConfig(map[string]interface{}{ + "model": model, + "messages": formattedMessages, + }, preparedParams) + + jsonBody, err := json.Marshal(requestBody) + if err != nil { + return nil, &schemas.BifrostError{ + IsBifrostError: true, + Error: schemas.ErrorField{ + Message: schemas.ErrProviderJSONMarshaling, + Error: err, + }, + } + } + + // Create request + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(provider.baseURL + "/v1/chat/completions") + req.Header.SetMethod("POST") + req.Header.SetContentType("application/json") + if key != "" { + req.Header.Set("Authorization", "Bearer "+key) + } + req.SetBody(jsonBody) + + // Make request + bifrostErr := makeRequestWithContext(ctx, provider.client, req, resp) + if bifrostErr != nil { + return nil, bifrostErr + } + + // Handle error response + if resp.StatusCode() != fasthttp.StatusOK { + provider.logger.Debug(fmt.Sprintf("error from ollama provider: %s", string(resp.Body()))) + + var errorResp map[string]interface{} + bifrostErr := handleProviderAPIError(resp, &errorResp) + bifrostErr.Error.Message = fmt.Sprintf("Ollama error: %v", errorResp) + return nil, bifrostErr + } + + responseBody := resp.Body() + + // Pre-allocate response structs from pools + response := acquireOllamaResponse() + defer releaseOllamaResponse(response) + + result := acquireBifrostResponse() + defer releaseBifrostResponse(result) + + // Use enhanced response handler with pre-allocated response + rawResponse, bifrostErr := handleProviderResponse(responseBody, response) + if bifrostErr != nil { + return nil, bifrostErr + } + + // Populate result from response + result.ID = response.ID + result.Choices = response.Choices + result.Object = response.Object + result.Usage = response.Usage + result.Model = response.Model + result.Created = response.Created + result.ExtraFields = schemas.BifrostResponseExtraFields{ + Provider: schemas.Ollama, + RawResponse: rawResponse, + } + + return result, nil +} diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 661591ee7f..6e1c0ac472 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -43,6 +43,7 @@ const ( Cohere ModelProvider = "cohere" Vertex ModelProvider = "vertex" Mistral ModelProvider = "mistral" + Ollama ModelProvider = "ollama" ) //* Request Structs diff --git a/core/schemas/provider.go b/core/schemas/provider.go index 646ea38d84..6427b0e821 100644 --- a/core/schemas/provider.go +++ b/core/schemas/provider.go @@ -27,7 +27,7 @@ const ( // NetworkConfig represents the network configuration for provider connections. type NetworkConfig struct { - // BaseURL is only supported for OpenAI, Anthropic and Cohere providers + // BaseURL is supported for OpenAI, Anthropic, Cohere, Mistral, and Ollama providers (required for Ollama) BaseURL string `json:"base_url,omitempty"` // Base URL for the provider (optional) DefaultRequestTimeoutInSeconds int `json:"default_request_timeout_in_seconds"` // Default timeout for requests MaxRetries int `json:"max_retries"` // Maximum number of retries diff --git a/core/tests/account.go b/core/tests/account.go index a46dab51d5..22aae0cad7 100644 --- a/core/tests/account.go +++ b/core/tests/account.go @@ -15,7 +15,7 @@ import ( // BaseAccount provides a test implementation of the Account interface. // It implements basic account functionality for testing purposes, supporting -// multiple AI providers including OpenAI, Anthropic, Bedrock, Cohere, and Azure. +// multiple AI providers including OpenAI, Anthropic, Bedrock, Cohere, Azure, Mistral, and Ollama. // The implementation uses environment variables from the .env file for API keys and provides // default configurations suitable for testing. type BaseAccount struct{} @@ -27,7 +27,16 @@ type BaseAccount struct{} // - []schemas.SupportedModelProvider: A slice containing the supported provider identifiers // - error: Always returns nil as this implementation doesn't produce errors func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) { - return []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic, schemas.Bedrock, schemas.Cohere, schemas.Azure, schemas.Vertex, schemas.Mistral}, nil + return []schemas.ModelProvider{ + schemas.OpenAI, + schemas.Anthropic, + schemas.Bedrock, + schemas.Cohere, + schemas.Azure, + schemas.Vertex, + schemas.Mistral, + schemas.Ollama, + }, nil } // GetKeysForProvider returns the API keys and associated models for a given provider. @@ -212,6 +221,20 @@ func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelPr NetworkConfig: schemas.DefaultNetworkConfig, ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize, }, nil + case schemas.Ollama: + return &schemas.ProviderConfig{ + NetworkConfig: schemas.NetworkConfig{ + BaseURL: os.Getenv("OLLAMA_BASE_URL"), + DefaultRequestTimeoutInSeconds: 30, + MaxRetries: 1, + RetryBackoffInitial: 100 * time.Millisecond, + RetryBackoffMax: 2 * time.Second, + }, + ConcurrencyAndBufferSize: schemas.ConcurrencyAndBufferSize{ + Concurrency: 3, + BufferSize: 10, + }, + }, nil default: return nil, fmt.Errorf("unsupported provider: %s", providerKey) } diff --git a/core/tests/mistral_test.go b/core/tests/mistral_test.go index 7e59d6fce7..7ae63d17bb 100644 --- a/core/tests/mistral_test.go +++ b/core/tests/mistral_test.go @@ -24,12 +24,6 @@ func TestMistral(t *testing.T) { SetupToolCalls: true, SetupImage: true, SetupBaseImage: true, - Fallbacks: []schemas.Fallback{ - { - Provider: schemas.Anthropic, - Model: "claude-3-7-sonnet-20250219", - }, - }, } SetupAllRequests(bifrost, config) diff --git a/core/tests/ollama_test.go b/core/tests/ollama_test.go new file mode 100644 index 0000000000..411c2abcaa --- /dev/null +++ b/core/tests/ollama_test.go @@ -0,0 +1,30 @@ +// Package tests provides test utilities and configurations for the Bifrost system. +// It includes test implementations of schemas, mock objects, and helper functions +// for testing the Bifrost functionality with various AI providers. +package tests + +import ( + "testing" + + schemas "github.com/maximhq/bifrost/core/schemas" +) + +func TestOllama(t *testing.T) { + bifrost, err := getBifrost() + if err != nil { + t.Fatalf("Error initializing bifrost: %v", err) + return + } + + config := TestConfig{ + Provider: schemas.Ollama, + TextModel: "llama3.2", + ChatModel: "llama3.2", + SetupText: false, // Ollama does not support text completion + SetupToolCalls: true, + SetupImage: true, + SetupBaseImage: true, + } + + SetupAllRequests(bifrost, config) +} diff --git a/docs/http-transport-api.md b/docs/http-transport-api.md index c013a7fb66..b1edf07cc9 100644 --- a/docs/http-transport-api.md +++ b/docs/http-transport-api.md @@ -227,14 +227,14 @@ Returns Prometheus metrics for monitoring and observability. The main request object for both chat and text completions. -| Field | Type | Required | Description | -| ----------- | ------------------------------------- | -------- | --------------------------------------------------------------------------------- | -| `provider` | `string` | ✅ | AI model provider (`openai`, `anthropic`, `azure`, `bedrock`, `cohere`, `vertex`) | -| `model` | `string` | ✅ | Model identifier (provider-specific) | -| `messages` | [`BifrostMessage[]`](#bifrostmessage) | ✅\* | Array of chat messages (required for chat completions) | -| `text` | `string` | ✅\* | Text prompt (required for text completions) | -| `params` | [`ModelParameters`](#modelparameters) | ❌ | Model parameters and configuration | -| `fallbacks` | [`Fallback[]`](#fallback) | ❌ | Fallback providers and models | +| Field | Type | Required | Description | +| ----------- | ------------------------------------- | -------- | ------------------------------------------------------------------------------------------------------ | +| `provider` | `string` | ✅ | AI model provider (`openai`, `anthropic`, `azure`, `bedrock`, `cohere`, `vertex`, `mistral`, `ollama`) | +| `model` | `string` | ✅ | Model identifier (provider-specific) | +| `messages` | [`BifrostMessage[]`](#bifrostmessage) | ✅\* | Array of chat messages (required for chat completions) | +| `text` | `string` | ✅\* | Text prompt (required for text completions) | +| `params` | [`ModelParameters`](#modelparameters) | ❌ | Model parameters and configuration | +| `fallbacks` | [`Fallback[]`](#fallback) | ❌ | Fallback providers and models | \*Either `messages` or `text` is required depending on the endpoint. @@ -451,6 +451,8 @@ Detailed error information. | AWS Bedrock | `bedrock` | | Cohere | `cohere` | | Google Vertex AI | `vertex` | +| Mistral | `mistral` | +| Ollama | `ollama` | ## Error Codes diff --git a/docs/openapi.json b/docs/openapi.json index f72a0ee464..6334df9243 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -2,7 +2,7 @@ "openapi": "3.0.3", "info": { "title": "Bifrost HTTP Transport API", - "description": "A unified HTTP API for accessing multiple AI model providers including OpenAI, Anthropic, Azure, Bedrock, Cohere, and Vertex AI. Bifrost provides standardized endpoints for text and chat completions with built-in fallback support and comprehensive monitoring.", + "description": "A unified HTTP API for accessing multiple AI model providers:\n\n• openai\n• anthropic\n• azure\n• bedrock\n• cohere\n• vertex\n• mistral\n• ollama\n\nBifrost provides standardized endpoints for text and chat completions with built-in fallback support and comprehensive monitoring.", "version": "1.0.0", "contact": { "name": "Bifrost API Support", @@ -432,7 +432,16 @@ }, "ModelProvider": { "type": "string", - "enum": ["openai", "anthropic", "azure", "bedrock", "cohere", "vertex"], + "enum": [ + "openai", + "anthropic", + "azure", + "bedrock", + "cohere", + "vertex", + "mistral", + "ollama" + ], "description": "AI model provider", "example": "openai" }, diff --git a/docs/providers.md b/docs/providers.md index 6e82476c8f..feba6600c7 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -12,6 +12,8 @@ Bifrost currently supports the following providers: - Bedrock - Cohere - Vertex +- Mistral +- Ollama ## 2. Provider Configuration diff --git a/transports/bifrost-http/main.go b/transports/bifrost-http/main.go index ed32e7cad1..0ebfb670b2 100644 --- a/transports/bifrost-http/main.go +++ b/transports/bifrost-http/main.go @@ -1,5 +1,5 @@ // Package http provides an HTTP service using FastHTTP that exposes endpoints -// for text and chat completions using various AI model providers (OpenAI, Anthropic, Bedrock, etc.). +// for text and chat completions using various AI model providers (OpenAI, Anthropic, Bedrock, Mistral, Ollama, etc.). // // The HTTP service provides two main endpoints: // - /v1/text/completions: For text completion requests