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

package controllers

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"time"

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

// 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) {
	// 构建请求URL
	apiURL := "https://aqai.shudaodsj.com:22000/admin/api/v1/knowledge/files/advanced-search"

	// 构建查询参数
	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
}