ZovOS_AI/internal/ai/ollama/ollama.go

package ollama

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/config"
)

type Provider struct {
	baseURL string
	model   string
	client  *http.Client
	log     *slog.Logger
}

func New(cfg config.OllamaConfig, log *slog.Logger) *Provider {
	return &Provider{
		baseURL: cfg.BaseURL,
		model:   cfg.Model,
		client:  &http.Client{},
		log:     log,
	}
}

func (p *Provider) Name() string {
	return "ollama"
}

func (p *Provider) Models() []string {
	return []string{p.model}
}

// ollamaMessage is Ollama's native message format.
type ollamaMessage struct {
	Role      string         `json:"role"`
	Content   string         `json:"content"`
	ToolCalls []ollamaToolCall `json:"tool_calls,omitempty"`
}

type ollamaToolCall struct {
	Function struct {
		Name      string                 `json:"name"`
		Arguments map[string]interface{} `json:"arguments"`
	} `json:"function"`
}

type ollamaTool struct {
	Type     string `json:"type"`
	Function struct {
		Name        string                 `json:"name"`
		Description string                 `json:"description"`
		Parameters  map[string]interface{} `json:"parameters"`
	} `json:"function"`
}

type chatRequest struct {
	Model    string          `json:"model"`
	Messages []ollamaMessage `json:"messages"`
	Stream   bool            `json:"stream"`
	Tools    []ollamaTool    `json:"tools,omitempty"`
}

type chatResponse struct {
	Message ollamaMessage `json:"message"`
	Done    bool          `json:"done"`
	Model   string        `json:"model"`
}

func (p *Provider) Chat(ctx context.Context, messages []ai.Message, tools []ai.Tool) (*ai.Response, error) {
	ollamaMsgs := convertMessages(messages)
	ollamaTools := convertTools(tools)

	reqBody := chatRequest{
		Model:    p.model,
		Messages: ollamaMsgs,
		Stream:   false,
		Tools:    ollamaTools,
	}

	data, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshaling request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL+"/api/chat", bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("creating request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := p.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("sending request to Ollama: %w", err)
	}
	defer resp.Body.Close()

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

	var chatResp chatResponse
	if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
		return nil, fmt.Errorf("decoding response: %w", err)
	}

	result := &ai.Response{
		Content: chatResp.Message.Content,
		Model:   chatResp.Model,
	}

	for _, tc := range chatResp.Message.ToolCalls {
		argsJSON, _ := json.Marshal(tc.Function.Arguments)
		result.ToolCalls = append(result.ToolCalls, ai.ToolCall{
			ID: tc.Function.Name,
			Function: struct {
				Name      string `json:"name"`
				Arguments string `json:"arguments"`
			}{
				Name:      tc.Function.Name,
				Arguments: string(argsJSON),
			},
		})
	}

	return result, nil
}

func (p *Provider) StreamChat(ctx context.Context, messages []ai.Message, tools []ai.Tool) (<-chan ai.StreamChunk, error) {
	ollamaMsgs := convertMessages(messages)
	ollamaTools := convertTools(tools)

	reqBody := chatRequest{
		Model:    p.model,
		Messages: ollamaMsgs,
		Stream:   true,
		Tools:    ollamaTools,
	}

	data, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshaling request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL+"/api/chat", bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("creating request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := p.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("sending request to Ollama: %w", err)
	}

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

	ch := make(chan ai.StreamChunk, 64)

	go func() {
		defer close(ch)
		defer resp.Body.Close()

		decoder := json.NewDecoder(resp.Body)
		for {
			var chunk chatResponse
			if err := decoder.Decode(&chunk); err != nil {
				if err != io.EOF {
					ch <- ai.StreamChunk{Error: err}
				}
				return
			}

			streamChunk := ai.StreamChunk{
				Content: chunk.Message.Content,
				Done:    chunk.Done,
			}

			for _, tc := range chunk.Message.ToolCalls {
				argsJSON, _ := json.Marshal(tc.Function.Arguments)
				streamChunk.ToolCalls = append(streamChunk.ToolCalls, ai.ToolCall{
					ID: tc.Function.Name,
					Function: struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					}{
						Name:      tc.Function.Name,
						Arguments: string(argsJSON),
					},
				})
			}

			select {
			case ch <- streamChunk:
			case <-ctx.Done():
				return
			}

			if chunk.Done {
				return
			}
		}
	}()

	return ch, nil
}

func convertMessages(msgs []ai.Message) []ollamaMessage {
	result := make([]ollamaMessage, len(msgs))
	for i, m := range msgs {
		result[i] = ollamaMessage{
			Role:    m.Role,
			Content: m.Content,
		}
	}
	return result
}

func convertTools(tools []ai.Tool) []ollamaTool {
	result := make([]ollamaTool, len(tools))
	for i, t := range tools {
		result[i] = ollamaTool{Type: "function"}
		result[i].Function.Name = t.Name
		result[i].Function.Description = t.Description
		result[i].Function.Parameters = t.Parameters
	}
	return result
}