shudao-main/shudao-go-backend/controllers/chroma.go

package controllers

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"shudao-chat-go/utils"
	"strconv"
	"time"

	"github.com/beego/beego/v2/server/web"
)

// chromaSearchHeartbeatTask 心跳任务
type chromaSearchHeartbeatTask struct {
	url        string
	interval   time.Duration
	httpClient *http.Client
	stopChan   chan struct{}
}

// ChromaController 知识库搜索控制器
type ChromaController struct {
	web.Controller
}

// SearchRequest 搜索请求结构体
type SearchRequest struct {
	QueryStr string `json:"query_str"`
	NResults int    `json:"n_results"`
}

// SearchResponse 搜索响应结构体
type SearchResponse struct {
	Total int        `json:"total"`
	Files []FileInfo `json:"files"`
}

// FileInfo 文件信息结构体
type FileInfo struct {
	DocumentID string `json:"document_id"`
	FileName   string `json:"file_name"`
	FilePath   string `json:"file_path"`
	SourceFile string `json:"source_file"`
	HasTitle   bool   `json:"has_title"`
}

// AdvancedSearch 知识库文件高级搜索接口
func (c *ChromaController) AdvancedSearch() {
	// 获取查询参数
	queryStr := c.GetString("query_str")
	nResultsStr := c.GetString("n_results", "50")

	// 参数验证
	if queryStr == "" {
		c.Data["json"] = map[string]interface{}{
			"code":    400,
			"message": "查询字符串不能为空",
			"data":    nil,
		}
		c.ServeJSON()
		return
	}

	// 转换n_results参数
	nResults, err := strconv.Atoi(nResultsStr)
	if err != nil || nResults <= 0 {
		nResults = 50 // 默认值
	}

	// 构建请求参数
	searchReq := SearchRequest{
		QueryStr: queryStr,
		NResults: nResults,
	}

	// 调用外部API
	result, err := c.callAdvancedSearchAPI(searchReq)
	if err != nil {
		c.Data["json"] = map[string]interface{}{
			"code":    500,
			"message": fmt.Sprintf("搜索失败: %v", err),
			"data":    nil,
		}
		c.ServeJSON()
		return
	}
	fmt.Println("知识库:", result)
	// 返回成功结果
	c.Data["json"] = map[string]interface{}{
		"code":    200,
		"message": "搜索成功",
		"data":    result,
	}
	c.ServeJSON()
}

// callAdvancedSearchAPI 调用外部高级搜索API
func (c *ChromaController) callAdvancedSearchAPI(req SearchRequest) (*SearchResponse, error) {
	// 从配置文件获取API URL
	apiURL := utils.GetKnowledgeSearchURL()

	// 构建查询参数
	params := url.Values{}
	params.Add("query_str", req.QueryStr)
	params.Add("n_results", strconv.Itoa(req.NResults))

	// 创建HTTP请求
	fullURL := apiURL + "?" + params.Encode()
	httpReq, err := http.NewRequest("GET", fullURL, nil)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %v", err)
	}

	// 设置请求头
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("User-Agent", "shudao-chat-go/1.0")

	// 创建HTTP客户端
	client := &http.Client{
		Timeout: 30 * time.Second,
	}

	// 发送请求
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("请求失败: %v", err)
	}
	defer resp.Body.Close()

	// 检查响应状态码
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API返回错误状态码: %d", resp.StatusCode)
	}

	// 读取响应体
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %v", err)
	}

	// 解析响应JSON
	var searchResp SearchResponse
	err = json.Unmarshal(body, &searchResp)
	if err != nil {
		return nil, fmt.Errorf("解析响应JSON失败: %v", err)
	}

	return &searchResp, nil
}

// StartChromaSearchHeartbeatTask 启动ChromaDB搜索服务心跳任务
func StartChromaSearchHeartbeatTask() {
	heartbeatURL := utils.GetConfigString("heartbeat_api_url", "")
	if heartbeatURL == "" {
		return
	}

	task := &chromaSearchHeartbeatTask{
		url:        heartbeatURL,
		interval:   10 * time.Minute,
		httpClient: &http.Client{Timeout: 30 * time.Second},
		stopChan:   make(chan struct{}),
	}

	go task.run()
}

func (h *chromaSearchHeartbeatTask) run() {
	ticker := time.NewTicker(h.interval)
	defer ticker.Stop()

	h.sendHeartbeat()

	for {
		select {
		case <-h.stopChan:
			return
		case <-ticker.C:
			h.sendHeartbeat()
		}
	}
}

func (h *chromaSearchHeartbeatTask) sendHeartbeat() {
	resp, err := h.httpClient.Get(h.url)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	io.ReadAll(resp.Body)
}