digital-psychology/apps/api/internal/llm/deepseek/client.go

// Package deepseek calls DeepSeek OpenAI-compatible Chat Completions API.
package deepseek

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

	"github.com/yuxingu/digital-psychology/apps/api/internal/config"
)

// Message is one chat turn.
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// Client talks to DeepSeek.
type Client struct {
	cfg    config.DeepSeekConfig
	http   *http.Client
}

// New builds a client from config.
func New(cfg config.DeepSeekConfig) *Client {
	sec := cfg.TimeoutSec
	if sec <= 0 {
		sec = 60
	}
	return &Client{
		cfg: cfg,
		http: &http.Client{
			Timeout: time.Duration(sec) * time.Second,
		},
	}
}

// Enabled mirrors config.
func (c *Client) Enabled() bool {
	return c != nil && c.cfg.Enabled()
}

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

type chatResponse struct {
	Choices []struct {
		Message Message `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

type streamChunk struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// Chat sends messages and returns assistant text.
func (c *Client) Chat(ctx context.Context, messages []Message) (string, error) {
	if !c.Enabled() {
		return "", fmt.Errorf("deepseek api_key not configured")
	}
	base := c.cfg.BaseURL
	if base == "" {
		base = "https://api.deepseek.com"
	}
	model := c.cfg.Model
	if model == "" {
		model = "deepseek-chat"
	}

	body, err := json.Marshal(chatRequest{Model: model, Messages: messages, Stream: false})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.cfg.APIKey)

	res, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(res.Body, 2<<20))
	if err != nil {
		return "", err
	}
	var out chatResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return "", fmt.Errorf("deepseek decode: %w body=%s", err, truncate(string(raw), 200))
	}
	if out.Error != nil && out.Error.Message != "" {
		return "", fmt.Errorf("deepseek: %s", out.Error.Message)
	}
	if res.StatusCode >= 300 {
		return "", fmt.Errorf("deepseek HTTP %d: %s", res.StatusCode, truncate(string(raw), 200))
	}
	if len(out.Choices) == 0 || strings.TrimSpace(out.Choices[0].Message.Content) == "" {
		return "", fmt.Errorf("deepseek empty choices")
	}
	return strings.TrimSpace(out.Choices[0].Message.Content), nil
}

// ChatStream streams assistant text; onDelta is called for each content piece.
// Returns the full assembled reply.
func (c *Client) ChatStream(ctx context.Context, messages []Message, onDelta func(delta string) error) (string, error) {
	if !c.Enabled() {
		return "", fmt.Errorf("deepseek api_key not configured")
	}
	if onDelta == nil {
		onDelta = func(string) error { return nil }
	}
	base := c.cfg.BaseURL
	if base == "" {
		base = "https://api.deepseek.com"
	}
	model := c.cfg.Model
	if model == "" {
		model = "deepseek-chat"
	}

	body, err := json.Marshal(chatRequest{Model: model, Messages: messages, Stream: true})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.cfg.APIKey)
	req.Header.Set("Accept", "text/event-stream")

	// Stream may outlive the default client timeout; use a dedicated client.
	sec := c.cfg.TimeoutSec
	if sec <= 0 {
		sec = 90
	}
	httpClient := &http.Client{Timeout: time.Duration(sec) * time.Second}
	res, err := httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	if res.StatusCode >= 300 {
		raw, _ := io.ReadAll(io.LimitReader(res.Body, 2<<20))
		return "", fmt.Errorf("deepseek HTTP %d: %s", res.StatusCode, truncate(string(raw), 200))
	}

	reader := bufio.NewReader(res.Body)
	var full strings.Builder
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return full.String(), err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" || strings.HasPrefix(line, ":") {
			continue
		}
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "[DONE]" {
			break
		}
		var chunk streamChunk
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			continue
		}
		if chunk.Error != nil && chunk.Error.Message != "" {
			return full.String(), fmt.Errorf("deepseek: %s", chunk.Error.Message)
		}
		if len(chunk.Choices) == 0 {
			continue
		}
		delta := chunk.Choices[0].Delta.Content
		if delta == "" {
			continue
		}
		full.WriteString(delta)
		if err := onDelta(delta); err != nil {
			return full.String(), err
		}
	}
	out := strings.TrimSpace(full.String())
	if out == "" {
		return "", fmt.Errorf("deepseek empty stream")
	}
	return out, nil
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}