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
	}
	if conv.Title == "新对话" && len(req.Content) > 0 {
		title := req.Content
		if len(title) > 30 {
			title = title[:30]
		}
		conv.Title = title
		_ = s.store.UpdateConversation(ctx, conv)
	}
	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.ID, assistant)
}

func (s *Service) persistAssistant(ctx context.Context, req StreamRequest, model string, providerID int64, assistant string) error {
	asst := &store.ChatMessageRow{ConversationID: req.ConversationID, Role: "assistant", Content: assistant}
	s.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error {
		_, err := st.InsertMessage(ctx, asst)
		return err
	})
	pid := providerID
	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 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
}