tm_back/pkg/ollama/client.go

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
}