orca.ai/pkg/llm/deepseek.go

package llm

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

type DeepSeekClient struct {
	baseURL    string
	model      string
	apiKey     string
	httpClient *http.Client
}

type DeepSeekOption func(*DeepSeekClient)

func WithDeepSeekBaseURL(url string) DeepSeekOption {
	return func(c *DeepSeekClient) {
		c.baseURL = strings.TrimRight(url, "/")
	}
}

func WithDeepSeekModel(model string) DeepSeekOption {
	return func(c *DeepSeekClient) {
		c.model = model
	}
}

func WithDeepSeekAPIKey(key string) DeepSeekOption {
	return func(c *DeepSeekClient) {
		c.apiKey = key
	}
}

func WithDeepSeekTimeout(timeout time.Duration) DeepSeekOption {
	return func(c *DeepSeekClient) {
		c.httpClient.Timeout = timeout
	}
}

func NewDeepSeekClient(opts ...DeepSeekOption) *DeepSeekClient {
	c := &DeepSeekClient{
		baseURL: "https://api.deepseek.com/v1",
		model:   "deepseek-chat",
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

func (c *DeepSeekClient) Chat(ctx context.Context, messages []Message) (*Response, error) {
	reqBody := c.buildChatRequest(messages, false)
	body, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("deepseek: failed to marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("deepseek: failed to create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.apiKey)

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

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

	var apiResp deepSeekChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
		return nil, fmt.Errorf("deepseek: failed to decode response: %w", err)
	}

	if len(apiResp.Choices) == 0 {
		return nil, fmt.Errorf("deepseek: no choices in response")
	}

	choice := apiResp.Choices[0]
	return &Response{
		Content:   choice.Message.Content,
		ToolCalls: choice.Message.ToolCalls,
	}, nil
}

func (c *DeepSeekClient) Stream(ctx context.Context, messages []Message, handler StreamHandler) error {
	reqBody := c.buildChatRequest(messages, true)
	body, err := json.Marshal(reqBody)
	if err != nil {
		return fmt.Errorf("deepseek: failed to marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("deepseek: failed to create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Accept", "text/event-stream")

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

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

	reader := bufio.NewReader(resp.Body)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("deepseek: error reading stream: %w", err)
		}

		line = strings.TrimSpace(line)
		if line == "" || line == "data: [DONE]" {
			continue
		}
		if !strings.HasPrefix(line, "data: ") {
			continue
		}

		data := strings.TrimPrefix(line, "data: ")
		var chunk deepSeekStreamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}

		if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
			if err := handler(chunk.Choices[0].Delta.Content); err != nil {
				return err
			}
		}
	}

	return nil
}

func (c *DeepSeekClient) buildChatRequest(messages []Message, stream bool) deepSeekChatRequest {
	return deepSeekChatRequest{
		Model:    c.model,
		Messages: messages,
		Stream:   stream,
	}
}

type deepSeekChatRequest struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
	Stream   bool      `json:"stream"`
}

type deepSeekChatResponse struct {
	Choices []struct {
		Message struct {
			Role      string     `json:"role"`
			Content   string     `json:"content"`
			ToolCalls []ToolCall `json:"tool_calls,omitempty"`
		} `json:"message"`
	} `json:"choices"`
}

type deepSeekStreamChunk struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
	} `json:"choices"`
}