feat: add Ollama provider support

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,15 +203,15 @@ 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
 - **Concurrency Control**: Manage rate limits and parallel requests effectively
 - **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):
 

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))

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
+}

core/schemas/bifrost.go

@@ -43,6 +43,7 @@ const (
 	Cohere    ModelProvider = "cohere"
 	Vertex    ModelProvider = "vertex"
 	Mistral   ModelProvider = "mistral"
+	Ollama    ModelProvider = "ollama"
 )
 
 //* Request Structs

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

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)
 	}

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)

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)
+}