113 lines
2.4 KiB
Go
113 lines
2.4 KiB
Go
package knowledge_base
|
|
|
|
import (
|
|
"ai_scheduler/internal/config"
|
|
"ai_scheduler/internal/pkg/l_request"
|
|
"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
|
|
}
|