package providers import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strings" "time" ) // OllamaClient implements the Provider interface for Ollama's local API type OllamaClient struct { model string baseURL string client *http.Client } // NewOllamaClient creates a new Ollama API client func NewOllamaClient(model, baseURL string) *OllamaClient { if baseURL == "" { baseURL = "http://localhost:11434" } // Normalize the URL: strip trailing slashes and /api suffix // Users sometimes enter http://host:11434/ or http://host:11434/api baseURL = strings.TrimSuffix(baseURL, "/") baseURL = strings.TrimSuffix(baseURL, "/api") baseURL = strings.TrimSuffix(baseURL, "/") // In case it was /api/ return &OllamaClient{ model: model, baseURL: baseURL, client: &http.Client{ Timeout: 300 * time.Second, // Local models can be slow }, } } // Name returns the provider name func (c *OllamaClient) Name() string { return "ollama" } // ollamaRequest is the request body for the Ollama API type ollamaRequest struct { Model string `json:"model"` Messages []ollamaMessage `json:"messages"` Stream bool `json:"stream"` Options *ollamaOptions `json:"options,omitempty"` } type ollamaMessage struct { Role string `json:"role"` Content string `json:"content"` } type ollamaOptions struct { NumPredict int `json:"num_predict,omitempty"` Temperature float64 `json:"temperature,omitempty"` } // ollamaResponse is the response from the Ollama API type ollamaResponse struct { Model string `json:"model"` CreatedAt string `json:"created_at"` Message ollamaMessage `json:"message"` Done bool `json:"done"` DoneReason string `json:"done_reason,omitempty"` TotalDuration int64 `json:"total_duration,omitempty"` LoadDuration int64 `json:"load_duration,omitempty"` PromptEvalCount int `json:"prompt_eval_count,omitempty"` EvalCount int `json:"eval_count,omitempty"` } // Chat sends a chat request to the Ollama API func (c *OllamaClient) Chat(ctx context.Context, req ChatRequest) (*ChatResponse, error) { // Convert messages to Ollama format messages := make([]ollamaMessage, 0, len(req.Messages)+1) // Add system message if provided if req.System != "" { messages = append(messages, ollamaMessage{ Role: "system", Content: req.System, }) } for _, m := range req.Messages { messages = append(messages, ollamaMessage{ Role: m.Role, Content: m.Content, }) } // Use provided model or fall back to client default model := req.Model // Strip "ollama:" prefix if present - callers may pass the full "provider:model" string if strings.HasPrefix(model, "ollama:") { model = strings.TrimPrefix(model, "ollama:") } if model == "" { model = c.model } // Ultimate fallback - if no model configured anywhere, use llama3 if model == "" { model = "llama3" } ollamaReq := ollamaRequest{ Model: model, Messages: messages, Stream: false, // Non-streaming for now } if req.MaxTokens > 0 || req.Temperature > 0 { ollamaReq.Options = &ollamaOptions{} if req.MaxTokens > 0 { ollamaReq.Options.NumPredict = req.MaxTokens } if req.Temperature > 0 { ollamaReq.Options.Temperature = req.Temperature } } body, err := json.Marshal(ollamaReq) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } url := c.baseURL + "/api/chat" httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } httpReq.Header.Set("Content-Type", "application/json") resp, err := c.client.Do(httpReq) if err != nil { return nil, fmt.Errorf("request failed: %w", err) } defer resp.Body.Close() respBody, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("failed to read response: %w", err) } if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("API error (%d): %s", resp.StatusCode, string(respBody)) } var ollamaResp ollamaResponse if err := json.Unmarshal(respBody, &ollamaResp); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } return &ChatResponse{ Content: ollamaResp.Message.Content, Model: ollamaResp.Model, StopReason: ollamaResp.DoneReason, InputTokens: ollamaResp.PromptEvalCount, OutputTokens: ollamaResp.EvalCount, }, nil } // TestConnection validates connectivity by checking the Ollama version endpoint func (c *OllamaClient) TestConnection(ctx context.Context) error { url := c.baseURL + "/api/version" httpReq, err := http.NewRequestWithContext(ctx, "GET", url, nil) if err != nil { return fmt.Errorf("failed to create request: %w", err) } resp, err := c.client.Do(httpReq) if err != nil { return fmt.Errorf("failed to connect to Ollama at %s: %w", c.baseURL, err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return fmt.Errorf("Ollama returned status %d", resp.StatusCode) } return nil } // ListModels fetches available models from the local Ollama instance func (c *OllamaClient) ListModels(ctx context.Context) ([]ModelInfo, error) { url := c.baseURL + "/api/tags" httpReq, err := http.NewRequestWithContext(ctx, "GET", url, nil) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } resp, err := c.client.Do(httpReq) if err != nil { return nil, fmt.Errorf("request failed: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) return nil, fmt.Errorf("API error (%d): %s", resp.StatusCode, string(body)) } var result struct { Models []struct { Name string `json:"name"` ModifiedAt string `json:"modified_at"` Size int64 `json:"size"` } `json:"models"` } if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } models := make([]ModelInfo, 0, len(result.Models)) for _, m := range result.Models { models = append(models, ModelInfo{ ID: m.Name, Name: m.Name, }) } return models, nil }