From 133b1e15fd7422969a232c2702567694b0392a37 Mon Sep 17 00:00:00 2001 From: qinling0210 <88864212+qinling0210@users.noreply.github.com> Date: Thu, 2 Jul 2026 09:45:01 +0800 Subject: [PATCH] Port agent PRs to GO (#16529) ### Summary Port the following PRs to GO in this PR https://github.com/infiniflow/ragflow/pull/14210 https://github.com/infiniflow/ragflow/pull/14641 https://github.com/infiniflow/ragflow/pull/15220 https://github.com/infiniflow/ragflow/pull/15228 https://github.com/infiniflow/ragflow/pull/15384 https://github.com/infiniflow/ragflow/pull/15754 https://github.com/infiniflow/ragflow/pull/16413 https://github.com/infiniflow/ragflow/pull/16483 https://github.com/infiniflow/ragflow/pull/16419 https://github.com/infiniflow/ragflow/pull/16361 https://github.com/infiniflow/ragflow/pull/16050 --- cmd/server_main.go | 18 +- internal/agent/canvas/runner.go | 47 ++- internal/agent/component/llm.go | 348 +++++++++++++++++++++- internal/agent/runtime/json_safe.go | 152 ++++++++++ internal/agent/runtime/state.go | 5 +- internal/agent/tool/bgpt.go | 245 +++++++++++++++ internal/agent/tool/registry.go | 1 + internal/common/time.go | 66 ++++- internal/dao/user_canvas.go | 12 +- internal/handler/agent.go | 28 ++ internal/handler/mcp_server.go | 230 ++++++++++++++ internal/mcp/connector.go | 130 ++++++++ internal/mcp/server.go | 445 ++++++++++++++++++++++++++++ internal/mcp/types.go | 225 ++++++++++++++ internal/router/router.go | 10 + internal/service/agent.go | 3 +- internal/service/agent_sessions.go | 4 +- internal/service/bot_completion.go | 6 +- internal/service/chat_pipeline.go | 103 +++++-- internal/service/dataset.go | 1 + internal/service/kg/scoring.go | 6 + internal/tokenizer/tokenizer.go | 31 ++ 22 files changed, 2062 insertions(+), 54 deletions(-) create mode 100644 internal/agent/runtime/json_safe.go create mode 100644 internal/agent/tool/bgpt.go create mode 100644 internal/handler/mcp_server.go create mode 100644 internal/mcp/connector.go create mode 100644 internal/mcp/server.go create mode 100644 internal/mcp/types.go diff --git a/cmd/server_main.go b/cmd/server_main.go index 6bc08138b1..c585af7f22 100644 --- a/cmd/server_main.go +++ b/cmd/server_main.go @@ -47,6 +47,7 @@ import ( "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/handler" + "ragflow/internal/mcp" "ragflow/internal/router" "ragflow/internal/service" "ragflow/internal/service/chunk" @@ -261,6 +262,21 @@ func startServer(config *server.Config) { fileHandler := handler.NewFileHandler(fileService, userService) memoryHandler := handler.NewMemoryHandler(memoryService) mcpHandler := handler.NewMCPHandler(mcpService) + + // MCP server endpoint — exposes RAGFlow capabilities as MCP tools + // (ragflow_retrieval, ragflow_list_datasets, ragflow_list_chats) to + // external AI clients via JSON-RPC over HTTP. + mcpServerHandler := handler.NewMCPServerHandler( + func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + return handler.MCPListDatasets(datasetsService, userID, page, pageSize, orderby, desc) + }, + func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + return handler.MCPListChats(chatService, userID, page, pageSize, orderby, desc) + }, + func(userID string, req mcp.RetrievalRequest) (string, error) { + return handler.MCPRetrieval(datasetsService, userID, req) + }, + ) skillSearchHandler := handler.NewSkillSearchHandler(docEngine) providerHandler := handler.NewProviderHandler(userService, modelProviderService) // Install the agent service's Redis-backed run infrastructure @@ -340,7 +356,7 @@ func startServer(config *server.Config) { adminRuntimeHandler := handler.NewAdminRuntimeHandler(adminRuntimeSelector) // Initialize router - r := router.NewRouter(authHandler, userHandler, tenantHandler, documentHandler, datasetsHandler, systemHandler, knowledgebaseHandler, chunkHandler, llmHandler, chatHandler, chatChannelHandler, langfuseHandler, chatSessionHandler, connectorHandler, searchHandler, fileHandler, memoryHandler, mcpHandler, skillSearchHandler, providerHandler, agentHandler, searchBotHandler, difyRetrievalHandler, pluginHandler, modelHandler, fileCommitHandler, adminRuntimeHandler, openaiChatHandler, botHandler) + r := router.NewRouter(authHandler, userHandler, tenantHandler, documentHandler, datasetsHandler, systemHandler, knowledgebaseHandler, chunkHandler, llmHandler, chatHandler, chatChannelHandler, langfuseHandler, chatSessionHandler, connectorHandler, searchHandler, fileHandler, memoryHandler, mcpHandler, mcpServerHandler, skillSearchHandler, providerHandler, agentHandler, searchBotHandler, difyRetrievalHandler, pluginHandler, modelHandler, fileCommitHandler, adminRuntimeHandler, openaiChatHandler, botHandler) // Create Gin engine ginEngine := gin.New() diff --git a/internal/agent/canvas/runner.go b/internal/agent/canvas/runner.go index 05254bc57a..961a2b8759 100644 --- a/internal/agent/canvas/runner.go +++ b/internal/agent/canvas/runner.go @@ -60,6 +60,7 @@ import ( "github.com/google/uuid" "go.uber.org/zap" + "ragflow/internal/agent/runtime" "ragflow/internal/common" ) @@ -334,14 +335,13 @@ func (r *Runner) Run( } } } - payload, _ := json.Marshal(waiting) - push(out, RunEvent{Type: "waiting_for_user", Data: string(payload), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) - // Always close a RunAgent call with the `done` - // terminator so the front-end can rely on a - // channel-end sentinel regardless of whether the run - // completed, errored, or paused for user input. - push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) - return + push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(waiting), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) + // Always close a RunAgent call with the `done` + // terminator so the front-end can rely on a + // channel-end sentinel regardless of whether the run + // completed, errored, or paused for user input. + push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) + return } if IsInterruptError(runErr) { // Raw InterruptSignal (no wrapped InterruptCtx list @@ -349,8 +349,7 @@ func (r *Runner) Run( // without a cpn id — the front-end falls back to // the first paused session it knows about. r.saveInterruptID(canvasID, sessionID, runErr.Error()) - payload, _ := json.Marshal(WaitingForUserEvent{CpnID: runErr.Error()}) - push(out, RunEvent{Type: "waiting_for_user", Data: string(payload), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) + push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(WaitingForUserEvent{CpnID: runErr.Error()}), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) return } @@ -457,10 +456,36 @@ func push(out chan<- RunEvent, ev RunEvent) { // pushErr serialises an ErrorEvent and pushes it on the channel. func pushErr(out chan<- RunEvent, msg string) { - payload, _ := json.Marshal(ErrorEvent{Message: msg}) + payload, err := json.Marshal(ErrorEvent{Message: msg}) + if err != nil { + common.Warn("runner: pushErr json.Marshal failed, falling back", + zap.Error(err)) + // ErrorEvent only has a string field; this should never fail. + // Fall back to a hard-coded minimal JSON. + payload = []byte(`{"message":"event serialization failed"}`) + } push(out, RunEvent{Type: "error", Data: string(payload)}) } +// safeEventJSON marshals v to a JSON string, falling back to +// runtime.SafeJSONMarshal when the value contains non-serializable +// types (funcs, channels). Mirrors the Python PR #14210 +// _canvas_json_default fallback for SSE event serialization. +func safeEventJSON(v any) string { + b, err := json.Marshal(v) + if err != nil { + common.Warn("runner: json.Marshal event payload failed, trying SafeJSONMarshal", + zap.Error(err)) + b, err = runtime.SafeJSONMarshal(v) + if err != nil { + common.Error("runner: SafeJSONMarshal also failed, using fallback", + err) + b = []byte(`{"message":"event serialization failed"}`) + } + } + return string(b) +} + // nowUnix returns the current Unix timestamp in seconds. func nowUnix() int64 { return time.Now().Unix() diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index 320ce3069f..b06f7d0c00 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -27,6 +27,7 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/common" "ragflow/internal/entity/models" + "ragflow/internal/tokenizer" "go.uber.org/zap" ) @@ -313,6 +314,8 @@ func chatURLSuffixFor(driver string) models.URLSuffix { switch strings.ToLower(driver) { case "anthropic": return models.URLSuffix{Chat: "v1/messages"} + case "ollama": + return models.URLSuffix{Chat: "api/chat"} default: return models.URLSuffix{Chat: "chat/completions"} } @@ -371,8 +374,63 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s if p.UserPrompt == "" { p.UserPrompt = p.SystemPrompt } + // Collect sys.files from canvas globals and inject their + // content into prompts and the image list. Mirrors Python's + // _collect_sys_files and the injection path in + // _prepare_prompt_variables (llm.py:225-281). + var sysFileTexts []string + var sysFileImgs []string + hasSysFilesPlaceholder := strings.Contains(p.SystemPrompt, "{sys.files}") || strings.Contains(p.UserPrompt, "{sys.files}") + if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { + sysFileTexts, sysFileImgs = collectSysFiles(state) + if len(sysFileImgs) > 0 { + p.VisualFiles = dedupStrings(append(p.VisualFiles, sysFileImgs...)) + } + } + // When the prompt contains an explicit {sys.files} placeholder, + // replace it with the collected file text and clear sysFileTexts + // so it is not injected again below. + if hasSysFilesPlaceholder { + joined := strings.Join(sysFileTexts, "\n\n") + p.SystemPrompt = strings.ReplaceAll(p.SystemPrompt, "{sys.files}", joined) + p.UserPrompt = strings.ReplaceAll(p.UserPrompt, "{sys.files}", joined) + sysFileTexts = nil + } msgs := buildMessagesWithImages(p.SystemPrompt, p.UserPrompt, p.VisualFiles, p.Cite) + // Inject sys.files text content into the last user message. + if len(sysFileTexts) > 0 { + joined := strings.Join(sysFileTexts, "\n\n") + if len(msgs) > 0 && msgs[len(msgs)-1].Role == schema.User { + last := &msgs[len(msgs)-1] + if len(last.UserInputMultiContent) > 0 { + inserted := false + for i := range last.UserInputMultiContent { + if last.UserInputMultiContent[i].Type == schema.ChatMessagePartTypeText { + if last.UserInputMultiContent[i].Text != "" { + last.UserInputMultiContent[i].Text += "\n\n" + joined + } else { + last.UserInputMultiContent[i].Text = joined + } + inserted = true + break + } + } + if !inserted { + last.UserInputMultiContent = append([]schema.MessageInputPart{{ + Type: schema.ChatMessagePartTypeText, + Text: joined, + }}, last.UserInputMultiContent...) + } + } else if last.Content != "" { + last.Content += "\n\n" + joined + } else { + last.Content = joined + } + } else { + msgs = append(msgs, schema.Message{Role: schema.User, Content: joined}) + } + } // Prepend the last N turns of conversation history from the // canvas state. Mirrors Python's `_get_chat_template_kwargs` / // `_fit_messages` path. When window size is 0 or history is @@ -383,6 +441,23 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s msgs = prependHistory(msgs, state.History, p.MessageHistoryWindowSize) } } + // Apply message fitting (trim to context window) after all + // prompt/history/sys.files augmentation and before invoking the + // LLM. Mirrors Python's message_fit_in in PR #16413. + { + maxCtx := 0 + if p.MaxTokens != nil { + maxCtx = *p.MaxTokens + } + // The system prompt is already embedded as the first message + // in msgs by buildMessagesWithImages; pass "" so fitMessages + // does not duplicate it. + fitted, fitErr := fitMessages("", msgs, maxCtx) + if fitErr != "" { + return map[string]any{"content": fitErr}, nil + } + msgs = fitted + } inv := getDefaultChatInvoker() // Param-level retry override. When MaxRetries OR // DelayAfterError is set on LLMParam, the user is asking @@ -697,6 +772,46 @@ func extractDataImages(values []string) []string { return out } +// collectSysFiles splits sys.files from canvas globals into text parts +// and image data URIs. The caller is responsible for handling any +// {sys.files} placeholder replacement in the prompts. +func collectSysFiles(state *runtime.CanvasState) (textParts, imageURIs []string) { + files, ok := state.Globals["sys.files"] + if !ok { + return nil, nil + } + fileList, ok := files.([]any) + if !ok || len(fileList) == 0 { + return nil, nil + } + for _, f := range fileList { + s, ok := f.(string) + if !ok { + continue + } + if strings.HasPrefix(s, "data:image/") { + imageURIs = append(imageURIs, s) + } else { + textParts = append(textParts, s) + } + } + return textParts, imageURIs +} + +// dedupStrings returns the deduplicated slice in first-seen order. +func dedupStrings(vals []string) []string { + seen := make(map[string]struct{}, len(vals)) + out := make([]string, 0, len(vals)) + for _, v := range vals { + if _, dup := seen[v]; dup { + continue + } + seen[v] = struct{}{} + out = append(out, v) + } + return out +} + // prependHistory inserts up to `window` prior turns from the canvas // history before the current system+user messages. Each history entry // is a {role, content} map; only the last `window` are kept, with @@ -860,10 +975,13 @@ func mergeLLMParam(base LLMParam, inputs map[string]any) LLMParam { p.MaxTokens = &i } if v, ok := stringFrom(inputs, "thinking"); ok { - // Only allow the two known sentinels through; an arbitrary - // string from the DSL is dropped to avoid surprising the LLM - // driver. Mirrors python llm.py:78-79 which gates on the - // same {"enabled","disabled"} set. + // Only allow "enabled" or "disabled"; arbitrary DSL + // strings are dropped. Python PR #15220 removed + // thinking from llm.py's gen_conf() — it is no + // longer forwarded to the model. The field is still + // parsed here to match the Python form parameter + // definition, but einoChatInvoker does not consume + // it, consistent with Python's behavior. if v == "enabled" || v == "disabled" { p.Thinking = v } @@ -871,6 +989,228 @@ func mergeLLMParam(base LLMParam, inputs map[string]any) LLMParam { return p } +// effectiveContextLength returns maxLength if positive, otherwise 8192. +// Mirrors Python's LLM.effective_context_length in PR #16413 — prevents +// zero/negative context windows from silently trimming all prompt content. +func effectiveContextLength(maxLength int) int { + if maxLength > 0 { + return maxLength + } + return 8192 +} + +// contextFitBudget returns 97% of the effective context length as the +// token budget for message_fit_in. Mirrors Python's LLM.context_fit_budget +// in PR #16413. +func contextFitBudget(maxLength int) int { + return int(float64(effectiveContextLength(maxLength)) * 0.97) +} + +// validateFittedMessages checks that the fitted message list is non-empty +// and the last message is a non-empty user turn (content or multi-modal +// parts). Returns an error string on failure, empty string on success. +// Python requires len >= 2 because the system prompt is always injected +// upstream; Go allows len >= 1 because the system message may be embedded +// inside msgs (from buildMessagesWithImages) or absent entirely. +func validateFittedMessages(msgFit []schema.Message) string { + if len(msgFit) == 0 { + return "**ERROR**: message_fit_in produced insufficient messages for LLM" + } + last := msgFit[len(msgFit)-1] + if last.Role != schema.User { + return "**ERROR**: LLM last message is not a user turn after prompt fitting; check model max_tokens context setting" + } + if strings.TrimSpace(last.Content) == "" && len(last.UserInputMultiContent) == 0 { + return "**ERROR**: LLM user message is empty after prompt fitting; check model max_tokens context setting" + } + return "" +} + +// fitMessages calls message_fit_in semantics on the given messages and +// validates that the result ends with a non-empty user turn. Returns the +// fitted messages and an error string (empty on success). +// Mirrors Python's LLM.fit_messages in PR #16413. +func fitMessages(systemPrompt string, msgs []schema.Message, maxLength int) ([]schema.Message, string) { + // Convert schema.Message → []map[string]interface{} for fitting. + all := make([]map[string]interface{}, 0, 1+len(msgs)) + // Deep-copy msgs (mirrors Python's deepcopy) to avoid mutating caller's slice. + copied := make([]schema.Message, len(msgs)) + for i, m := range msgs { + cloned := slices.Clone(m.UserInputMultiContent) + for j, p := range cloned { + if p.Image != nil { + imgCopy := *p.Image + if p.Image.URL != nil { + u := *p.Image.URL + imgCopy.URL = &u + } + cloned[j].Image = &imgCopy + } + } + copied[i] = schema.Message{ + Role: m.Role, + Content: m.Content, + UserInputMultiContent: cloned, + } + } + // System prompt first (when not already embedded in msgs). + if systemPrompt != "" { + all = append(all, map[string]interface{}{ + "role": "system", + "content": systemPrompt, + }) + } + for _, m := range copied { + entry := map[string]interface{}{ + "role": string(m.Role), + "content": m.Content, + } + if len(m.UserInputMultiContent) > 0 { + entry["user_input_multi_content"] = m.UserInputMultiContent + } + all = append(all, entry) + } + // Use 97% of effective context as the token budget. + budget := contextFitBudget(maxLength) + _, fitted := messageFitInRaw(all, budget) + + // Convert back to []schema.Message. + result := make([]schema.Message, 0, len(fitted)) + for _, m := range fitted { + role, _ := m["role"].(string) + content, _ := m["content"].(string) + msg := schema.Message{ + Role: schema.RoleType(role), + Content: content, + } + if multi, ok := m["user_input_multi_content"].([]schema.MessageInputPart); ok { + msg.UserInputMultiContent = multi + } + result = append(result, msg) + } + return result, validateFittedMessages(result) +} + +// messageFitInRaw trims messages to fit within a token budget. Operates on +// raw []map[string]interface{} (role + content). Returns the token count +// used and the trimmed slice. Mirrors Python's message_fit_in in +// rag/prompts/generator.py. +// +// Strategy: +// 1. If everything fits → return as-is. +// 2. Keep all system messages + the last user/assistant message. +// 3. If still too large, trim content proportionally: +// - System dominates (>80%) → preserve last message first. +// - Otherwise → preserve system first. +func messageFitInRaw(messages []map[string]interface{}, maxTokens int) (int, []map[string]interface{}) { + if maxTokens <= 0 { + maxTokens = 8192 + } + + // Step 1: everything fits. + totalTokens := countAllTokens(messages) + if totalTokens < maxTokens { + return totalTokens, messages + } + + // Step 2: keep all system messages + the last non-system message. + result := make([]map[string]interface{}, 0) + for _, m := range messages { + if role, _ := m["role"].(string); role == "system" { + result = append(result, m) + } + } + if len(messages) > 0 { + last := messages[len(messages)-1] + if role, _ := last["role"].(string); role != "system" { + result = append(result, last) + } + } + if len(result) == 0 { + return 0, result + } + + totalTokens = countAllTokens(result) + if totalTokens < maxTokens { + return totalTokens, result + } + + // Step 3: trim content to fit. + ll := tokenizer.NumTokensFromString(stringContent(result[0])) + ll2 := tokenizer.NumTokensFromString(stringContent(result[len(result)-1])) + total := ll + ll2 + if total <= 0 { + return 0, result + } + + if len(result) == 1 { + setContent(result[0], tokenizer.TrimContentToTokenLimit(stringContent(result[0]), maxTokens)) + return countAllTokens(result), result + } + + if float64(ll)/float64(total) > 0.8 { + preservedLast := min(ll2, maxTokens) + setContent(result[len(result)-1], tokenizer.TrimContentToTokenLimit(stringContent(result[len(result)-1]), preservedLast)) + remaining := max(0, maxTokens-preservedLast) + setContent(result[0], tokenizer.TrimContentToTokenLimit(stringContent(result[0]), remaining)) + } else { + preservedSystem := min(ll, maxTokens) + setContent(result[0], tokenizer.TrimContentToTokenLimit(stringContent(result[0]), preservedSystem)) + remaining := max(0, maxTokens-preservedSystem) + setContent(result[len(result)-1], tokenizer.TrimContentToTokenLimit(stringContent(result[len(result)-1]), remaining)) + } + + return countAllTokens(result), result +} + +// countAllTokens returns the total token count across all messages. +func countAllTokens(messages []map[string]interface{}) int { + total := 0 + for _, m := range messages { + total += tokenizer.NumTokensFromString(stringContent(m)) + } + return total +} + +// stringContent extracts the display text from a message map, or "". +// For plain messages the text lives in "content"; for multimodal messages +// (with images) it lives in a text part of "user_input_multi_content". +func stringContent(m map[string]interface{}) string { + s, _ := m["content"].(string) + if s != "" { + return s + } + if multi, ok := m["user_input_multi_content"].([]schema.MessageInputPart); ok { + for _, part := range multi { + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + return part.Text + } + } + } + return "" +} + +// setContent writes text into a message map's display field. When the +// message has user_input_multi_content parts, the first text part is +// updated; otherwise the plain "content" field is set. +func setContent(m map[string]interface{}, text string) { + if multi, ok := m["user_input_multi_content"].([]schema.MessageInputPart); ok { + for i, part := range multi { + if part.Type == schema.ChatMessagePartTypeText { + multi[i].Text = text + return + } + } + // No existing text part – prepend one. + m["user_input_multi_content"] = append( + []schema.MessageInputPart{{Type: schema.ChatMessagePartTypeText, Text: text}}, + multi..., + ) + return + } + m["content"] = text +} + // stringFrom extracts a string from inputs[name], accepting both string and // fmt.Stringer-able values. func stringFrom(inputs map[string]any, name string) (string, bool) { diff --git a/internal/agent/runtime/json_safe.go b/internal/agent/runtime/json_safe.go new file mode 100644 index 0000000000..f458ef14f1 --- /dev/null +++ b/internal/agent/runtime/json_safe.go @@ -0,0 +1,152 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// json_safe.go — Safe JSON marshaling with fallback for non-serializable +// values, mirroring the Python PR #14210 approach. +// +// Agent components may store functools.partial-like objects (function +// closures) or non-copyable client objects in their output slots. These +// values leak into canvas state maps and must not crash json.Marshal. +// This file provides helpers that recursively clean state maps before +// serialization, converting function values and channels to nil. +package runtime + +import ( + "encoding/json" + "fmt" + "reflect" + "strings" + + "go.uber.org/zap" + + "ragflow/internal/common" +) + +// SafeJSONMarshal marshals v to JSON after recursively replacing +// non-serializable Go values (func, chan) with nil. This mirrors the +// Python PR #14210 _serialize_default / _canvas_json_default fallback +// that converts callables to None. +// +// Use this in preference to json.Marshal when the input is a canvas +// state map (map[string]any) that may contain function closures, +// channels, or other types encoding/json cannot handle. +func SafeJSONMarshal(v any) ([]byte, error) { + cleaned := cleanForJSON(v) + return json.Marshal(cleaned) +} + +// cleanForJSON recursively replaces non-JSON-serializable values with nil. +func cleanForJSON(v any) any { + if v == nil { + return nil + } + switch val := v.(type) { + case map[string]any: + return cleanMap(val) + case map[string]map[string]any: + out := make(map[string]map[string]any, len(val)) + for k, vv := range val { + out[k] = cleanMap(vv) + } + return out + case []map[string]any: + out := make([]any, len(val)) + for i, vv := range val { + out[i] = cleanMap(vv) + } + return out + case []any: + out := make([]any, len(val)) + for i, vv := range val { + out[i] = cleanForJSON(vv) + } + return out + default: + rv := reflect.ValueOf(v) + if !rv.IsValid() { + return nil + } + switch rv.Kind() { + case reflect.Func, reflect.Chan: + common.Info("cleanForJSON: converting non-serializable value to nil", + zap.String("type", rv.Type().String())) + return nil + case reflect.Map: + // Generic map fallback (e.g., map[string]int, etc.). + // Build as map[string]any because cleaned values may + // change type (func→nil, typed slice→[]any, etc.) and + // SetMapIndex on an original-typed map would panic. + cleaned := make(map[string]any, rv.Len()) + for iter := rv.MapRange(); iter.Next(); { + key := fmt.Sprint(iter.Key().Interface()) + cleaned[key] = cleanForJSON(iter.Value().Interface()) + } + return cleaned + case reflect.Slice, reflect.Array: + length := rv.Len() + cleaned := make([]any, length) + for i := 0; i < length; i++ { + cleaned[i] = cleanForJSON(rv.Index(i).Interface()) + } + return cleaned + case reflect.Ptr: + if rv.IsNil() { + return nil + } + return cleanForJSON(rv.Elem().Interface()) + case reflect.Interface: + if rv.IsNil() { + return nil + } + return cleanForJSON(rv.Elem().Interface()) + case reflect.Struct: + cleaned := make(map[string]any, rv.NumField()) + t := rv.Type() + for i := 0; i < rv.NumField(); i++ { + field := t.Field(i) + if !field.IsExported() { + continue + } + fv := rv.Field(i) + if !fv.CanInterface() { + continue + } + key := field.Name + if tag := field.Tag.Get("json"); tag != "" { + name, _, _ := strings.Cut(tag, ",") + if name == "-" { + continue + } + if name != "" { + key = name + } + } + cleaned[key] = cleanForJSON(fv.Interface()) + } + return cleaned + default: + return v + } + } +} + +func cleanMap(m map[string]any) map[string]any { + out := make(map[string]any, len(m)) + for k, v := range m { + out[k] = cleanForJSON(v) + } + return out +} diff --git a/internal/agent/runtime/state.go b/internal/agent/runtime/state.go index d56ff647b7..b727f53dfb 100644 --- a/internal/agent/runtime/state.go +++ b/internal/agent/runtime/state.go @@ -146,7 +146,10 @@ func (s *CanvasState) MarshalJSON() ([]byte, error) { RunID: s.RunID, TaskID: s.TaskID, } - return json.Marshal(snap) + // Use SafeJSONMarshal to handle non-serializable values (funcs, + // channels) that may have leaked into state maps. Mirrors the + // Python PR #14210 _serialize_default fallback in Graph.__str__. + return SafeJSONMarshal(snap) } // UnmarshalJSON restores the wire shape produced by MarshalJSON. diff --git a/internal/agent/tool/bgpt.go b/internal/agent/tool/bgpt.go new file mode 100644 index 0000000000..106bd7cb8d --- /dev/null +++ b/internal/agent/tool/bgpt.go @@ -0,0 +1,245 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tool + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +const bgptToolName = "bgpt_search" + +const bgptToolDescription = "Search scientific papers via BGPT (bgpt.pro) and return structured evidence " + + "extracted from full-text studies: methods, sample sizes, results, limitations, conflicts of interest, " + + "data/code availability, study blind spots, quality scores, and falsification prompts." + +// bgptEndpoint is the BGPT search URL. Exposed as a package var for tests. +var bgptEndpoint = "https://bgpt.pro/api/mcp-search" + +// bgptParams is the JSON shape the model sends into InvokableRun. +type bgptParams struct { + Query string `json:"query"` + MaxResults int `json:"num_results"` + APIKey string `json:"api_key,omitempty"` + DaysBack int `json:"days_back,omitempty"` +} + +// bgptResult is one paper in the result list. +type bgptResult struct { + Title string `json:"title"` + Authors string `json:"authors"` + Journal string `json:"journal"` + Year string `json:"year"` + DOI string `json:"doi"` + URL string `json:"url"` + Abstract string `json:"abstract"` + Methods string `json:"methods"` + SampleSize string `json:"sample_size"` + Results string `json:"results"` + Limitations string `json:"limitations"` + ConflictOfInterest string `json:"conflict_of_interest"` + DataAvailability string `json:"data_availability"` + BlindSpots string `json:"blind_spots"` + Falsify string `json:"falsify"` +} + +// bgptEnv is what the model sees. +type bgptEnv struct { + Results []bgptResult `json:"results"` + Error string `json:"_ERROR,omitempty"` +} + +// bgptResponse is the upstream BGPT API envelope. +type bgptResponse struct { + Results []map[string]interface{} `json:"results"` + Error string `json:"error,omitempty"` +} + +// BGPTTool searches scientific papers via bgpt.pro. +type BGPTTool struct { + helper *HTTPHelper +} + +// NewBGPTTool returns a BGPTTool using the default HTTPHelper. +func NewBGPTTool() *BGPTTool { + return NewBGPTToolWith(NewHTTPHelper()) +} + +// NewBGPTToolWith returns a BGPTTool with the given helper. +func NewBGPTToolWith(h *HTTPHelper) *BGPTTool { + if h == nil { + h = NewHTTPHelper() + } + return &BGPTTool{helper: h} +} + +// Info returns the tool's metadata for the chat model. +func (b *BGPTTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: bgptToolName, + Desc: bgptToolDescription, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Type: schema.String, + Desc: "Natural-language scientific search query.", + Required: true, + }, + "num_results": { + Type: schema.Integer, + Desc: "Maximum number of results. Defaults to 10.", + Required: false, + }, + "api_key": { + Type: schema.String, + Desc: "Optional BGPT API key. Leave blank for the free tier.", + Required: false, + }, + "days_back": { + Type: schema.Integer, + Desc: "Optional recency filter (e.g. 365 for last year).", + Required: false, + }, + }), + }, nil +} + +// InvokableRun performs the BGPT search. +func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { + var p bgptParams + if err := json.Unmarshal([]byte(argsJSON), &p); err != nil { + return bgptErrJSON(fmt.Errorf("bgpt: parse arguments: %w", err)), + fmt.Errorf("bgpt: parse arguments: %w", err) + } + if strings.TrimSpace(p.Query) == "" { + return bgptErrJSON(fmt.Errorf("query is required")), + fmt.Errorf("bgpt: query is required") + } + if p.MaxResults <= 0 { + p.MaxResults = 10 + } + + reqBody := map[string]interface{}{ + "query": strings.TrimSpace(p.Query), + "num_results": p.MaxResults, + } + if p.APIKey != "" { + reqBody["api_key"] = p.APIKey + } + if p.DaysBack > 0 { + reqBody["days_back"] = p.DaysBack + } + + bodyBytes, err := json.Marshal(reqBody) + if err != nil { + return bgptErrJSON(err), err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, bgptEndpoint, bytes.NewReader(bodyBytes)) + if err != nil { + return bgptErrJSON(err), err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + + client := &http.Client{Timeout: 25 * time.Second} + resp, err := client.Do(req) + if err != nil { + return bgptErrJSON(fmt.Errorf("bgpt: request failed: %w", err)), + fmt.Errorf("bgpt: request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return bgptErrJSON(fmt.Errorf("bgpt: read response: %w", err)), + fmt.Errorf("bgpt: read response: %w", err) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return bgptErrJSON(fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body))), + fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body)) + } + + var raw bgptResponse + if err := json.Unmarshal(body, &raw); err != nil { + return bgptErrJSON(fmt.Errorf("bgpt: parse response: %w", err)), + fmt.Errorf("bgpt: parse response: %w", err) + } + if strings.TrimSpace(raw.Error) != "" { + err := fmt.Errorf("bgpt: %s", strings.TrimSpace(raw.Error)) + return bgptErrJSON(err), err + } + + results := make([]bgptResult, 0, len(raw.Results)) + for _, r := range raw.Results { + results = append(results, bgptResult{ + Title: strVal(r["title"]), + Authors: strVal(r["authors"]), + Journal: strVal(r["journal"]), + Year: strVal(r["year"]), + DOI: strVal(r["doi"]), + URL: strVal(r["url"]), + Abstract: strVal(r["abstract"]), + Methods: firstStr(r, "methods_and_experimental_techniques", "methods"), + SampleSize: firstStr(r, "sample_size_and_population_characteristics", "sample_size_and_population"), + Results: firstStr(r, "results_and_conclusions", "results"), + Limitations: firstStr(r, "paper_limitations_and_biases", "limitations"), + ConflictOfInterest: firstStr(r, "conflict_of_interest_statements", "conflict_of_interest"), + DataAvailability: firstStr(r, "data_availability_statements", "data_availability"), + BlindSpots: strVal(r["study_blindspots"]), + Falsify: strVal(r["how_to_falsify"]), + }) + } + + env := bgptEnv{Results: results} + out, _ := json.Marshal(env) + return string(out), nil +} + +func bgptErrJSON(err error) string { + env := bgptEnv{Error: err.Error()} + b, _ := json.Marshal(env) + return string(b) +} + +func strVal(v interface{}) string { + if v == nil { + return "" + } + s, _ := v.(string) + return s +} + +func firstStr(m map[string]interface{}, keys ...string) string { + for _, k := range keys { + if s := strVal(m[k]); s != "" { + return s + } + } + return "" +} + + diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 565d331361..ecc0efd004 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -31,6 +31,7 @@ type Factory func(params map[string]any) (einotool.BaseTool, error) var registry = map[string]Factory{ "akshare": noConfig("akshare", func() einotool.BaseTool { return NewAkShareTool() }), "arxiv": noConfig("arxiv", func() einotool.BaseTool { return NewArxivTool() }), + "bgpt": noConfig("bgpt", func() einotool.BaseTool { return NewBGPTTool() }), "code_exec": noConfig("code_exec", func() einotool.BaseTool { return NewCodeExecTool() }), "crawler": noConfig("crawler", func() einotool.BaseTool { return NewCrawlerTool() }), "deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }), diff --git a/internal/common/time.go b/internal/common/time.go index 25d04588b1..ed2127315b 100644 --- a/internal/common/time.go +++ b/internal/common/time.go @@ -17,9 +17,71 @@ package common import ( + "log/slog" + "strings" "time" ) +// ParseISO8601 parses a date string trying multiple ISO 8601 / RFC3339 +// variants, mirroring the flexibility of Python's dateutil.isoparse. +// +// Supported formats (tried in order): +// - RFC3339Nano: "2006-01-02T15:04:05.999999999Z07:00" +// - RFC3339: "2006-01-02T15:04:05Z07:00" +// - No timezone: "2006-01-02T15:04:05" +// - Date only: "2006-01-02" +// +// Z suffix is automatically normalised to +00:00 before parsing (PR #16483). +func ParseISO8601(dateString string) (time.Time, error) { + // Normalise Z → +00:00 for compatibility with time.RFC3339. + normalized := dateString + if strings.HasSuffix(dateString, "Z") { + normalized = dateString[:len(dateString)-1] + "+00:00" + } + + layouts := []string{ + time.RFC3339Nano, // "2006-01-02T15:04:05.999999999Z07:00" + time.RFC3339, // "2006-01-02T15:04:05Z07:00" + "2006-01-02T15:04:05", // no timezone + "2006-01-02", // date only + } + for _, layout := range layouts { + var t time.Time + var err error + if strings.Contains(layout, "Z07:00") || strings.Contains(layout, "MST") { + t, err = time.Parse(layout, normalized) + } else { + t, err = time.ParseInLocation(layout, normalized, time.Local) + } + if err == nil { + return t, nil + } + } + return time.Time{}, &time.ParseError{ + Layout: "ISO 8601", + Value: dateString, + LayoutElem: "", + ValueElem: dateString, + Message: "failed to parse as any supported ISO 8601 variant", + } +} + +// FormatISO8601ToYMDHMS parses an ISO 8601 / RFC3339 date string and +// returns it formatted as "YYYY-MM-DD HH:MM:SS". If parsing fails the +// original string is returned unchanged. +// +// Mirrors Python's format_iso_8601_to_ymd_hms in common/time_utils.py, +// with the fix from PR #16483 (single-parse using dateutil.isoparse +// instead of a broken double-parse). +func FormatISO8601ToYMDHMS(timeStr string) string { + dt, err := ParseISO8601(timeStr) + if err != nil { + slog.Error("FormatISO8601ToYMDHMS parse error", "input", timeStr, "error", err) + return timeStr + } + return dt.Format("2006-01-02 15:04:05") +} + // DeltaSeconds calculates seconds elapsed from a given date string to now. // // Supports multiple time formats: @@ -39,8 +101,8 @@ import ( // DeltaSeconds("2024-01-01 12:00:00") // DeltaSeconds("2026-04-09T18:55:46+08:00") func DeltaSeconds(dateString string) (float64, error) { - // Try RFC3339 format first (ISO 8601 with timezone, e.g., "2026-04-09T18:55:46+08:00") - dt, err := time.Parse(time.RFC3339, dateString) + // Try ISO 8601 / RFC3339 with flexible parsing (PR #16483). + dt, err := ParseISO8601(dateString) if err == nil { return time.Since(dt).Seconds(), nil } diff --git a/internal/dao/user_canvas.go b/internal/dao/user_canvas.go index 0ca8babdfd..67adf7538f 100644 --- a/internal/dao/user_canvas.go +++ b/internal/dao/user_canvas.go @@ -275,7 +275,7 @@ func (dao *UserCanvasDAO) GetByUserAndTitle(userID, title, canvasCategory string // GetList get canvases list with pagination and filtering // Similar to Python UserCanvasService.get_list -func (dao *UserCanvasDAO) GetList(tenantID string, pageNumber, itemsPerPage int, orderby string, desc bool, id, title string, canvasCategory string) ([]*entity.UserCanvas, error) { +func (dao *UserCanvasDAO) GetList(tenantID string, pageNumber, itemsPerPage int, orderby string, desc bool, id, title string, canvasCategory, canvasType string) ([]*entity.UserCanvas, error) { query := DB.Model(&entity.UserCanvas{}). Where("user_id = ?", tenantID) @@ -293,6 +293,10 @@ func (dao *UserCanvasDAO) GetList(tenantID string, pageNumber, itemsPerPage int, query = query.Where("canvas_category = ?", "agent_canvas") } + if canvasType != "" { + query = query.Where("canvas_type = ?", canvasType) + } + // Order by // Route orderby through userCanvasOrderClause above so user-supplied // query params can never reach Order() verbatim. The helper validates @@ -348,7 +352,7 @@ type UserCanvasListItem struct { // ListByTenantIDs lists agent canvases accessible to the given owner IDs with optional // keyword filter, tag filter, pagination, and ordering. // Mirrors Python UserCanvasService.get_by_tenant_ids (list route only). -func (dao *UserCanvasDAO) ListByTenantIDs(ownerIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords string, canvasCategory string, tags []string) ([]*UserCanvasListItem, int64, error) { +func (dao *UserCanvasDAO) ListByTenantIDs(ownerIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords, canvasCategory, canvasType string, tags []string) ([]*UserCanvasListItem, int64, error) { if len(ownerIDs) == 0 { return nil, 0, nil } @@ -385,6 +389,10 @@ func (dao *UserCanvasDAO) ListByTenantIDs(ownerIDs []string, userID string, page base = base.Where("user_canvas.canvas_category = ?", "agent_canvas") } + if canvasType != "" { + base = base.Where("canvas_type = ?", canvasType) + } + if keywords != "" { like := "%" + keywords + "%" base = base.Where("user_canvas.title LIKE ?", like) diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 55961c9010..1e42e3798b 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -27,6 +27,7 @@ import ( "ragflow/internal/engine/redis" "strconv" "strings" + "time" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -136,6 +137,7 @@ func (h *AgentHandler) ListAgents(c *gin.Context) { keywords := c.Query("keywords") canvasCategory := c.Query("canvas_category") + canvasType := c.Query("canvas_type") page := 0 if v := c.Query("page"); v != "" { @@ -186,6 +188,7 @@ func (h *AgentHandler) ListAgents(c *gin.Context) { desc, ownerIDs, canvasCategory, + canvasType, tags, ) if err != nil { @@ -1103,7 +1106,9 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { // parser feeds each `data:` line directly into JSON.parse and // breaks on the `e` of `event:` (browser console: "SyntaxError: // Unexpected token 'e', \"event: mes\"…"). + emitted := false for ev := range events { + emitted = true common.Debug("agent chat completions: streaming event", zap.String("agent_id", req.AgentID), zap.String("session_id", req.SessionID), @@ -1119,6 +1124,29 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { return } } + if !emitted { + // Canvas produced no events (e.g. empty query). Echo the + // session_id so the client can resume the conversation + // (fixes #15169). The [DONE] terminator must be emitted + // here explicitly because the canvas never sends a + // "done" event on this path. + common.Info("empty agent output - returning session_id", + zap.String("agent_id", req.AgentID), + zap.String("session_id", req.SessionID), + ) + event := canvas.RunEvent{ + Type: "", + Data: "{}", + CreatedAt: time.Now().Unix(), + SessionID: req.SessionID, + } + _ = service.WriteChatbotRunEvent(c.Writer, event) + if _, err := c.Writer.Write([]byte("data: [DONE]\n\n")); err != nil { + common.Debug("agent chat completions: failed to write [DONE]", + zap.Error(err), + ) + } + } common.Debug("agent chat completions: stream closed", zap.String("agent_id", req.AgentID), zap.String("session_id", req.SessionID), diff --git a/internal/handler/mcp_server.go b/internal/handler/mcp_server.go new file mode 100644 index 0000000000..bba0fbf9e9 --- /dev/null +++ b/internal/handler/mcp_server.go @@ -0,0 +1,230 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package handler + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/gin-gonic/gin" + + "ragflow/internal/common" + "ragflow/internal/mcp" + "ragflow/internal/service" +) + +// MCPRetrievalService abstracts the dataset retrieval operations needed +// by the MCP server handler. +type MCPRetrievalService interface { + SearchDatasets(req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) + ListDatasets(id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) +} + +// MCPServerHandler handles MCP protocol requests (JSON-RPC over HTTP). +// It exposes RAGFlow capabilities as MCP tools to external AI clients. +type MCPServerHandler struct { + listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + retrievalFunc func(userID string, req mcp.RetrievalRequest) (string, error) +} + +// NewMCPServerHandler creates a new MCPServerHandler. +// The service functions are passed as closures to avoid importing the service +// package directly from the handler layer. +func NewMCPServerHandler( + listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + retrievalFunc func(userID string, req mcp.RetrievalRequest) (string, error), +) *MCPServerHandler { + return &MCPServerHandler{ + listDatasetsFunc: listDatasetsFunc, + listChatsFunc: listChatsFunc, + retrievalFunc: retrievalFunc, + } +} + +// HandleMCP is the Gin handler for the MCP endpoint. It reads the JSON-RPC +// request body, creates a connector for the authenticated user, and returns +// the JSON-RPC response. The endpoint is placed behind BetaAuthMiddleware +// so the user is already resolved from the Authorization header. +// +// @Summary MCP Endpoint (JSON-RPC over HTTP) +// @Tags mcp +// @Accept json +// @Produce json +// @Router /api/v1/mcp [post] +func (h *MCPServerHandler) HandleMCP(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + const maxMCPBodyBytes = 1 << 20 // 1 MiB + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxMCPBodyBytes) + body, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": common.CodeBadRequest, + "message": "Failed to read request body: " + err.Error(), + }) + return + } + + // Create a connector for this user. Each request gets its own connector + // so that user context is always correct. + connector := mcp.NewServiceConnector( + user.ID, + h.listDatasetsFunc, + h.listChatsFunc, + h.retrievalFunc, + ) + + server := mcp.NewServer(connector) + respBody, hasResponse, err := server.HandleRequest(body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "code": common.CodeServerError, + "message": "MCP server error: " + err.Error(), + }) + return + } + + if !hasResponse { + // Notification — no response per JSON-RPC spec. + c.Status(http.StatusAccepted) + return + } + + // The MCP protocol uses application/json with JSON-RPC responses. + c.Data(http.StatusOK, "application/json", respBody) +} + +// MCPListDatasets wraps DatasetService.ListDatasets for the MCP tool handler, +// filling in default values for parameters that the MCP tool does not expose. +func MCPListDatasets(ds *service.DatasetService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + data, total, _, err := ds.ListDatasets( + "", "", page, pageSize, orderby, desc, + "", nil, "", userID, + ) + return data, total, err +} + +// MCPListChats wraps ChatService.ListChats for the MCP tool handler, +// converting the typed response into a generic []map[string]interface{}. +func MCPListChats(cs *service.ChatService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) { + resp, err := cs.ListChats(userID, "1", "", page, pageSize, orderby, desc) + if err != nil { + return nil, 0, err + } + var chatList []map[string]interface{} + for _, chat := range resp.Chats { + chatList = append(chatList, map[string]interface{}{ + "id": chat.ID, + "name": chat.Name, + "description": chat.Description, + }) + } + return chatList, resp.Total, nil +} + +// MCPRetrieval executes a retrieval request on behalf of the MCP tool handler. +// It translates the mcp.RetrievalRequest into a service.SearchDatasetsRequest +// and calls DatasetService.SearchDatasets. The result is serialized as JSON. +func MCPRetrieval(ds *service.DatasetService, userID string, req mcp.RetrievalRequest) (string, error) { + // Resolve dataset IDs: if none provided, fetch ALL accessible datasets + // across all pages (matching Python _fetch_all_datasets behaviour). + datasetIDs := req.DatasetIDs + if len(datasetIDs) == 0 { + const maxPageSize = 100 + page := 1 + for { + data, _, _, err := ds.ListDatasets( + "", "", page, maxPageSize, "create_time", true, + "", nil, "", userID, + ) + if err != nil { + return "", fmt.Errorf("cannot resolve accessible datasets: %w", err) + } + if len(data) == 0 { + break + } + for _, d := range data { + if id, ok := d["id"].(string); ok && id != "" { + datasetIDs = append(datasetIDs, id) + } + } + // A page smaller than maxPageSize is the last page. + if len(data) < maxPageSize { + break + } + page++ + } + if len(datasetIDs) == 0 { + return "", fmt.Errorf("no accessible datasets found") + } + } + + searchReq := &service.SearchDatasetsRequest{ + DatasetIDs: datasetIDs, + Question: req.Question, + DocIDs: req.DocumentIDs, + ForceRefresh: req.ForceRefresh, + } + + if req.Page > 0 { + v := req.Page + searchReq.Page = &v + } + if req.PageSize > 0 { + v := req.PageSize + searchReq.Size = &v + } + if req.TopK > 0 { + v := req.TopK + searchReq.TopK = &v + } + { + v := req.SimilarityThreshold + searchReq.SimilarityThreshold = &v + } + { + v := req.VectorSimilarityWeight + searchReq.VectorSimilarityWeight = &v + } + if req.RerankID != "" { + v := req.RerankID + searchReq.RerankID = &v + } + { + v := req.Keyword + searchReq.Keyword = &v + } + + resp, err := ds.SearchDatasets(searchReq, userID) + if err != nil { + return "", err + } + + result, err := json.Marshal(resp) + if err != nil { + return "", fmt.Errorf("failed to serialize retrieval result: %w", err) + } + return string(result), nil +} diff --git a/internal/mcp/connector.go b/internal/mcp/connector.go new file mode 100644 index 0000000000..fbdfd23c88 --- /dev/null +++ b/internal/mcp/connector.go @@ -0,0 +1,130 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package mcp + +import ( + "encoding/json" + "fmt" + "strings" +) + +// ServiceConnector implements the Connector interface using in-process +// service layer calls, avoiding HTTP round-trips to self. +type ServiceConnector struct { + userID string + listDatasets func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + listChats func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) + retrieval func(userID string, req RetrievalRequest) (string, error) +} + +// NewServiceConnector creates a ServiceConnector. +// The function arguments abstract the service dependencies so this package +// does not import the service layer directly. +func NewServiceConnector( + userID string, + listDatasetsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + listChatsFunc func(userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error), + retrievalFunc func(userID string, req RetrievalRequest) (string, error), +) *ServiceConnector { + return &ServiceConnector{ + userID: userID, + listDatasets: listDatasetsFunc, + listChats: listChatsFunc, + retrieval: retrievalFunc, + } +} + +// ListDatasets returns newline-delimited JSON, each line being +// {"id": "...", "description": "..."} for a dataset. +func (c *ServiceConnector) ListDatasets(page, pageSize int, orderby string, desc bool) (string, error) { + data, _, err := c.listDatasets(c.userID, page, pageSize, orderby, desc) + if err != nil { + return "", fmt.Errorf("list datasets: %w", err) + } + + var lines []string + for _, d := range data { + id, _ := d["id"].(string) + name := "" + if v, ok := d["name"]; ok { + if s, ok := v.(string); ok { + name = s + } + } + desc := "" + if v, ok := d["description"]; ok { + if s, ok := v.(string); ok { + desc = s + } + } + // Match Python output: {"id": "...", "name": "...", "description": "..."} + item := map[string]interface{}{ + "id": id, + "name": name, + "description": desc, + } + b, err := json.Marshal(item) + if err != nil { + continue + } + lines = append(lines, string(b)) + } + return strings.Join(lines, "\n"), nil +} + +// ListChats returns newline-delimited JSON, each line being +// {"id": "...", "name": "...", "description": "..."} for a chat assistant. +func (c *ServiceConnector) ListChats(page, pageSize int, orderby string, desc bool) (string, error) { + data, _, err := c.listChats(c.userID, page, pageSize, orderby, desc) + if err != nil { + return "", fmt.Errorf("list chats: %w", err) + } + + var lines []string + for _, d := range data { + id, _ := d["id"].(string) + name := "" + if v, ok := d["name"]; ok { + if s, ok := v.(string); ok { + name = s + } + } + description := "" + if v, ok := d["description"]; ok { + if s, ok := v.(string); ok { + description = s + } + } + item := map[string]interface{}{ + "id": id, + "name": name, + "description": description, + } + b, err := json.Marshal(item) + if err != nil { + continue + } + lines = append(lines, string(b)) + } + return strings.Join(lines, "\n"), nil +} + +// Retrieval executes a retrieval via the in-process service and returns +// the result as a JSON string. +func (c *ServiceConnector) Retrieval(req RetrievalRequest) (string, error) { + return c.retrieval(c.userID, req) +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go new file mode 100644 index 0000000000..12b1a25aea --- /dev/null +++ b/internal/mcp/server.go @@ -0,0 +1,445 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package mcp + +import ( + "encoding/json" + "fmt" + "strings" +) + +// Connector provides the data-access operations the MCP tools need. +// It abstracts the RAGFlow backend so that tool implementations can use +// either in-process service calls or out-of-process HTTP calls. +type Connector interface { + // ListDatasets returns newline-delimited JSON lines, each containing + // at minimum {"id": "...", "description": "..."}. + ListDatasets(page, pageSize int, orderby string, desc bool) (string, error) + + // ListChats returns newline-delimited JSON lines, each containing + // at minimum {"id": "...", "name": "...", "description": "..."}. + ListChats(page, pageSize int, orderby string, desc bool) (string, error) + + // Retrieval executes a retrieval request and returns the result as + // a JSON string. + Retrieval(req RetrievalRequest) (string, error) +} + +// RetrievalRequest carries all parameters for a retrieval query. +type RetrievalRequest struct { + DatasetIDs []string `json:"dataset_ids"` + DocumentIDs []string `json:"document_ids"` + Question string `json:"question"` + Page int `json:"page"` + PageSize int `json:"page_size"` + SimilarityThreshold float64 `json:"similarity_threshold"` + VectorSimilarityWeight float64 `json:"vector_similarity_weight"` + TopK int `json:"top_k"` + RerankID string `json:"rerank_id,omitempty"` + Keyword bool `json:"keyword"` + ForceRefresh bool `json:"force_refresh"` +} + +// Server handles MCP JSON-RPC requests. +type Server struct { + connector Connector + version string +} + +// NewServer creates a new MCP Server. +func NewServer(connector Connector) *Server { + return &Server{ + connector: connector, + version: "1.0.0", + } +} + +// HandleRequest dispatches a raw JSON-RPC request body and returns the +// serialized JSON-RPC response. Returns nil if the request is a +// notification (no id) and requires no response. +func (s *Server) HandleRequest(body []byte) ([]byte, bool, error) { + // Try to decode as a request (with an id) first. + var req JSONRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + resp := NewParseError() + data, _ := json.Marshal(resp) + return data, true, nil + } + + // Notifications have no id — do not send a response. + if req.ID == nil || string(req.ID) == "null" { + return nil, false, nil + } + + // Validate jsonrpc field. + if req.JSONRPC != JSONRPCVersion { + resp := NewInvalidRequestError(req.ID, "jsonrpc must be \"2.0\"") + data, _ := json.Marshal(resp) + return data, true, nil + } + + var resp JSONRPCResponse + + switch req.Method { + case "initialize": + resp = s.handleInitialize(req.ID) + case "tools/list": + resp = s.handleListTools(req.ID) + case "tools/call": + resp = s.handleCallTool(req.ID, req.Params) + case "ping": + resp = s.handlePing(req.ID) + default: + resp = NewErrorResponse(req.ID, ErrCodeMethodNotFound, + fmt.Sprintf("Method not found: %s", req.Method)) + } + + data, err := json.Marshal(resp) + if err != nil { + return nil, true, fmt.Errorf("failed to marshal response: %w", err) + } + return data, true, nil +} + +func (s *Server) handleInitialize(id json.RawMessage) JSONRPCResponse { + result := InitializeResult{ + ProtocolVersion: MCPProtocolVersion, + Capabilities: Capabilities{ + Tools: &ToolsCapability{ListChanged: false}, + }, + ServerInfo: ServerInfo{ + Name: ServerName, + Version: s.version, + }, + } + return NewSuccessResponse(id, result) +} + +func (s *Server) handlePing(id json.RawMessage) JSONRPCResponse { + return NewSuccessResponse(id, struct{}{}) +} + +func (s *Server) handleListTools(id json.RawMessage) JSONRPCResponse { + // Fetch dataset and chat descriptions for embedding into tool descriptions, + // matching the Python MCP server behavior. + datasetDescription, err := s.connector.ListDatasets(1, 100, "create_time", true) + if err != nil { + datasetDescription = "" + } + chatDescription, err := s.connector.ListChats(1, 30, "create_time", true) + if err != nil { + chatDescription = "" + } + + tools := []Tool{ + { + Name: "ragflow_retrieval", + Description: "Retrieve relevant chunks from the RAGFlow retrieve interface based on the question. You can optionally specify dataset_ids to search only specific datasets, or omit dataset_ids entirely to search across ALL available datasets. You can also optionally specify document_ids to search within specific documents. When dataset_ids is not provided or is empty, the system will automatically search across all available datasets. Below is the list of all available datasets, including their descriptions and IDs:" + datasetDescription, + InputSchema: InputSchema{ + Type: "object", + Properties: map[string]Property{ + "dataset_ids": { + Type: "array", + Description: "Optional array of dataset IDs to search. If not provided or empty, all datasets will be searched.", + Items: &Items{Type: "string"}, + }, + "document_ids": { + Type: "array", + Description: "Optional array of document IDs to search within.", + Items: &Items{Type: "string"}, + }, + "question": { + Type: "string", + Description: "The question or query to search for.", + }, + "page": { + Type: "integer", + Description: "Page number for pagination", + Default: 1, + Minimum: float64Ptr(1), + }, + "page_size": { + Type: "integer", + Description: "Number of results to return per page (default: 10, max recommended: 50 to avoid token limits)", + Default: 10, + Minimum: float64Ptr(1), + Maximum: float64Ptr(100), + }, + "similarity_threshold": { + Type: "number", + Description: "Minimum similarity threshold for results", + Default: 0.2, + Minimum: float64Ptr(0), + Maximum: float64Ptr(1), + }, + "vector_similarity_weight": { + Type: "number", + Description: "Weight for vector similarity vs term similarity", + Default: 0.3, + Minimum: float64Ptr(0), + Maximum: float64Ptr(1), + }, + "keyword": { + Type: "boolean", + Description: "Enable keyword-based search", + Default: false, + }, + "top_k": { + Type: "integer", + Description: "Maximum results to consider before ranking", + Default: 1024, + Minimum: float64Ptr(1), + Maximum: float64Ptr(1024), + }, + "rerank_id": { + Type: "string", + Description: "Optional reranking model identifier", + }, + "force_refresh": { + Type: "boolean", + Description: "Set to true only if fresh dataset and document metadata is explicitly required. Otherwise, cached metadata is used (default: false).", + Default: false, + }, + }, + Required: []string{"question"}, + }, + }, + { + Name: "ragflow_list_datasets", + Description: "List all accessible datasets (knowledge bases) in RAGFlow. Returns dataset IDs, names, and descriptions. Use this tool to discover which datasets are available before performing retrieval." + datasetDescription, + InputSchema: InputSchema{ + Type: "object", + Properties: map[string]Property{ + "page": { + Type: "integer", + Description: "Page number", + Default: 1, + Minimum: float64Ptr(1), + }, + "page_size": { + Type: "integer", + Description: "Results per page", + Default: 100, + Minimum: float64Ptr(1), + Maximum: float64Ptr(1000), + }, + }, + }, + }, + { + Name: "ragflow_list_chats", + Description: "List all accessible chat assistants in RAGFlow. Returns chat assistant IDs, names, and descriptions. Use this tool to discover available chat assistants that can be used for conversations." + chatDescription, + InputSchema: InputSchema{ + Type: "object", + Properties: map[string]Property{ + "page": { + Type: "integer", + Description: "Page number", + Default: 1, + Minimum: float64Ptr(1), + }, + "page_size": { + Type: "integer", + Description: "Results per page", + Default: 30, + Minimum: float64Ptr(1), + Maximum: float64Ptr(100), + }, + }, + }, + }, + } + + return NewSuccessResponse(id, ListToolsResult{Tools: tools}) +} + +func (s *Server) handleCallTool(id json.RawMessage, rawParams json.RawMessage) JSONRPCResponse { + var params CallToolParams + if err := json.Unmarshal(rawParams, ¶ms); err != nil { + return NewErrorResponse(id, ErrCodeInvalidParams, + fmt.Sprintf("Invalid params: %s", err.Error())) + } + + if params.Arguments == nil { + params.Arguments = make(map[string]interface{}) + } + + switch params.Name { + case "ragflow_retrieval": + return s.callRagflowRetrieval(id, params.Arguments) + case "ragflow_list_datasets": + return s.callListDatasets(id, params.Arguments) + case "ragflow_list_chats": + return s.callListChats(id, params.Arguments) + default: + return NewErrorResponse(id, ErrCodeMethodNotFound, + fmt.Sprintf("Tool not found: %s", params.Name)) + } +} + +func (s *Server) callRagflowRetrieval(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { + req := RetrievalRequest{ + Page: getBoundedIntArg(args, "page", 1, 1, 1_000_000), + PageSize: getBoundedIntArg(args, "page_size", 10, 1, 100), + SimilarityThreshold: getBoundedFloat64Arg(args, "similarity_threshold", 0.2, 0, 1), + VectorSimilarityWeight: getBoundedFloat64Arg(args, "vector_similarity_weight", 0.3, 0, 1), + Keyword: getBoolArg(args, "keyword", false), + TopK: getBoundedIntArg(args, "top_k", 1024, 1, 1024), + ForceRefresh: getBoolArg(args, "force_refresh", false), + Question: getStringArg(args, "question", ""), + RerankID: getStringArg(args, "rerank_id", ""), + } + + if v, ok := args["dataset_ids"]; ok { + req.DatasetIDs = toStringSlice(v) + } + if v, ok := args["document_ids"]; ok { + req.DocumentIDs = toStringSlice(v) + } + + if strings.TrimSpace(req.Question) == "" { + return NewSuccessResponse(id, NewErrorResult("question is required")) + } + + result, err := s.connector.Retrieval(req) + if err != nil { + return NewSuccessResponse(id, NewErrorResult(err.Error())) + } + return NewSuccessResponse(id, NewTextResult(result)) +} + +func (s *Server) callListDatasets(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { + page := getBoundedIntArg(args, "page", 1, 1, 1_000_000) + pageSize := getBoundedIntArg(args, "page_size", 100, 1, 1000) + + result, err := s.connector.ListDatasets(page, pageSize, "create_time", true) + if err != nil { + return NewSuccessResponse(id, NewErrorResult(err.Error())) + } + return NewSuccessResponse(id, NewTextResult(result)) +} + +func (s *Server) callListChats(id json.RawMessage, args map[string]interface{}) JSONRPCResponse { + page := getBoundedIntArg(args, "page", 1, 1, 1_000_000) + pageSize := getBoundedIntArg(args, "page_size", 30, 1, 100) + + result, err := s.connector.ListChats(page, pageSize, "create_time", true) + if err != nil { + return NewSuccessResponse(id, NewErrorResult(err.Error())) + } + return NewSuccessResponse(id, NewTextResult(result)) +} + +// --- argument extraction helpers --- + +func getStringArg(args map[string]interface{}, key, defaultVal string) string { + if v, ok := args[key]; ok { + if s, ok := v.(string); ok { + return s + } + } + return defaultVal +} + +func getIntArg(args map[string]interface{}, key string, defaultVal int) int { + if v, ok := args[key]; ok { + switch n := v.(type) { + case float64: + return int(n) + case int: + return n + case int64: + return int(n) + case json.Number: + if i, err := n.Int64(); err == nil { + return int(i) + } + } + } + return defaultVal +} + +func getBoundedIntArg(args map[string]interface{}, key string, defaultVal, minVal, maxVal int) int { + v := getIntArg(args, key, defaultVal) + if v < minVal { + return minVal + } + if v > maxVal { + return maxVal + } + return v +} + +func getFloat64Arg(args map[string]interface{}, key string, defaultVal float64) float64 { + if v, ok := args[key]; ok { + switch f := v.(type) { + case float64: + return f + case int: + return float64(f) + case int64: + return float64(f) + case json.Number: + if fl, err := f.Float64(); err == nil { + return fl + } + } + } + return defaultVal +} + +func getBoundedFloat64Arg(args map[string]interface{}, key string, defaultVal, minVal, maxVal float64) float64 { + v := getFloat64Arg(args, key, defaultVal) + if v < minVal { + return minVal + } + if v > maxVal { + return maxVal + } + return v +} + +func getBoolArg(args map[string]interface{}, key string, defaultVal bool) bool { + if v, ok := args[key]; ok { + if b, ok := v.(bool); ok { + return b + } + if s, ok := v.(string); ok { + switch strings.ToLower(s) { + case "true", "1", "yes": + return true + case "false", "0", "no": + return false + } + } + } + return defaultVal +} + +// toStringSlice converts an interface{} value (expected to be a JSON array) +// into a []string. Values are converted via fmt.Sprintf. +func toStringSlice(v interface{}) []string { + arr, ok := v.([]interface{}) + if !ok { + return nil + } + result := make([]string, 0, len(arr)) + for _, item := range arr { + result = append(result, fmt.Sprintf("%v", item)) + } + return result +} diff --git a/internal/mcp/types.go b/internal/mcp/types.go new file mode 100644 index 0000000000..fa7d407e31 --- /dev/null +++ b/internal/mcp/types.go @@ -0,0 +1,225 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// Package mcp implements the MCP (Model Context Protocol) server embedded in +// the RAGFlow Go backend. It exposes RAGFlow capabilities (retrieval, dataset +// listing, chat listing) as MCP tools that external AI clients can discover +// and invoke via JSON-RPC over HTTP. +package mcp + +import ( + "encoding/json" + "fmt" +) + +// JSONRPCVersion is the protocol version string. +const JSONRPCVersion = "2.0" + +// MCPProtocolVersion is the MCP protocol version this server implements. +const MCPProtocolVersion = "2024-11-05" + +// ServerName identifies this MCP server instance. +const ServerName = "ragflow-mcp-server" + +// JSONRPCRequest represents a JSON-RPC 2.0 request. +type JSONRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +// JSONRPCResponse represents a JSON-RPC 2.0 response. +type JSONRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result interface{} `json:"result,omitempty"` + Error *JSONRPCError `json:"error,omitempty"` +} + +// JSONRPCNotification represents a JSON-RPC 2.0 notification +// (a request without an id, which requires no response). +type JSONRPCNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +// JSONRPCError represents a JSON-RPC 2.0 error object. +type JSONRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data interface{} `json:"data,omitempty"` +} + +// Predefined JSON-RPC error codes. +const ( + ErrCodeParseError = -32700 + ErrCodeInvalidRequest = -32600 + ErrCodeMethodNotFound = -32601 + ErrCodeInvalidParams = -32602 + ErrCodeInternalError = -32603 +) + +// InitializeResult is the response payload for the "initialize" method. +type InitializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities Capabilities `json:"capabilities"` + ServerInfo ServerInfo `json:"serverInfo"` +} + +// Capabilities describes the set of MCP capabilities this server supports. +type Capabilities struct { + Tools *ToolsCapability `json:"tools,omitempty"` +} + +// ToolsCapability indicates that the server supports tools and optionally +// whether tool list changes are notified. +type ToolsCapability struct { + ListChanged bool `json:"listChanged,omitempty"` +} + +// ServerInfo provides identifying information about the MCP server. +type ServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// Tool represents an MCP tool definition. +type Tool struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema InputSchema `json:"inputSchema"` +} + +// InputSchema is the JSON Schema for a tool's input parameters. +type InputSchema struct { + Type string `json:"type"` + Properties map[string]Property `json:"properties,omitempty"` + Required []string `json:"required,omitempty"` +} + +// Property describes a single parameter in a tool's InputSchema. +type Property struct { + Type string `json:"type"` + Description string `json:"description,omitempty"` + Default interface{} `json:"default,omitempty"` + Minimum *float64 `json:"minimum,omitempty"` + Maximum *float64 `json:"maximum,omitempty"` + Items *Items `json:"items,omitempty"` +} + +// Items describes the expected element type for array-typed properties. +type Items struct { + Type string `json:"type"` +} + +// ListToolsResult is the result of the "tools/list" method. +type ListToolsResult struct { + Tools []Tool `json:"tools"` +} + +// CallToolParams is the params payload for the "tools/call" method. +type CallToolParams struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments,omitempty"` +} + +// CallToolResult is the result of the "tools/call" method. +type CallToolResult struct { + Content []ContentItem `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// ContentItem is a single item in a CallToolResult's content array. +type ContentItem struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + Data string `json:"data,omitempty"` + // MIME Type of the data, if present. + MIMEType string `json:"mimeType,omitempty"` +} + +// NewErrorResponse creates a JSONRPCResponse with an error. +func NewErrorResponse(id json.RawMessage, code int, message string) JSONRPCResponse { + return JSONRPCResponse{ + JSONRPC: JSONRPCVersion, + ID: id, + Error: &JSONRPCError{ + Code: code, + Message: message, + }, + } +} + +// NewSuccessResponse creates a JSONRPCResponse with a result. +func NewSuccessResponse(id json.RawMessage, result interface{}) JSONRPCResponse { + return JSONRPCResponse{ + JSONRPC: JSONRPCVersion, + ID: id, + Result: result, + } +} + +// NewTextContent creates a ContentItem of type "text". +func NewTextContent(text string) ContentItem { + return ContentItem{Type: "text", Text: text} +} + +// NewTextResult creates a CallToolResult with a single text content item. +func NewTextResult(text string) *CallToolResult { + return &CallToolResult{ + Content: []ContentItem{NewTextContent(text)}, + } +} + +// NewErrorResult creates a CallToolResult indicating a tool execution error. +func NewErrorResult(errMsg string) *CallToolResult { + return &CallToolResult{ + Content: []ContentItem{NewTextContent(errMsg)}, + IsError: true, + } +} + +// NewParseError creates a standard Parse Error response (id is set to null). +func NewParseError() JSONRPCResponse { + return JSONRPCResponse{ + JSONRPC: JSONRPCVersion, + ID: nil, + Error: &JSONRPCError{ + Code: ErrCodeParseError, + Message: "Parse error", + }, + } +} + +// NewInvalidRequestError creates a standard Invalid Request response. +func NewInvalidRequestError(id json.RawMessage, msg string) JSONRPCResponse { + return JSONRPCResponse{ + JSONRPC: JSONRPCVersion, + ID: id, + Error: &JSONRPCError{ + Code: ErrCodeInvalidRequest, + Message: fmt.Sprintf("Invalid Request: %s", msg), + }, + } +} + +// float64Ptr returns a pointer to a float64 value, used for Property +// minimum/maximum defaults. +func float64Ptr(v float64) *float64 { + return &v +} diff --git a/internal/router/router.go b/internal/router/router.go index 826ed379dc..5b3a8967fc 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -43,6 +43,7 @@ type Router struct { fileHandler *handler.FileHandler memoryHandler *handler.MemoryHandler mcpHandler *handler.MCPHandler + mcpServerHandler *handler.MCPServerHandler skillSearchHandler *handler.SkillSearchHandler providerHandler *handler.ProviderHandler agentHandler *handler.AgentHandler @@ -75,6 +76,7 @@ func NewRouter( fileHandler *handler.FileHandler, memoryHandler *handler.MemoryHandler, mcpHandler *handler.MCPHandler, + mcpServerHandler *handler.MCPServerHandler, skillSearchHandler *handler.SkillSearchHandler, providerHandler *handler.ProviderHandler, agentHandler *handler.AgentHandler, @@ -107,6 +109,7 @@ func NewRouter( fileHandler: fileHandler, memoryHandler: memoryHandler, mcpHandler: mcpHandler, + mcpServerHandler: mcpServerHandler, skillSearchHandler: skillSearchHandler, providerHandler: providerHandler, agentHandler: agentHandler, @@ -208,6 +211,13 @@ func (r *Router) Setup(engine *gin.Engine) { apiBetaAuth.GET("/documents/images/:image_id", r.documentHandler.GetDocumentImage) apiBetaAuth.GET("/documents/:id/preview", r.documentHandler.GetDocumentPreview) apiBetaAuth.GET("/thumbnails", r.documentHandler.GetThumbnail) + + // MCP server endpoint — exposes RAGFlow capabilities as MCP tools. + // Uses BetaAuthMiddleware to resolve the user from the + // Authorization header. + if r.mcpServerHandler != nil { + apiBetaAuth.POST("/mcp", r.mcpServerHandler.HandleMCP) + } } // Protected routes diff --git a/internal/service/agent.go b/internal/service/agent.go index 0bdca754f6..4cf98ed062 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -256,7 +256,7 @@ func toAgentItem(c *dao.UserCanvasListItem) *AgentItem { // ListAgents returns agent canvases visible to userID. // Mirrors Python agent_api.list_agents — validates owner_ids against joined tenants, // then delegates to the DAO. -func (s *AgentService) ListAgents(userID string, keywords string, page, pageSize int, orderby string, desc bool, ownerIDs []string, canvasCategory string, tags []string) (*ListAgentsResponse, common.ErrorCode, error) { +func (s *AgentService) ListAgents(userID string, keywords string, page, pageSize int, orderby string, desc bool, ownerIDs []string, canvasCategory, canvasType string, tags []string) (*ListAgentsResponse, common.ErrorCode, error) { // Build the set of tenant IDs the user is authorised to query. tenantIDs, err := s.userTenantDAO.GetTenantIDsByUserID(userID) if err != nil { @@ -292,6 +292,7 @@ func (s *AgentService) ListAgents(userID string, keywords string, page, pageSize desc, keywords, canvasCategory, + canvasType, tags, ) if err != nil { diff --git a/internal/service/agent_sessions.go b/internal/service/agent_sessions.go index 593b59b8b9..c2f3e6de38 100644 --- a/internal/service/agent_sessions.go +++ b/internal/service/agent_sessions.go @@ -101,9 +101,9 @@ func parseAgentSessionDate(value string, isEnd bool) (*time.Time, error) { return nil, nil } + // Try flexible ISO 8601 / RFC3339 parsing first (PR #16483). if strings.Contains(value, "T") { - normalized := strings.ReplaceAll(value, "Z", "+00:00") - parsed, err := time.Parse(time.RFC3339, normalized) + parsed, err := common.ParseISO8601(value) if err != nil { return nil, err } diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index 40f6960cf0..3677eb39c6 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -44,6 +44,7 @@ import ( "go.uber.org/zap" "ragflow/internal/agent/canvas" + "ragflow/internal/agent/runtime" "ragflow/internal/common" "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" @@ -115,7 +116,10 @@ func WriteChatbotFrame(w http.ResponseWriter, f ChatbotSSEFrame) error { "data": data, } } - b, err := json.Marshal(payload) + // Use SafeJSONMarshal to handle non-serializable values (funcs, + // channels) that may have leaked into SSE payload maps. Mirrors + // the Python PR #14210 _canvas_json_default fallback in agent_api.py. + b, err := runtime.SafeJSONMarshal(payload) if err != nil { return err } diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index 97c120d65d..256d918eb3 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -213,8 +213,10 @@ func (s *ChatPipelineService) AsyncChat( } modelMaxTokens := 8192 if llmModelConfig != nil { - if mt, ok := llmModelConfig["max_tokens"].(float64); ok { - modelMaxTokens = int(mt) + // Treat max_tokens=0 as unset (default 8192) — mirrors + // PR #16413 Python fix: model_extra.get("max_tokens") or 8192 + if mt, ok := llmModelConfig["max_tokens"].(int); ok && mt > 0 { + modelMaxTokens = mt } } timer.Exit(common.PhaseCheckLLM) @@ -2385,38 +2387,81 @@ func (s *ChatPipelineService) formatPrompt(template string, kwargs map[string]in // messageFitIn trims messages to fit within a token budget. // Mirrors Python's message_fit_in() in rag/prompts/generator.py. +// +// Strategy: +// 1. If everything fits → return as-is. +// 2. Keep all system messages + the last user/assistant message. +// 3. If still too large, trim content proportionally: +// - System dominates (>80%) → preserve last message first. +// - Otherwise → preserve system first. func (s *ChatPipelineService) messageFitIn(messages []map[string]interface{}, maxTokens int) (int, []map[string]interface{}) { - totalTokens := 0 - var result []map[string]interface{} - - // Always keep the system message first. - if len(messages) > 0 { - if role, _ := messages[0]["role"].(string); role == "system" { - if content, ok := messages[0]["content"].(string); ok { - sysTokens := kg.NumTokensFromString(content) - if sysTokens <= maxTokens { - totalTokens = sysTokens - result = append(result, messages[0]) - } - } - } + if maxTokens <= 0 { + maxTokens = 8192 } - // Add user/assistant messages from the end, working backwards. - rest := messages[1:] - for i := len(rest) - 1; i >= 0; i-- { - m := rest[i] - content, _ := m["content"].(string) - tokens := kg.NumTokensFromString(content) - if totalTokens+tokens > maxTokens { - break - } - totalTokens += tokens - // Prepend to maintain order. - result = append([]map[string]interface{}{m}, result...) + // Step 1: everything fits. + totalTokens := s.countAllTokens(messages) + if totalTokens < maxTokens { + return totalTokens, messages } - return totalTokens, result + // Step 2: keep all system messages + the last message. + result := make([]map[string]interface{}, 0) + for _, m := range messages { + if role, _ := m["role"].(string); role == "system" { + result = append(result, m) + } + } + if len(messages) > 1 { + result = append(result, messages[len(messages)-1]) + } + + totalTokens = s.countAllTokens(result) + if totalTokens < maxTokens { + return totalTokens, result + } + + // Step 3: trim content to fit. + ll := kg.NumTokensFromString(s.stringContent(result[0])) + ll2 := kg.NumTokensFromString(s.stringContent(result[len(result)-1])) + total := ll + ll2 + if total <= 0 { + return 0, result + } + + if len(result) == 1 { + result[0]["content"] = kg.TrimContentToTokenLimit(s.stringContent(result[0]), maxTokens) + return s.countAllTokens(result), result + } + + if float64(ll)/float64(total) > 0.8 { + preservedLast := min(ll2, maxTokens) + result[len(result)-1]["content"] = kg.TrimContentToTokenLimit(s.stringContent(result[len(result)-1]), preservedLast) + remaining := max(0, maxTokens-preservedLast) + result[0]["content"] = kg.TrimContentToTokenLimit(s.stringContent(result[0]), remaining) + } else { + preservedSystem := min(ll, maxTokens) + result[0]["content"] = kg.TrimContentToTokenLimit(s.stringContent(result[0]), preservedSystem) + remaining := max(0, maxTokens-preservedSystem) + result[len(result)-1]["content"] = kg.TrimContentToTokenLimit(s.stringContent(result[len(result)-1]), remaining) + } + + return s.countAllTokens(result), result +} + +// countAllTokens returns the total token count across all messages. +func (s *ChatPipelineService) countAllTokens(messages []map[string]interface{}) int { + total := 0 + for _, m := range messages { + total += kg.NumTokensFromString(s.stringContent(m)) + } + return total +} + +// stringContent extracts the "content" string from a message map, or "". +func (s *ChatPipelineService) stringContent(m map[string]interface{}) string { + c, _ := m["content"].(string) + return c } // buildChatMessages converts the internal message representation to diff --git a/internal/service/dataset.go b/internal/service/dataset.go index 041c092f2e..aade8907de 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -1708,6 +1708,7 @@ type SearchDatasetsRequest struct { Keyword *bool `json:"keyword,omitempty"` SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"` VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"` + ForceRefresh bool `json:"force_refresh"` } // SearchDatasetsResponse is the response structure for dataset search results. diff --git a/internal/service/kg/scoring.go b/internal/service/kg/scoring.go index 61896c9942..b6e4c339a8 100644 --- a/internal/service/kg/scoring.go +++ b/internal/service/kg/scoring.go @@ -153,6 +153,12 @@ func NumTokensFromString(s string) int { return tokenizer.NumTokensFromString(s) } +// TrimContentToTokenLimit truncates s to at most limit tokens. +// Delegates to the shared implementation in the tokenizer package. +func TrimContentToTokenLimit(s string, limit int) string { + return tokenizer.TrimContentToTokenLimit(s, limit) +} + // formatCSVLine formats fields as a single CSV record with trailing newline. // Handles commas, quotes, and newlines in field values correctly — unlike fmt.Sprintf. // Matches Python: pd.DataFrame(...).to_csv() quoting behavior. diff --git a/internal/tokenizer/tokenizer.go b/internal/tokenizer/tokenizer.go index f9bde3b7e8..b303210542 100644 --- a/internal/tokenizer/tokenizer.go +++ b/internal/tokenizer/tokenizer.go @@ -25,6 +25,7 @@ import ( "sync" "sync/atomic" "time" + "unicode/utf8" "github.com/pkoukk/tiktoken-go" "go.uber.org/zap" @@ -533,3 +534,33 @@ func NumTokensFromString(s string) int { } return len(enc.Encode(s, nil, nil)) } + +// TrimContentToTokenLimit truncates s to at most limit tokens using the +// cl100k_base encoder. Mirrors Python's trim_content helper in +// rag/prompts/generator.py: encoder.decode(encoder.encode(content)[:limit]). +// Returns the original string if it already fits. +func TrimContentToTokenLimit(s string, limit int) string { + if limit < 0 { + limit = 0 + } + enc, err := getCL100KEncoder() + if err != nil { + // Fail closed: fall back to byte-length trimming with UTF-8 safety. + if limit <= 0 { + return "" + } + b := []byte(s) + if len(b) <= limit { + return s + } + for limit > 0 && !utf8.Valid(b[:limit]) { + limit-- + } + return string(b[:limit]) + } + tokens := enc.Encode(s, nil, nil) + if len(tokens) <= limit { + return s + } + return enc.Decode(tokens[:limit]) +}