ZovOS_AI/internal/ai/anthropic/anthropic.go

package anthropic

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

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

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

func New(cfg config.AnthropicConfig, log *slog.Logger) *Provider {
	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "https://api.anthropic.com"
	}
	return &Provider{
		apiKey:  cfg.APIKey,
		model:   cfg.Model,
		baseURL: strings.TrimRight(baseURL, "/"),
		client:  &http.Client{},
		log:     log,
	}
}

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

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

type claudeRequest struct {
	Model     string         `json:"model"`
	MaxTokens int            `json:"max_tokens"`
	System    string         `json:"system,omitempty"`
	Messages  []claudeMsg    `json:"messages"`
	Tools     []claudeTool   `json:"tools,omitempty"`
	Stream    bool           `json:"stream"`
}

type claudeMsg struct {
	Role    string        `json:"role"`
	Content []contentBlock `json:"content"`
}

type contentBlock struct {
	Type      string `json:"type"`
	Text      string `json:"text,omitempty"`
	ID        string `json:"id,omitempty"`
	Name      string `json:"name,omitempty"`
	Input     json.RawMessage `json:"input,omitempty"`
	ToolUseID string `json:"tool_use_id,omitempty"`
	Content   string `json:"content,omitempty"`
}

type claudeTool struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description"`
	InputSchema map[string]interface{} `json:"input_schema"`
}

type claudeResponse struct {
	Content []contentBlock `json:"content"`
	Model   string         `json:"model"`
	Usage   struct {
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
	} `json:"usage"`
	StopReason string `json:"stop_reason"`
}

type claudeStreamEvent struct {
	Type  string `json:"type"`
	Index int    `json:"index"`
	Delta struct {
		Type string `json:"type"`
		Text string `json:"text"`
	} `json:"delta"`
	ContentBlock *contentBlock `json:"content_block"`
}

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

	reqBody := claudeRequest{
		Model:     p.model,
		MaxTokens: 4096,
		System:    system,
		Messages:  claudeMsgs,
		Tools:     convertTools(tools),
		Stream:    false,
	}

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

	req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL+"/v1/messages", bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("creating request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", p.apiKey)
	req.Header.Set("anthropic-version", "2023-06-01")

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

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

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

	result := &ai.Response{
		Model: cResp.Model,
		Usage: ai.Usage{
			PromptTokens:     cResp.Usage.InputTokens,
			CompletionTokens: cResp.Usage.OutputTokens,
			TotalTokens:      cResp.Usage.InputTokens + cResp.Usage.OutputTokens,
		},
	}

	for _, block := range cResp.Content {
		switch block.Type {
		case "text":
			result.Content += block.Text
		case "tool_use":
			result.ToolCalls = append(result.ToolCalls, ai.ToolCall{
				ID: block.ID,
				Function: struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				}{
					Name:      block.Name,
					Arguments: string(block.Input),
				},
			})
		}
	}

	return result, nil
}

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

	reqBody := claudeRequest{
		Model:     p.model,
		MaxTokens: 4096,
		System:    system,
		Messages:  claudeMsgs,
		Tools:     convertTools(tools),
		Stream:    true,
	}

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

	req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL+"/v1/messages", bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("creating request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", p.apiKey)
	req.Header.Set("anthropic-version", "2023-06-01")

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

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

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

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

		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			line := scanner.Text()
			if !strings.HasPrefix(line, "data: ") {
				continue
			}
			payload := strings.TrimPrefix(line, "data: ")

			var event claudeStreamEvent
			if err := json.Unmarshal([]byte(payload), &event); err != nil {
				continue
			}

			switch event.Type {
			case "content_block_delta":
				if event.Delta.Type == "text_delta" {
					select {
					case ch <- ai.StreamChunk{Content: event.Delta.Text}:
					case <-ctx.Done():
						return
					}
				}
			case "message_stop":
				ch <- ai.StreamChunk{Done: true}
				return
			}
		}
	}()

	return ch, nil
}

func convertMessages(msgs []ai.Message) (string, []claudeMsg) {
	var system string
	var result []claudeMsg

	for _, m := range msgs {
		if m.Role == "system" {
			system = m.Content
			continue
		}

		role := m.Role
		if role == "tool" {
			role = "user"
			result = append(result, claudeMsg{
				Role: role,
				Content: []contentBlock{{
					Type:      "tool_result",
					ToolUseID: m.ToolCallID,
					Content:   m.Content,
				}},
			})
			continue
		}

		result = append(result, claudeMsg{
			Role: role,
			Content: []contentBlock{{
				Type: "text",
				Text: m.Content,
			}},
		})
	}

	return system, result
}

func convertTools(tools []ai.Tool) []claudeTool {
	result := make([]claudeTool, len(tools))
	for i, t := range tools {
		result[i] = claudeTool{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: t.Parameters,
		}
	}
	return result
}