ai_scheduler/internal/domain/tools/common/knowledge_base/client.go

package knowledge_base

import (
	"ai_scheduler/internal/config"
	"ai_scheduler/internal/pkg/l_request"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

type Client struct {
	cfg config.KnowledgeConfig
}

func New(cfg config.KnowledgeConfig) *Client {
	return &Client{cfg: cfg}
}

// 查询知识库
func (c *Client) Query(req *QueryRequest) (io.ReadCloser, error) {
	if req == nil {
		return nil, fmt.Errorf("req is nil")
	}
	if req.TenantID == "" {
		return nil, fmt.Errorf("tenantID is empty")
	}
	if req.Query == "" {
		return nil, fmt.Errorf("query is empty")
	}
	if req.Mode == "" {
		req.Mode = c.cfg.Mode
	}
	if !req.Think {
		req.Think = c.cfg.Think
	}
	if !req.OnlyRAG {
		req.OnlyRAG = c.cfg.OnlyRAG
	}

	baseURL := strings.TrimRight(c.cfg.BaseURL, "/")

	rsp, err := (&l_request.Request{
		Method: "POST",
		Url:    baseURL + "/query",
		Headers: map[string]string{
			"Content-Type": "application/json",
			"X-Tenant-ID":  req.TenantID,
			"Accept":       "text/event-stream",
		},
		Json: map[string]interface{}{
			"query":    req.Query,
			"mode":     req.Mode,
			"stream":   req.Stream,
			"think":    req.Think,
			"only_rag": req.OnlyRAG,
		},
	}).SendNoParseResponse()
	if err != nil {
		return nil, err
	}
	if rsp == nil || rsp.Body == nil {
		return nil, fmt.Errorf("empty response")
	}

	if rsp.StatusCode != http.StatusOK {
		defer rsp.Body.Close()
		bodyPreview, _ := io.ReadAll(io.LimitReader(rsp.Body, 4096))
		if len(bodyPreview) > 0 {
			return nil, fmt.Errorf("knowledge base returned status %d: %s", rsp.StatusCode, string(bodyPreview))
		}
		return nil, fmt.Errorf("knowledge base returned status %d", rsp.StatusCode)
	}

	return rsp.Body, nil
}

// IngestText 向知识库中注入文本
func (c *Client) IngestText(req *IngestTextRequest) error {
	if req == nil {
		return fmt.Errorf("req is nil")
	}
	if req.TenantID == "" {
		return fmt.Errorf("tenantID is empty")
	}
	if req.Text == "" {
		return fmt.Errorf("text is empty")
	}

	baseURL := strings.TrimRight(c.cfg.BaseURL, "/")

	rsp, err := (&l_request.Request{
		Method: "POST",
		Url:    baseURL + "/ingest/text",
		Headers: map[string]string{
			"Content-Type": "application/json",
			"X-Tenant-ID":  req.TenantID,
		},
		Json: map[string]interface{}{
			"text": req.Text,
		},
	}).Send()
	if err != nil {
		return err
	}

	if rsp.StatusCode != http.StatusOK {
		return fmt.Errorf("knowledge base returned status %d: %s", rsp.StatusCode, rsp.Text)
	}

	return nil
}

// IngestBatchQA 向知识库中注入问答对
func (c *Client) IngestBatchQA(req *IngestBacthQARequest) error {
	if req == nil {
		return fmt.Errorf("req is nil")
	}
	if req.TenantID == "" {
		return fmt.Errorf("tenantID is empty")
	}
	for _, item := range req.QAList {
		if item.Question == "" {
			return fmt.Errorf("question is empty")
		}
		if item.Answer == "" {
			return fmt.Errorf("answer is empty")
		}
	}
	data := []map[string]string{}
	for _, item := range req.QAList {
		data = append(data, map[string]string{
			"question": item.Question,
			"answer":   item.Answer,
		})
	}
	jsonByte, err := json.Marshal(data)
	if err != nil {
		return err
	}

	baseURL := strings.TrimRight(c.cfg.BaseURL, "/")

	rsp, err := (&l_request.Request{
		Method: "POST",
		Url:    baseURL + "/ingest/batch_qa",
		Headers: map[string]string{
			"Content-Type": "application/json",
			"X-Tenant-ID":  req.TenantID,
		},
		JsonByte: jsonByte,
	}).Send()
	if err != nil {
		return err
	}

	if rsp.StatusCode != http.StatusOK {
		return fmt.Errorf("knowledge base returned status %d: %s", rsp.StatusCode, rsp.Text)
	}

	return nil
}