package ollama import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "time" "tm/pkg/logger" ) // Client implements the HTTPClient interface for Ollama API type Client struct { config *Config httpClient *http.Client logger logger.Logger } // NewClient creates a new Ollama HTTP client func NewClient(config *Config, logger logger.Logger) (*Client, error) { if config == nil { config = DefaultConfig() } if err := config.Validate(); err != nil { return nil, err } httpClient := &http.Client{ Timeout: config.Timeout, } return &Client{ config: config, httpClient: httpClient, logger: logger, }, nil } // NewClientWithHTTPClient creates a new client with a custom HTTP client func NewClientWithHTTPClient(config *Config, httpClient *http.Client, logger logger.Logger) (*Client, error) { if config == nil { config = DefaultConfig() } if err := config.Validate(); err != nil { return nil, err } return &Client{ config: config, httpClient: httpClient, logger: logger, }, nil } // Post performs a POST request to the Ollama API func (c *Client) Post(ctx context.Context, url string, body interface{}, result interface{}) error { return c.makeRequest(ctx, "POST", url, body, result) } // Get performs a GET request to the Ollama API func (c *Client) Get(ctx context.Context, url string, result interface{}) error { return c.makeRequest(ctx, "GET", url, nil, result) } // Delete performs a DELETE request to the Ollama API func (c *Client) Delete(ctx context.Context, url string, body interface{}, result interface{}) error { return c.makeRequest(ctx, "DELETE", url, body, result) } // makeRequest performs an HTTP request with retry logic and logging func (c *Client) makeRequest(ctx context.Context, method, endpoint string, body interface{}, result interface{}) error { url := c.config.BaseURL + endpoint var bodyReader io.Reader if body != nil { bodyBytes, err := json.Marshal(body) if err != nil { return fmt.Errorf("failed to marshal request body: %w", err) } bodyReader = bytes.NewReader(bodyBytes) } var lastErr error for attempt := 0; attempt <= c.config.RetryAttempts; attempt++ { if attempt > 0 { // Wait before retry select { case <-ctx.Done(): return ctx.Err() case <-time.After(c.config.RetryDelay): } } err := c.doRequest(ctx, method, url, bodyReader, result) if err == nil { return nil } lastErr = err // Don't retry on certain errors if c.shouldNotRetry(err) { break } if c.logger != nil { c.logger.Warn("Request failed, retrying", map[string]interface{}{ "attempt": attempt + 1, "max_attempts": c.config.RetryAttempts + 1, "error": err.Error(), "url": url, "method": method, }) } } return fmt.Errorf("request failed after %d attempts: %w", c.config.RetryAttempts+1, lastErr) } // doRequest performs a single HTTP request func (c *Client) doRequest(ctx context.Context, method, url string, body io.Reader, result interface{}) error { req, err := http.NewRequestWithContext(ctx, method, url, body) if err != nil { return fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", c.config.UserAgent) if c.logger != nil && c.config.EnableLogging { c.logger.Debug("Making HTTP request", map[string]interface{}{ "method": method, "url": url, "user_agent": c.config.UserAgent, }) } resp, err := c.httpClient.Do(req) if err != nil { return fmt.Errorf("failed to execute request: %w", err) } defer resp.Body.Close() if c.logger != nil && c.config.EnableLogging { c.logger.Debug("Received HTTP response", map[string]interface{}{ "status_code": resp.StatusCode, "url": url, }) } if resp.StatusCode >= 400 { bodyBytes, _ := io.ReadAll(resp.Body) return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) } if result != nil { if err := json.NewDecoder(resp.Body).Decode(result); err != nil { return fmt.Errorf("failed to decode response: %w", err) } } return nil } // shouldNotRetry determines if an error should not be retried func (c *Client) shouldNotRetry(err error) bool { // Don't retry on context cancellation if err == context.Canceled || err == context.DeadlineExceeded { return true } // Don't retry on certain HTTP status codes if err.Error() == "HTTP 400" || err.Error() == "HTTP 401" || err.Error() == "HTTP 403" || err.Error() == "HTTP 404" { return true } return false } // StreamGenerate performs a streaming text generation request func (c *Client) StreamGenerate(ctx context.Context, req *GenerateRequest, onChunk func(*StreamChunk) error) error { url := c.config.BaseURL + "/api/generate" // Force streaming stream := true req.Stream = &stream bodyBytes, err := json.Marshal(req) if err != nil { return fmt.Errorf("failed to marshal request body: %w", err) } httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes)) if err != nil { return fmt.Errorf("failed to create request: %w", err) } httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("User-Agent", c.config.UserAgent) resp, err := c.httpClient.Do(httpReq) if err != nil { return fmt.Errorf("failed to execute request: %w", err) } defer resp.Body.Close() if resp.StatusCode >= 400 { bodyBytes, _ := io.ReadAll(resp.Body) return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) } decoder := json.NewDecoder(resp.Body) for { var chunk StreamChunk if err := decoder.Decode(&chunk); err != nil { if err == io.EOF { break } return fmt.Errorf("failed to decode stream chunk: %w", err) } if err := onChunk(&chunk); err != nil { return fmt.Errorf("error processing stream chunk: %w", err) } if chunk.Done { break } } return nil } // StreamChat performs a streaming chat request func (c *Client) StreamChat(ctx context.Context, req *ChatRequest, onChunk func(*StreamChunk) error) error { url := c.config.BaseURL + "/api/chat" // Force streaming stream := true req.Stream = &stream bodyBytes, err := json.Marshal(req) if err != nil { return fmt.Errorf("failed to marshal request body: %w", err) } httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes)) if err != nil { return fmt.Errorf("failed to create request: %w", err) } httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("User-Agent", c.config.UserAgent) resp, err := c.httpClient.Do(httpReq) if err != nil { return fmt.Errorf("failed to execute request: %w", err) } defer resp.Body.Close() if resp.StatusCode >= 400 { bodyBytes, _ := io.ReadAll(resp.Body) return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) } decoder := json.NewDecoder(resp.Body) for { var chunk StreamChunk if err := decoder.Decode(&chunk); err != nil { if err == io.EOF { break } return fmt.Errorf("failed to decode stream chunk: %w", err) } if err := onChunk(&chunk); err != nil { return fmt.Errorf("error processing stream chunk: %w", err) } if chunk.Done { break } } return nil }