AIRouter/internal/adapter/ollama.go

package adapter

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

// newHTTPClient 创建带连接池优化的 HTTP 客户端，复用 TCP 连接以减少延迟。
// timeout 为整体请求超时，connectTimeout 为 TCP 连接阶段超时。
func newHTTPClient(timeout time.Duration) *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			MaxIdleConns:        100,
			MaxIdleConnsPerHost: 20,
			IdleConnTimeout:     90 * time.Second,
			MaxConnsPerHost:     0, // 不限制
			DialContext: (&net.Dialer{
				Timeout:   5 * time.Second,
				KeepAlive: 30 * time.Second,
			}).DialContext,
		},
	}
}

// OllamaAdapter implements ModelAdapter for Ollama inference engine.
type OllamaAdapter struct {
	endpoint   string
	httpClient *http.Client
}

// NewOllamaAdapter creates a new Ollama adapter.
func NewOllamaAdapter(endpoint string) *OllamaAdapter {
	return &OllamaAdapter{
		endpoint:   strings.TrimRight(endpoint, "/"),
		httpClient: newHTTPClient(120 * time.Second),
	}
}

func (a *OllamaAdapter) Name() string {
	return "ollama"
}

// ollamaChatRequest is the Ollama /api/chat request format.
type ollamaChatRequest struct {
	Model    string        `json:"model"`
	Messages []ollamaMsg   `json:"messages"`
	Stream   bool          `json:"stream"`
	Options  ollamaOptions `json:"options,omitempty"`
}

type ollamaMsg struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ollamaOptions struct {
	Temperature float64 `json:"temperature,omitempty"`
	TopP        float64 `json:"top_p,omitempty"`
	NumPredict  int     `json:"num_predict,omitempty"`
}

// ollamaChatResponse is the Ollama /api/chat non-streaming response.
type ollamaChatResponse struct {
	Model           string    `json:"model"`
	Message         ollamaMsg `json:"message"`
	Done            bool      `json:"done"`
	PromptEvalCount int       `json:"prompt_eval_count"`
	EvalCount       int       `json:"eval_count"`
}

// ollamaChatStreamResponse is a single chunk in Ollama streaming response.
type ollamaChatStreamResponse struct {
	Model           string    `json:"model"`
	Message         ollamaMsg `json:"message"`
	Done            bool      `json:"done"`
	PromptEvalCount int       `json:"prompt_eval_count,omitempty"`
	EvalCount       int       `json:"eval_count,omitempty"`
}

func (a *OllamaAdapter) ChatCompletion(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
	ollamaReq := a.buildRequest(req, false)

	body, err := json.Marshal(ollamaReq)
	if err != nil {
		return nil, fmt.Errorf("marshal ollama request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", a.endpoint+"/api/chat", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create ollama request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if req.RequestID != "" {
		httpReq.Header.Set("X-Request-ID", req.RequestID)
	}

	resp, err := a.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("ollama request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(bodyBytes))
	}

	var ollamaResp ollamaChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&ollamaResp); err != nil {
		return nil, fmt.Errorf("decode ollama response: %w", err)
	}

	return &ChatResponse{
		Content:      ollamaResp.Message.Content,
		FinishReason: "stop",
		InputTokens:  ollamaResp.PromptEvalCount,
		OutputTokens: ollamaResp.EvalCount,
		ActualModel:  ollamaResp.Model,
	}, nil
}

func (a *OllamaAdapter) ChatCompletionStream(ctx context.Context, req *ChatRequest) (<-chan StreamChunk, error) {
	ollamaReq := a.buildRequest(req, true)

	body, err := json.Marshal(ollamaReq)
	if err != nil {
		return nil, fmt.Errorf("marshal ollama stream request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", a.endpoint+"/api/chat", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create ollama stream request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if req.RequestID != "" {
		httpReq.Header.Set("X-Request-ID", req.RequestID)
	}

	resp, err := a.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("ollama stream request failed: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("ollama stream returned status %d: %s", resp.StatusCode, string(bodyBytes))
	}

	ch := make(chan StreamChunk, 100)
	go func() {
		defer close(ch)
		defer resp.Body.Close()

		decoder := json.NewDecoder(resp.Body)
		for {
			var chunk ollamaChatStreamResponse
			if err := decoder.Decode(&chunk); err != nil {
				if err == io.EOF {
					ch <- StreamChunk{Done: true, FinishReason: "stop"}
					return
				}
				ch <- StreamChunk{Error: fmt.Errorf("decode stream chunk: %w", err)}
				return
			}

			// Check for cancellation
			select {
			case <-req.CancelCh:
				ch <- StreamChunk{Done: true, FinishReason: "cancelled"}
				return
			default:
			}

			if chunk.Done {
				ch <- StreamChunk{
					Done:         true,
					FinishReason: "stop",
					InputTokens:  chunk.PromptEvalCount,
					OutputTokens: chunk.EvalCount,
				}
				return
			}

			if chunk.Message.Content != "" {
				ch <- StreamChunk{Delta: chunk.Message.Content}
			}
		}
	}()

	return ch, nil
}

func (a *OllamaAdapter) ListModels(ctx context.Context) ([]ModelInfo, error) {
	httpReq, err := http.NewRequestWithContext(ctx, "GET", a.endpoint+"/api/tags", nil)
	if err != nil {
		return nil, fmt.Errorf("create list models request: %w", err)
	}

	resp, err := a.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("list models failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("list models returned status %d", resp.StatusCode)
	}

	var tagsResp struct {
		Models []struct {
			Name string `json:"name"`
		} `json:"models"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&tagsResp); err != nil {
		return nil, fmt.Errorf("decode tags response: %w", err)
	}

	models := make([]ModelInfo, len(tagsResp.Models))
	for i, m := range tagsResp.Models {
		models[i] = ModelInfo{Name: m.Name}
	}
	return models, nil
}

func (a *OllamaAdapter) HealthCheck(ctx context.Context) error {
	httpReq, err := http.NewRequestWithContext(ctx, "GET", a.endpoint+"/api/tags", nil)
	if err != nil {
		return fmt.Errorf("create health check request: %w", err)
	}

	resp, err := a.httpClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("health check failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("health check returned status %d", resp.StatusCode)
	}
	return nil
}

func (a *OllamaAdapter) Cancel(requestID string) error {
	// Ollama doesn't support request cancellation by ID in the API.
	// Cancellation is handled by closing the HTTP connection (context cancellation).
	return nil
}

func (a *OllamaAdapter) buildRequest(req *ChatRequest, stream bool) ollamaChatRequest {
	msgs := make([]ollamaMsg, len(req.Messages))
	for i, m := range req.Messages {
		content, _ := m.Content.(string)
		msgs[i] = ollamaMsg{Role: m.Role, Content: content}
	}

	ollamaReq := ollamaChatRequest{
		Model:    req.Model,
		Messages: msgs,
		Stream:   stream,
	}

	if req.MaxTokens > 0 {
		ollamaReq.Options.NumPredict = req.MaxTokens
	}
	if req.Temperature != nil {
		ollamaReq.Options.Temperature = *req.Temperature
	}
	if req.TopP != nil {
		ollamaReq.Options.TopP = *req.TopP
	}

	return ollamaReq
}