X-Agents/server/internal/service/agent_service.go

package service

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

	"x-agents/server/internal/repository"

	"github.com/gin-gonic/gin"
)

// AgentChatRequest Python Agent 对话请求
type AgentChatRequest struct {
	AgentID        int    `json:"agent_id"`
	Message        string `json:"message"`
	UserID         int    `json:"user_id"`
	SessionID      string `json:"session_id,omitempty"`
	ModelID        string `json:"model_id,omitempty"`
	ModelName      string `json:"model_name,omitempty"`
	ModelProvider  string `json:"model_provider,omitempty"`
	APIKey         string `json:"api_key,omitempty"`
	BaseURL        string `json:"base_url,omitempty"`
	UseXBot        bool   `json:"use_xbot"`
}

// AgentChatResponse Python Agent 对话响应
type AgentChatResponse struct {
	AgentID    int             `json:"agent_id"`
	Response   string          `json:"response"`
	ToolCalls  []interface{}   `json:"tool_calls"`
	TokensUsed int            `json:"tokens_used"`
	DurationMs int            `json:"duration_ms"`
	SessionID  string         `json:"session_id"`
}

// TeamChatRequest 多智能体群聊请求
type TeamChatRequest struct {
	SupervisorAgentID int      `json:"supervisor_agent_id"`
	MemberAgentIDs    []int    `json:"member_agent_ids"`
	Message          string   `json:"message"`
	UserID           int      `json:"user_id"`
	SessionID        string   `json:"session_id,omitempty"`
	Strategy         string   `json:"strategy,omitempty"`
}

// TeamChatResponse 多智能体群聊响应
type TeamChatResponse struct {
	SupervisorAgentID int             `json:"supervisor_agent_id"`
	Response         string          `json:"response"`
	SubtaskResults   []interface{}   `json:"subtask_results"`
	Strategy         string          `json:"strategy"`
	DurationMs       int             `json:"duration_ms"`
	SessionID        string          `json:"session_id"`
}

// AgentService Python Agent 服务
type AgentService struct {
	pythonURL string
	client    *http.Client
	modelRepo *repository.ModelRepository
}

// NewAgentService 创建 Agent 服务
func NewAgentService(pythonURL string, modelRepo *repository.ModelRepository) *AgentService {
	return &AgentService{
		pythonURL: pythonURL,
		client: &http.Client{
			Timeout: 120 * time.Second, // Agent 可能需要较长时间
		},
		modelRepo: modelRepo,
	}
}

// Chat 单智能体对话
func (s *AgentService) Chat(req AgentChatRequest) (*AgentChatResponse, error) {
	// 如果传入了 model_id，查询模型配置获取 api_key 和 base_url
	log.Printf("[AgentService] Chat called, model_id: %s, modelRepo: %v", req.ModelID, s.modelRepo != nil)

	if req.ModelID != "" && s.modelRepo != nil {
		model, err := s.modelRepo.FindByID(req.ModelID)
		if err != nil {
			log.Printf("[AgentService] Error finding model: %v", err)
		} else if model != nil {
			log.Printf("[AgentService] Found model: id=%s, provider=%s, model=%s, base_url=%s, api_key_len=%d",
				model.ID, model.Provider, model.Model, model.BaseURL, len(model.APIKey))
			req.APIKey = model.APIKey
			req.BaseURL = model.BaseURL
			req.ModelProvider = model.Provider
			req.ModelName = model.Model
			log.Printf("[AgentService] Set req.APIKey=%s, req.BaseURL=%s", req.APIKey[:10]+"...", req.BaseURL)
		} else {
			log.Printf("[AgentService] Model not found for id: %s", req.ModelID)
		}
	} else if s.modelRepo == nil {
		log.Printf("[AgentService] WARNING: modelRepo is nil!")
	}

	// 打印传给 Python 的请求内容
	apiKeyPreview := ""
	if req.APIKey != "" {
		apiKeyPreview = req.APIKey
		if len(apiKeyPreview) > 10 {
			apiKeyPreview = apiKeyPreview[:10] + "..."
		}
	}
	log.Printf("[AgentService] Sending to Python: model_id=%s, api_key=%s, base_url=%s, provider=%s, model=%s",
		req.ModelID, apiKeyPreview, req.BaseURL, req.ModelProvider, req.ModelName)

	url := fmt.Sprintf("%s/agent/chat", s.pythonURL)

	jsonData, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := s.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("failed to call python agent: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("python agent error: %s", string(body))
	}

	var result AgentChatResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
	}

	return &result, nil
}

// TeamChat 多智能体群聊
func (s *AgentService) TeamChat(req TeamChatRequest) (*TeamChatResponse, error) {
	url := fmt.Sprintf("%s/agent/team/chat", s.pythonURL)

	// 设置默认策略
	if req.Strategy == "" {
		req.Strategy = "parallel"
	}

	jsonData, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := s.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("failed to call python agent: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("python agent error: %s", string(body))
	}

	var result TeamChatResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
	}

	return &result, nil
}

// ChatStream 流式对话
func (s *AgentService) ChatStream(c interface{}, agentID int, message, sessionID, modelID string, userID int) error {
	// 获取 gin.Context
	ginCtx, ok := c.(*gin.Context)
	if !ok {
		return fmt.Errorf("invalid context type")
	}

	// 初始化请求体
	reqBody := map[string]interface{}{
		"agent_id":   agentID,
		"message":    message,
		"user_id":    userID,
		"session_id": sessionID,
		"use_xbot":   false,
	}

	// 如果传入了 model_id，查询模型配置获取 api_key 和 base_url
	if modelID != "" && s.modelRepo != nil {
		model, err := s.modelRepo.FindByID(modelID)
		if err != nil {
			log.Printf("[ChatStream] Model not found: %s, error: %v", modelID, err)
		} else if model != nil {
			log.Printf("[ChatStream] Using model: provider=%s, model=%s, base_url=%s", model.Provider, model.Model, model.BaseURL)
			// 将模型配置添加到请求体
			reqBody["model_provider"] = model.Provider
			reqBody["model_name"] = model.Model
			reqBody["api_key"] = model.APIKey
			reqBody["base_url"] = model.BaseURL
		}
	} else {
		log.Printf("[ChatStream] modelID is empty or modelRepo is nil: modelID=%s, modelRepo=%v", modelID, s.modelRepo != nil)
	}

	streamURL := fmt.Sprintf("%s/agent/chat/stream", s.pythonURL)

	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return fmt.Errorf("failed to marshal request: %w", err)
	}

	// 创建 HTTP 请求，设置不缓冲
	httpReq, err := http.NewRequest("POST", streamURL, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "text/event-stream")
	httpReq.Header.Set("Cache-Control", "no-cache")

	// 创建不缓冲的 HTTP 客户端
	client := &http.Client{
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		return fmt.Errorf("failed to call python agent: %w", err)
	}
	defer resp.Body.Close()

	// 设置 SSE 响应头
	ginCtx.Header("Content-Type", "text/event-stream")
	ginCtx.Header("Cache-Control", "no-cache")
	ginCtx.Header("Connection", "keep-alive")
	ginCtx.Header("X-Accel-Buffering", "no")

	// 分块读取并转发，使用小 buffer 减少延迟
	buf := make([]byte, 1024)
	for {
		n, err := resp.Body.Read(buf)
		if n > 0 {
			_, writeErr := ginCtx.Writer.Write(buf[:n])
			if writeErr != nil {
				break
			}
			// 强制刷新到客户端
			if flusher, ok := ginCtx.Writer.(interface{ Flush() }); ok {
				flusher.Flush()
			}
		}
		if err != nil {
			break
		}
	}

	return nil
}