llm_service.go 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. package facade
  2. import (
  3. "eta/eta_api/services/llm/facade/bus_response"
  4. "eta/eta_api/utils/llm"
  5. "eta/eta_api/utils/llm/eta_llm/eta_llm_http"
  6. "eta/eta_api/utils/ws"
  7. "fmt"
  8. "github.com/gorilla/websocket"
  9. "github.com/rdlucklib/rdluck_tools/uuid"
  10. )
  11. var (
  12. llmService, _ = llm.GetInstance(llm.ETA_LLM_CLIENT)
  13. )
  14. func generateSessionCode() (code string) {
  15. return fmt.Sprintf("%s%s", "llm_session_", uuid.NewUUID().Hex32())
  16. }
  17. // AddSession 创建会话session
  18. func AddSession(userId int, conn *websocket.Conn) {
  19. sessionId := generateSessionCode()
  20. session := ws.NewSession(userId, sessionId, conn)
  21. ws.Manager().AddSession(session)
  22. }
  23. // LLMKnowledgeBaseSearchDocs 搜索知识库
  24. func LLMKnowledgeBaseSearchDocs(search LLMKnowledgeSearch) (resp bus_response.SearchDocsEtaResponse, err error) {
  25. docs, err := llmService.SearchKbDocs(search.Query, search.KnowledgeBaseName)
  26. if err != nil {
  27. return
  28. }
  29. for _, doc := range docs.([]eta_llm_http.SearchDocsResponse) {
  30. resp.Content = resp.Content + doc.PageContent
  31. }
  32. resp.Docs = docs.([]eta_llm_http.SearchDocsResponse)
  33. return
  34. }
  35. type LLMKnowledgeSearch struct {
  36. Query string `json:"Query"`
  37. KnowledgeBaseName string `json:"KnowledgeBaseName"`
  38. }