mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-23 00:46:42 +08:00
feat[Go]: implement searches/<search_id>/completions POST (#16440)
### Summary As title ### Type of change - [x] New Feature (non-breaking change which adds functionality) - [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
@@ -30,6 +30,9 @@ import (
|
||||
type SearchHandler struct {
|
||||
searchService *service.SearchService
|
||||
userService *service.UserService
|
||||
streamLLM streamingLLM
|
||||
askService *service.AskService
|
||||
sseWriter SSEWriter
|
||||
}
|
||||
|
||||
// NewSearchHandler create search handler
|
||||
@@ -37,9 +40,16 @@ func NewSearchHandler(searchService *service.SearchService, userService *service
|
||||
return &SearchHandler{
|
||||
searchService: searchService,
|
||||
userService: userService,
|
||||
sseWriter: &ginSSEWriter{},
|
||||
}
|
||||
}
|
||||
|
||||
// SetCompletionDependencies wires the streaming search completion runtime.
|
||||
func (h *SearchHandler) SetCompletionDependencies(streamLLM streamingLLM, askService *service.AskService) {
|
||||
h.streamLLM = streamLLM
|
||||
h.askService = askService
|
||||
}
|
||||
|
||||
// ListSearches list search apps
|
||||
// @Summary List Search Apps
|
||||
// @Description Get list of search apps for the current user with filtering, pagination and sorting
|
||||
@@ -421,3 +431,87 @@ func (h *SearchHandler) UpdateSearch(c *gin.Context) {
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
|
||||
func (h *SearchHandler) Completion(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
jsonError(c, errorCode, errorMessage)
|
||||
return
|
||||
}
|
||||
|
||||
var req service.SearchCompletionsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
jsonError(c, common.CodeArgumentError, "question is required")
|
||||
return
|
||||
}
|
||||
|
||||
searchSvc := h.searchService
|
||||
if searchSvc == nil {
|
||||
searchSvc = service.NewSearchService()
|
||||
}
|
||||
|
||||
plan, code, err := searchSvc.PrepareCompletion(user.ID, c.Param("search_id"), &req)
|
||||
if err != nil {
|
||||
if code == common.CodeAuthenticationError {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": code,
|
||||
"data": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if code == common.CodeServerError {
|
||||
jsonInternalError(c, err)
|
||||
return
|
||||
}
|
||||
jsonError(c, code, err.Error())
|
||||
return
|
||||
}
|
||||
if plan == nil {
|
||||
jsonError(c, common.CodeServerError, "completion plan is nil")
|
||||
return
|
||||
}
|
||||
|
||||
disableWriteDeadlineForSSE(c)
|
||||
c.Header("Content-Type", "text/event-stream; charset=utf-8")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
|
||||
writer := h.sseWriter
|
||||
if writer == nil {
|
||||
writer = &ginSSEWriter{}
|
||||
}
|
||||
if plan.ModelID == "" {
|
||||
writer.Write(c, sseError("chat model not configured"))
|
||||
return
|
||||
}
|
||||
if h.askService == nil {
|
||||
writer.Write(c, sseError("ask service not configured"))
|
||||
return
|
||||
}
|
||||
if h.streamLLM == nil {
|
||||
writer.Write(c, sseError("streaming LLM not configured"))
|
||||
return
|
||||
}
|
||||
|
||||
adapter := &askStreamAdapter{llm: h.streamLLM, tenantID: plan.UserID, modelID: plan.ModelID}
|
||||
|
||||
hadError := false
|
||||
for delta := range h.askService.StreamWithOptions(c.Request.Context(), adapter, plan.UserID, plan.Question, plan.KBIDs, plan.Options) {
|
||||
switch delta.Kind {
|
||||
case service.AskDeltaAnswer:
|
||||
writer.Write(c, sseAnswer(delta.Value, nil, false))
|
||||
case service.AskDeltaMarker:
|
||||
writer.Write(c, sseMarker(delta.Value))
|
||||
case service.AskDeltaError:
|
||||
hadError = true
|
||||
writer.Write(c, sseError(delta.Value))
|
||||
case service.AskDeltaFinal:
|
||||
writer.Write(c, sseAnswer(delta.Value, delta.Refs, true))
|
||||
}
|
||||
}
|
||||
if !hadError {
|
||||
writer.Write(c, "data: {\"code\": 0, \"message\": \"\", \"data\": true}\n\n")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user