feat: implement Pydantic AI dependency injection with RunContext and deps_type #190
Conversation
Summary of ChangesHello @vitali87, 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 significantly refactors the codebase to implement a robust dependency injection pattern using Pydantic AI's 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 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 counter productive. 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. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. 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 significant refactoring of dependency management and tool integration within the codebase, primarily by implementing a RAGDeps dataclass for explicit, type-safe dependency injection. The LLMGenerationError exception was moved to its own file. Many docstrings were removed from various functions, including Typer commands and internal utility functions, which a reviewer noted degrades the CLI's self-documentation and removes useful explanations of non-obvious behavior. The main.py and mcp/tools.py files were updated to leverage the new RAGDeps structure, passing dependencies explicitly to functions and agents, and directly calling service methods instead of using factory functions to create pydantic_ai.Tool objects. The create_rag_orchestrator function was modified to directly list tool functions, now accepting RAGDeps as a context. Correspondingly, all tool files (code_retrieval.py, codebase_query.py, directory_lister.py, document_analyzer.py, file_editor.py, file_reader.py, file_writer.py, semantic_search.py, shell_command.py) were refactored to accept RunContext[RAGDeps] as their first argument, removing their create_*_tool factory functions and their internal service class instantiations. Integration tests for MCP tools were updated to reflect these changes, directly interacting with service methods and asserting on ingestor.fetch_all calls, though one reviewer suggested strengthening a test to also verify cypher_gen.generate calls.
| cypher_generator: Any | ||
| code_retriever: Any | ||
| file_reader: Any | ||
| file_writer: Any | ||
| file_editor: Any | ||
| shell_commander: Any | ||
| directory_lister: Any | ||
| document_analyzer: Any |
There was a problem hiding this comment.
The RAGDeps dataclass uses Any for most of its fields. This undermines the goal of "Explicit, type-safe dependency injection" mentioned in the PR description. To improve type safety and developer experience, you should use specific class types for these dependencies.
To avoid circular imports that this might cause, you can use a typing.TYPE_CHECKING block and string forward references.
First, add these imports at the top of the file:
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .services.llm import CypherGenerator
from .tools.code_retrieval import CodeRetriever
from .tools.directory_lister import DirectoryLister
from .tools.document_analyzer import DocumentAnalyzer
from .tools.file_editor import FileEditor
from .tools.file_reader import FileReader
from .tools.file_writer import FileWriter
from .tools.shell_command import ShellCommanderThen, update the dataclass fields:
| cypher_generator: Any | |
| code_retriever: Any | |
| file_reader: Any | |
| file_writer: Any | |
| file_editor: Any | |
| shell_commander: Any | |
| directory_lister: Any | |
| document_analyzer: Any | |
| cypher_generator: "CypherGenerator" | |
| code_retriever: "CodeRetriever" | |
| file_reader: "FileReader" | |
| file_writer: "FileWriter" | |
| file_editor: "FileEditor" | |
| shell_commander: "ShellCommander" | |
| directory_lister: "DirectoryLister" | |
| document_analyzer: "DocumentAnalyzer" |
| help="Number of buffered nodes/relationships before flushing to Memgraph", | ||
| ), | ||
| ) -> None: | ||
| """Starts the Codebase RAG CLI.""" |
There was a problem hiding this comment.
The docstring for this Typer command has been removed. Docstrings on Typer command functions are used to generate the help text for the CLI (e.g., graph-code start --help). Removing them degrades the user experience by providing no description for the command. This issue also affects other Typer commands in this file: index, export, optimize, mcp-server, graph-loader-command, and language-command. Please restore the docstrings for all Typer commands to ensure the CLI remains self-documenting.
|
|
||
|
|
||
| def get_multiline_input(prompt_text: str = "Ask a question") -> str: | ||
| """Get multiline input from user with Ctrl+J to submit.""" |
There was a problem hiding this comment.
The docstring for get_multiline_input was useful as it explained the key bindings for submitting (Ctrl+J) versus creating a new line (Enter). This is non-obvious behavior that is important for maintainers to understand. Please consider restoring it.
| await mcp_registry.query_code_graph(query) | ||
|
|
||
| mcp_registry._query_tool.function.assert_called_once_with(query) | ||
| mcp_registry.ingestor.fetch_all.assert_called_once() # type: ignore[attr-defined] |
There was a problem hiding this comment.
This test is weaker than its previous version. It only asserts that ingestor.fetch_all was called, but doesn't verify that the natural language query was passed to the cypher generator. You could strengthen this test by also asserting that cypher_gen.generate was called with the correct query.
| mcp_registry.ingestor.fetch_all.assert_called_once() # type: ignore[attr-defined] | |
| mcp_registry.cypher_gen.generate.assert_called_once_with(query) | |
| mcp_registry.ingestor.fetch_all.assert_called_once() |
Summary
RAGDepsdataclass to centralize all tool dependenciesRunContext[RAGDeps]instead of closure-based factoriescreate_rag_orchestratorto usedeps_typeparameterChanges
codebase_rag/deps.pywithRAGDepsdataclasscodebase_rag/exceptions.pyto break circular importsRunContext[RAGDeps]as first parametermain.pyto create deps at runtime and pass to agentBenefits
Test plan