Atlas/internal/service/chat/service.go

package chat

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/rose_cat707/Atlas/internal/domain"
	"github.com/rose_cat707/Atlas/internal/provider"
	"github.com/rose_cat707/Atlas/internal/store"
)

type Service struct {
	store  *store.Store
	router *provider.Router
	models *provider.ModelsService
}

func NewService(st *store.Store, router *provider.Router, models *provider.ModelsService) *Service {
	return &Service{store: st, router: router, models: models}
}

type StreamRequest struct {
	SessionID      string
	ConversationID int64
	Content        string
	Model          string
	ProviderID     int64
}

type streamCapture struct {
	io.Writer
	buf []byte
}

func (s *streamCapture) Write(p []byte) (int, error) {
	s.buf = append(s.buf, p...)
	return s.Writer.Write(p)
}

func (s *Service) StreamMessage(ctx context.Context, req StreamRequest, w io.Writer, flusher func()) error {
	conv, err := s.store.GetConversation(ctx, req.SessionID, req.ConversationID)
	if err != nil {
		return err
	}
	userMsg := &store.ChatMessageRow{ConversationID: req.ConversationID, Role: "user", Content: req.Content}
	if _, err := s.store.InsertMessage(ctx, userMsg); err != nil {
		return err
	}
	msgs, err := s.store.ListMessages(ctx, req.ConversationID, 40)
	if err != nil {
		return err
	}
	model := req.Model
	if model == "" {
		model = conv.Model
	}
	cfg, err := s.router.Resolve(ctx, string(domain.ServiceChat), req.ProviderID)
	if err != nil {
		return err
	}
	if model == "" {
		models, _ := s.models.ListEffective(ctx, cfg)
		model = provider.DefaultModel(models)
	}
	body, err := json.Marshal(map[string]any{
		"model":    model,
		"messages": toChatMessages(msgs),
		"stream":   true,
	})
	if err != nil {
		return err
	}
	capture := &streamCapture{Writer: w}
	meta := provider.CallMeta{
		SessionID: req.SessionID,
		Method:    "POST",
		Path:      "/v1/chat/completions",
		Model:     model,
	}
	if err := s.router.ProxyStream(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta, capture, flusher); err != nil {
		return err
	}
	assistant := extractAssistantFromSSE(string(capture.buf))
	return s.persistAssistant(ctx, req, model, cfg, assistant)
}

func (s *Service) persistAssistant(ctx context.Context, req StreamRequest, model string, cfg *provider.Config, assistant string) error {
	asst := &store.ChatMessageRow{ConversationID: req.ConversationID, Role: "assistant", Content: assistant}
	if _, err := s.store.InsertMessage(ctx, asst); err != nil {
		return err
	}
	_ = s.maybeSummarizeTitle(ctx, req, model, cfg)

	pid := cfg.ID
	out, _ := json.Marshal(map[string]any{"content": assistant})
	gen := &store.Generation{
		SessionID:   req.SessionID,
		ServiceType: string(domain.ServiceChat),
		ProviderID:  &pid,
		Model:       model,
		InputJSON:   fmt.Sprintf(`{"content":%q}`, req.Content),
		OutputJSON:  string(out),
	}
	s.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error {
		_, err := st.InsertGeneration(ctx, gen)
		return err
	})
	return nil
}

func (s *Service) maybeSummarizeTitle(ctx context.Context, req StreamRequest, model string, cfg *provider.Config) error {
	conv, err := s.store.GetConversation(ctx, req.SessionID, req.ConversationID)
	if err != nil {
		return err
	}
	params := conversationParams(conv.ParamsJSON)
	if titleAlreadySummarized(params) {
		return nil
	}
	userCount, err := s.store.CountUserMessages(ctx, req.ConversationID)
	if err != nil {
		return err
	}
	if userCount <= titleSummaryRoundThreshold {
		return nil
	}

	msgs, err := s.store.ListMessages(ctx, req.ConversationID, 24)
	if err != nil {
		return err
	}
	title, err := s.generateTitle(ctx, req, model, cfg, msgs)
	if err != nil || title == "" {
		return err
	}
	conv.Title = title
	if model != "" {
		conv.Model = model
	}
	params["title_summarized"] = true
	conv.ParamsJSON = conversationParamsJSON(params)
	return s.store.UpdateConversation(ctx, conv)
}

func (s *Service) generateTitle(ctx context.Context, req StreamRequest, model string, cfg *provider.Config, msgs []store.ChatMessageRow) (string, error) {
	chatMsgs := []map[string]string{{"role": "system", "content": titleSummarySystemPrompt}}
	chatMsgs = append(chatMsgs, toChatMessages(msgs)...)
	body, err := json.Marshal(map[string]any{
		"model":       model,
		"messages":    chatMsgs,
		"max_tokens":  40,
		"temperature": 0.3,
		"stream":      false,
	})
	if err != nil {
		return "", err
	}
	meta := provider.CallMeta{
		SessionID: req.SessionID,
		Method:    "POST",
		Path:      "/v1/chat/completions",
		Model:     model,
	}
	res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta)
	if err != nil {
		return "", err
	}
	if res.StatusCode >= 400 {
		return "", fmt.Errorf("title summarize upstream %d", res.StatusCode)
	}
	return sanitizeTitle(parseChatCompletionContent(res.Body)), nil
}

func extractAssistantFromSSE(raw string) string {
	var assistant strings.Builder
	for _, line := range strings.Split(raw, "\n") {
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			break
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}
		if len(chunk.Choices) > 0 {
			assistant.WriteString(chunk.Choices[0].Delta.Content)
		}
	}
	return assistant.String()
}

func toChatMessages(rows []store.ChatMessageRow) []map[string]string {
	out := make([]map[string]string, 0, len(rows))
	for _, m := range rows {
		out = append(out, map[string]string{"role": m.Role, "content": m.Content})
	}
	return out
}