Add get_processed_document MCP tool and chain it from get_chunks#286
Add get_processed_document MCP tool and chain it from get_chunks#286PuneetPunamiya wants to merge 1 commit into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the unstructured data MCP server by introducing a mechanism to retrieve full processed documents. By chaining the existing chunk search tool with a new document retrieval tool, the system now allows LLM agents to fetch complete markdown content after identifying relevant chunks. The changes also include architectural improvements to pipeline query configuration and standardized error messaging. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new MCP tool, get_processed_document, which retrieves full processed documents from a pipeline's DocumentProcessor stage in Snowflake. It also updates the get_chunks_for_embeddings tool to return file_id values and direct users to the new tool, and modifies the Kubernetes client to fetch query configurations by stage type. The review feedback recommends only appending the 'NEXT STEP' instruction when chunks are actually found to prevent LLM hallucinations, and moving the shared errOAuthTokenNotFound constant to a common file to avoid implicit cross-file dependencies.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| "github.com/redhat-data-and-ai/unstructured-data-controller/pkg/snowflake" | ||
| ) | ||
|
|
||
| const errOAuthTokenNotFound = "Error: oauth token not found in context" |
There was a problem hiding this comment.
Defining the shared package-level constant errOAuthTokenNotFound in a specific tool file (get_processed_document.go) creates an implicit dependency for other files in the same package (like get_chunks.go and list_pipelines.go). If this file is excluded, renamed, or deleted, compilation of other tools will break. Consider moving shared constants to a dedicated common.go or constants.go file within the tools package to improve maintainability and modularity.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a new MCP tool get_processed_document to retrieve full processed documents from Snowflake by file_id. It updates the existing get_chunks_for_embeddings tool to return the file_id for each chunk and suggest calling get_processed_document as a next step. It also refactors GetPipelineQueryConfig to support querying configurations for different stage types. Feedback was provided to ensure the NEXT STEP instruction in get_chunks_for_embeddings is only appended when a valid file_id is actually found in the results, preventing misleading tool calls when no chunks are returned.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
- Only show get_processed_document TIP when chunks are non-empty - Soften tool description to let the agent decide when to fetch full documents - Use user-friendly error message for missing query config
PuneetPunamiya
force-pushed
the
worktree-trim-pipeline-info
branch
from
d4c9276 to
4a52198
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a new MCP tool, get_processed_document, which allows users to retrieve full processed documents from a pipeline's DocumentProcessor stage in Snowflake. It also updates the get_chunks_for_embeddings tool to return a file_id and suggest using the new tool for additional context, and refactors GetPipelineQueryConfig to support querying configurations for different stage types. The review feedback suggests several key improvements: simplifying the next-step tip logic to only trigger when a valid file_id is present, returning a more precise and actionable error message when the pipeline query configuration is incomplete, and outputting the raw markdown content directly instead of JSON-marshaling the result to preserve readability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| var nextStep string | ||
| if len(chunks) > 0 && chunks[0].FileID != "" { | ||
| nextStep = fmt.Sprintf( | ||
| "\n\nTIP: To retrieve the full processed document, call get_processed_document with pipeline_name=%q and file_id=%q.", | ||
| args.PipelineName, chunks[0].FileID, | ||
| ) | ||
| } else if len(chunks) > 0 { | ||
| nextStep = fmt.Sprintf( | ||
| "\n\nTIP: To retrieve the full processed document, call get_processed_document with pipeline_name=%q and the file_id from the highest-scoring chunk above.", | ||
| args.PipelineName, | ||
| ) | ||
| } |
There was a problem hiding this comment.
If chunks[0].FileID is empty, then the highest-scoring chunk does not have a valid file_id. In this scenario, prompting the LLM or user to call get_processed_document with "the file_id from the highest-scoring chunk above" is misleading and will fail because no file_id is available.
We should simplify this logic to only suggest calling get_processed_document if we actually have a non-empty FileID in the top chunk.
var nextStep string
if len(chunks) > 0 && chunks[0].FileID != "" {
nextStep = fmt.Sprintf(
"\n\nTIP: To retrieve the full processed document, call get_processed_document with pipeline_name=%q and file_id=%q.",
args.PipelineName, chunks[0].FileID,
)
}| if qc.Database == "" || qc.Schema == "" || qc.Table == "" { | ||
| log.Error("pipeline query config has empty fields", "database", qc.Database, "schema", qc.Schema, "table", qc.Table) | ||
| return &mcp.CallToolResult{ | ||
| Content: []mcp.Content{&mcp.TextContent{Text: "Unable to retrieve the document, please check again in some time"}}, | ||
| IsError: true, | ||
| }, nil, nil | ||
| } |
There was a problem hiding this comment.
The error message returned when the pipeline query configuration is incomplete is vague ("Unable to retrieve the document, please check again in some time"). This implies a transient error, which might cause the LLM to retry or wait.
Instead, we should return a clear, actionable error message indicating that the pipeline has an incomplete Snowflake query configuration, matching the behavior in get_chunks.go.
| if qc.Database == "" || qc.Schema == "" || qc.Table == "" { | |
| log.Error("pipeline query config has empty fields", "database", qc.Database, "schema", qc.Schema, "table", qc.Table) | |
| return &mcp.CallToolResult{ | |
| Content: []mcp.Content{&mcp.TextContent{Text: "Unable to retrieve the document, please check again in some time"}}, | |
| IsError: true, | |
| }, nil, nil | |
| } | |
| if qc.Database == "" || qc.Schema == "" || qc.Table == "" { | |
| log.Error("pipeline query config has empty fields", "database", qc.Database, "schema", qc.Schema, "table", qc.Table) | |
| return &mcp.CallToolResult{ | |
| Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Error: pipeline %q has incomplete Snowflake query configuration", args.PipelineName)}}, | |
| IsError: true, | |
| }, nil, nil | |
| } |
| jsonBytes, err := json.Marshal(doc) | ||
| if err != nil { | ||
| log.Error("failed to marshal result", "error", err) | ||
| return &mcp.CallToolResult{ | ||
| Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Error marshaling result: %v", err)}}, | ||
| IsError: true, | ||
| }, nil, nil | ||
| } | ||
|
|
||
| log.Info("completed successfully", "pipeline", args.PipelineName, "file_id", args.FileID) | ||
| return &mcp.CallToolResult{ | ||
| Content: []mcp.Content{&mcp.TextContent{ | ||
| Text: fmt.Sprintf("Processed document for file_id %q in pipeline %q:\n%s", args.FileID, args.PipelineName, string(jsonBytes)), | ||
| }}, | ||
| }, nil, nil |
There was a problem hiding this comment.
Since get_processed_document retrieves a single document's markdown content, JSON-marshaling the entire ProcessedDocumentResult struct escapes all newlines (\n) and special characters. This makes the output extremely hard to read for humans and less clean for LLM processing.
Instead, we should return the raw MarkdownContent directly in a nicely formatted string. This also eliminates the need for json.Marshal and its error handling.
Note: Remember to also remove the unused "encoding/json" import from the top of the file.
log.Info("completed successfully", "pipeline", args.PipelineName, "file_id", args.FileID)
return &mcp.CallToolResult{
Content: []mcp.Content{&mcp.TextContent{
Text: fmt.Sprintf("Processed document for file_id %q in pipeline %q:\n\n%s", args.FileID, args.PipelineName, doc.MarkdownContent),
}},
}, nil, nil
| var nextStep string | ||
| if len(chunks) > 0 && chunks[0].FileID != "" { | ||
| nextStep = fmt.Sprintf( | ||
| "\n\nTIP: To retrieve the full processed document, call get_processed_document with pipeline_name=%q and file_id=%q.", | ||
| args.PipelineName, chunks[0].FileID, | ||
| ) | ||
| } else if len(chunks) > 0 { | ||
| nextStep = fmt.Sprintf( | ||
| "\n\nTIP: To retrieve the full processed document, call get_processed_document with pipeline_name=%q and the file_id from the highest-scoring chunk above.", |
There was a problem hiding this comment.
I wonder if fileID field is going to be empty for any record. We should probably return an error in such scenarios. @concaf WDYT?
2 participants
Summary
get_processed_documentMCP tool retrieves full markdown content from the DocumentProcessor stage table byfile_idget_chunks_for_embeddingsnow returnsfile_idin results and prompts the LLM agent to callget_processed_documentas a follow-up for the full documentGetPipelineQueryConfigto accept a stage type parameter so it works for both VectorEmbeddingsGenerator and DocumentProcessor stagessnowflake.GetProcessedDocument()to query the processed document tableerrOAuthTokenNotFoundconstant across MCP toolsTest plan
get_chunks_for_embeddingsreturnsfile_idin chunk resultsget_chunks_for_embeddingsresponse includes the NEXT STEP instruction to callget_processed_documentget_processed_documentretrieves the correct document byfile_idfrom the DocumentProcessor stage tableGetPipelineQueryConfigcorrectly resolves query config for bothStageTypeVectorEmbeddingsGeneratorandStageTypeDocumentProcessor