diff --git a/.gitignore b/.gitignore index 7431db0f4..eb1dc6b65 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,8 @@ wheels/ .pytest_cache/ .claude/ .qdrant_code_embeddings/ + +# Popular VibeCoding Agents +.roo/ +.augment +.vscode diff --git a/README.md b/README.md index bf8a1c628..0f171df23 100644 --- a/README.md +++ b/README.md @@ -255,12 +255,23 @@ The system automatically detects and processes files for all supported languages ### Step 2: Query the Codebase +**Interactive mode:** + Start the interactive RAG CLI: ```bash python -m codebase_rag.main start --repo-path /path/to/your/repo ``` +**Non-interactive mode (single query):** + +Run a single query and exit, with output sent to stdout (useful for scripting): + +```bash +python -m codebase_rag.main start --repo-path /path/to/your/repo \ + --ask-agent "What functions call UserService.create_user?" +``` + ### Step 2.5: Real-Time Graph Updates (Optional) For active development, you can keep your knowledge graph automatically synchronized with code changes using the realtime updater. This is particularly useful when you're actively modifying code and want the AI assistant to always work with the latest codebase structure. @@ -503,6 +514,10 @@ claude mcp add --transport stdio graph-code \ - **surgical_replace_code** - Precise code edits - **read_file / write_file** - File operations - **list_directory** - Browse project structure +- **shell_command** - Execute terminal commands +- **document_analyzer** - Analyze PDFs and documents +- **semantic_search** - Find functions by intent +- **get_function_source** - Retrieve function source by ID ### Example Usage @@ -607,6 +622,9 @@ The agent has access to a suite of tools to understand and interact with the cod - **`create_new_file`**: Creates a new file with specified content. - **`replace_code_surgically`**: Surgically replaces specific code blocks in files. Requires exact target code and replacement. Only modifies the specified block, leaving rest of file unchanged. True surgical patching. - **`execute_shell_command`**: Executes a shell command in the project's environment. +- **`analyze_document`**: Analyzes PDFs and documents to answer questions about their content. +- **`semantic_search_functions`**: Finds functions by natural language intent using embeddings. +- **`get_function_source_by_id`**: Retrieves source code for a function by its node ID. ### Intelligent and Safe File Editing diff --git a/codebase_rag/main.py b/codebase_rag/main.py index 0586a9f3c..5e57ad3d3 100644 --- a/codebase_rag/main.py +++ b/codebase_rag/main.py @@ -159,11 +159,21 @@ def is_edit_operation_response(response_text: str) -> bool: return tool_usage or content_indicators or pattern_match -def _setup_common_initialization(repo_path: str) -> Path: - """Common setup logic for both main and optimize functions.""" +def _setup_common_initialization(repo_path: str, question_mode: bool = False) -> Path: + """Common setup logic for both main and optimize functions. + + Args: + repo_path: Path to the repository + question_mode: If True, suppress INFO/DEBUG/WARNING logs (only show errors and direct output) + """ # Logger initialization logger.remove() - logger.add(sys.stdout, format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {message}") + if question_mode: + # In question mode, only show ERROR level logs + logger.add(sys.stderr, level="ERROR", format="{message}") + else: + # In interactive mode, show all logs + logger.add(sys.stdout, format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {message}") # Temporary directory cleanup project_root = Path(repo_path).resolve() @@ -774,6 +784,29 @@ def _validate_provider_config(role: str, config: Any) -> None: return rag_agent +async def main_async_single_query( + repo_path: str, batch_size: int, question: str +) -> None: + """Initializes services and runs a single query in non-interactive mode.""" + project_root = _setup_common_initialization(repo_path, question_mode=True) + + with MemgraphIngestor( + host=settings.MEMGRAPH_HOST, + port=settings.MEMGRAPH_PORT, + batch_size=batch_size, + ) as ingestor: + rag_agent = _initialize_services_and_agent(repo_path, ingestor) + + # Handle images in the question + question_with_context = _handle_chat_images(question, project_root) + + # Run the query + response = await rag_agent.run(question_with_context, message_history=[]) + + # Output response to stdout + print(response.output) + + async def main_async(repo_path: str, batch_size: int) -> None: """Initializes services and runs the main application loop.""" project_root = _setup_common_initialization(repo_path) @@ -840,6 +873,12 @@ def start( min=1, help="Number of buffered nodes/relationships before flushing to Memgraph", ), + ask_agent: str | None = typer.Option( + None, + "-a", + "--ask-agent", + help="Run a single query and exit (non-interactive mode). Output is sent to stdout.", + ), ) -> None: """Starts the Codebase RAG CLI.""" global confirm_edits_globally @@ -892,7 +931,16 @@ def start( return try: - asyncio.run(main_async(target_repo_path, effective_batch_size)) + if ask_agent: + # Non-interactive mode: run single query and exit + asyncio.run( + main_async_single_query( + target_repo_path, effective_batch_size, ask_agent + ) + ) + else: + # Interactive mode: run chat loop + asyncio.run(main_async(target_repo_path, effective_batch_size)) except KeyboardInterrupt: console.print("\n[bold red]Application terminated by user.[/bold red]") except ValueError as e: diff --git a/codebase_rag/mcp/client.py b/codebase_rag/mcp/client.py new file mode 100644 index 000000000..69b1fc176 --- /dev/null +++ b/codebase_rag/mcp/client.py @@ -0,0 +1,87 @@ +"""MCP client for querying the code graph via the MCP server. + +This module provides a simple CLI client that connects to the MCP server +and executes the ask_agent tool with a provided question. +""" + +import asyncio +import json +import os +import sys +from typing import Any + +import typer +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +app = typer.Typer() + + +async def query_mcp_server(question: str) -> dict[str, Any]: + """Query the MCP server with a question. + + Args: + question: The question to ask about the codebase + + Returns: + Dictionary with the response from the server + """ + # Start the MCP server as a subprocess with stderr redirected to /dev/null + # This suppresses all server logs while keeping stdout/stdin for MCP communication + with open(os.devnull, "w") as devnull: + server_params = StdioServerParameters( + command="python", + args=["-m", "codebase_rag.main", "mcp-server"], + ) + + async with stdio_client(server=server_params, errlog=devnull) as (read, write): + async with ClientSession(read, write) as session: + # Initialize the session + await session.initialize() + + # Call the ask_agent tool + result = await session.call_tool("ask_agent", {"question": question}) + + # Extract the response text + if result.content: + response_text = result.content[0].text + # Parse JSON response + try: + parsed = json.loads(response_text) + if isinstance(parsed, dict): + return parsed + return {"output": str(parsed)} + except json.JSONDecodeError: + return {"output": response_text} + return {"output": "No response from server"} + + +@app.command() +def main( + question: str = typer.Option( + ..., "--ask-agent", "-a", help="Question to ask about the codebase" + ), +) -> None: + """Query the code graph via MCP server. + + Example: + python -m codebase_rag.mcp.client --ask-agent "What functions call UserService.create_user?" + """ + try: + # Run the async query + result = asyncio.run(query_mcp_server(question)) + + # Print only the output (clean for scripting) + if isinstance(result, dict) and "output" in result: + print(result["output"]) + else: + print(json.dumps(result)) + + except Exception as e: + # Print error to stderr and exit with error code + print(f"Error: {str(e)}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + app() diff --git a/codebase_rag/mcp/server.py b/codebase_rag/mcp/server.py index 9e05e36e0..f82380eb5 100644 --- a/codebase_rag/mcp/server.py +++ b/codebase_rag/mcp/server.py @@ -6,7 +6,6 @@ import json import os -import sys from pathlib import Path from loguru import logger @@ -20,14 +19,54 @@ from codebase_rag.services.llm import CypherGenerator -def setup_logging() -> None: - """Configure logging to stderr for MCP stdio transport.""" +def setup_logging(enable_logging: bool = False) -> None: + """Configure logging for MCP stdio transport. + + By default, logging is disabled to prevent token waste in LLM context. + Can be enabled via environment variable MCP_ENABLE_LOGGING=1 for debugging. + + When enabled, logs are written to a file to avoid polluting STDIO transport. + The log file path can be configured via MCP_LOG_FILE environment variable. + + Args: + enable_logging: Whether to enable logging output. Defaults to False. + Can also be controlled via MCP_ENABLE_LOGGING environment variable. + """ logger.remove() # Remove default handler - logger.add( - sys.stderr, - level="INFO", - format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {message}", + + # Check environment variable to override enable_logging parameter + env_enable = os.environ.get("MCP_ENABLE_LOGGING", "").lower() in ( + "1", + "true", + "yes", ) + should_enable = enable_logging or env_enable + + if should_enable: + # Get log file path from environment or use default + log_file = os.environ.get("MCP_LOG_FILE") + if not log_file: + # Use ~/.cache/code-graph-rag/mcp.log as default + cache_dir = Path.home() / ".cache" / "code-graph-rag" + cache_dir.mkdir(parents=True, exist_ok=True) + log_file = str(cache_dir / "mcp.log") + + # Ensure log file directory exists + log_path = Path(log_file) + log_path.parent.mkdir(parents=True, exist_ok=True) + + # Add file handler - logs go to file, not STDERR/STDOUT + logger.add( + log_file, + level="INFO", + format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {message}", + colorize=False, # Disable ANSI color codes + rotation="10 MB", # Rotate when file reaches 10MB + retention="7 days", # Keep logs for 7 days + ) + else: + # Disable all logging by default for MCP mode + logger.disable("codebase_rag") def get_project_root() -> Path: @@ -143,34 +182,45 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]: Tool handlers are dynamically resolved from the MCPToolsRegistry, ensuring consistency with tool definitions. + + Logging is suppressed during tool execution to prevent token waste in LLM context. """ - logger.info(f"[GraphCode MCP] Calling tool: {name}") + import io + from contextlib import redirect_stderr, redirect_stdout try: # Resolve handler from registry handler_info = tools.get_tool_handler(name) if not handler_info: - error_msg = f"Unknown tool: {name}" - logger.error(f"[GraphCode MCP] {error_msg}") + error_msg = "Unknown tool" return [TextContent(type="text", text=f"Error: {error_msg}")] handler, returns_json = handler_info - # Call handler with unpacked arguments - result = await handler(**arguments) - - # Format result based on output type - if returns_json: - result_text = json.dumps(result, indent=2) - else: - result_text = str(result) - - return [TextContent(type="text", text=result_text)] - - except Exception as e: - error_msg = f"Error executing tool '{name}': {str(e)}" - logger.error(f"[GraphCode MCP] {error_msg}", exc_info=True) - return [TextContent(type="text", text=f"Error: {error_msg}")] + # Suppress all logging output during tool execution + with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()): + logger.disable("codebase_rag") + try: + # Call handler with unpacked arguments + result = await handler(**arguments) + + # Format result based on output type + if returns_json: + result_text = json.dumps(result, indent=2) + else: + result_text = str(result) + + return [TextContent(type="text", text=result_text)] + finally: + logger.enable("codebase_rag") + + except Exception: + # Fail silently without logging or printing error details + return [ + TextContent( + type="text", text="Error: There was an error executing the tool" + ) + ] return server, ingestor diff --git a/codebase_rag/mcp/tools.py b/codebase_rag/mcp/tools.py index 483f725e6..f70d495fd 100644 --- a/codebase_rag/mcp/tools.py +++ b/codebase_rag/mcp/tools.py @@ -4,26 +4,37 @@ """ import itertools +import sys from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Any, cast from loguru import logger +from rich.console import Console from codebase_rag.graph_updater import GraphUpdater from codebase_rag.parser_loader import load_parsers from codebase_rag.services.graph_service import MemgraphIngestor -from codebase_rag.services.llm import CypherGenerator +from codebase_rag.services.llm import CypherGenerator, create_rag_orchestrator from codebase_rag.tools.code_retrieval import CodeRetriever, create_code_retrieval_tool from codebase_rag.tools.codebase_query import create_query_tool from codebase_rag.tools.directory_lister import ( DirectoryLister, create_directory_lister_tool, ) +from codebase_rag.tools.document_analyzer import ( + DocumentAnalyzer, + create_document_analyzer_tool, +) from codebase_rag.tools.file_editor import FileEditor, create_file_editor_tool from codebase_rag.tools.file_reader import FileReader, create_file_reader_tool from codebase_rag.tools.file_writer import FileWriter, create_file_writer_tool +from codebase_rag.tools.semantic_search import ( + create_get_function_source_tool, + create_semantic_search_tool, +) +from codebase_rag.tools.shell_command import ShellCommander, create_shell_command_tool @dataclass @@ -66,10 +77,14 @@ def __init__( self.file_reader = FileReader(project_root=project_root) self.file_writer = FileWriter(project_root=project_root) self.directory_lister = DirectoryLister(project_root=project_root) + self.shell_commander = ShellCommander(project_root=project_root) + self.document_analyzer = DocumentAnalyzer(project_root=project_root) # Create pydantic-ai tools - we'll call the underlying functions directly + # Use a Console that outputs to stderr to avoid corrupting JSONRPC on stdout + stderr_console = Console(file=sys.stderr, width=None, force_terminal=True) self._query_tool = create_query_tool( - ingestor=ingestor, cypher_gen=cypher_gen, console=None + ingestor=ingestor, cypher_gen=cypher_gen, console=stderr_console ) self._code_tool = create_code_retrieval_tool(code_retriever=self.code_retriever) self._file_editor_tool = create_file_editor_tool(file_editor=self.file_editor) @@ -78,6 +93,17 @@ def __init__( self._directory_lister_tool = create_directory_lister_tool( directory_lister=self.directory_lister ) + self._shell_command_tool = create_shell_command_tool( + shell_commander=self.shell_commander + ) + self._document_analyzer_tool = create_document_analyzer_tool( + self.document_analyzer + ) + self._semantic_search_tool = create_semantic_search_tool() + self._function_source_tool = create_get_function_source_tool() + + # Create RAG orchestrator agent (lazy initialization for testing) + self._rag_agent: Any = None # Build tool registry - single source of truth for all tool metadata self._tools: dict[str, ToolMetadata] = { @@ -214,8 +240,57 @@ def __init__( handler=self.list_directory, returns_json=False, ), + "ask_agent": ToolMetadata( + name="ask_agent", + description="Ask the Code Graph RAG agent a question about the codebase. " + "Use this tool for general questions about the codebase, architecture, functionality, and code relationships. " + "Examples: 'How is the authentication implemented?', " + "'What are the main components of the system?', 'Where is the database connection configured?'", + input_schema={ + "type": "object", + "properties": { + "question": { + "type": "string", + "description": "A question about the codebase, architecture, functionality, and code relationships. " + "Examples: 'What functions call UserService.create_user?', " + "'How is error handling implemented?', 'What are the main entry points?'", + } + }, + "required": ["question"], + }, + handler=self.ask_agent, + returns_json=True, + ), } + @property + def rag_agent(self) -> Any: + """Lazy-initialize the RAG orchestrator agent on first access. + + This allows tests to mock the agent without triggering LLM initialization. + """ + if self._rag_agent is None: + self._rag_agent = create_rag_orchestrator( + tools=[ + self._query_tool, + self._code_tool, + self._file_reader_tool, + self._file_writer_tool, + self._file_editor_tool, + self._shell_command_tool, + self._directory_lister_tool, + self._document_analyzer_tool, + self._semantic_search_tool, + self._function_source_tool, + ] + ) + return self._rag_agent + + @rag_agent.setter + def rag_agent(self, value: Any) -> None: + """Allow setting the RAG agent (useful for testing).""" + self._rag_agent = value + async def index_repository(self) -> str: """Parse and ingest the repository into the Memgraph knowledge graph. @@ -439,6 +514,43 @@ async def list_directory(self, directory_path: str = ".") -> str: logger.error(f"[MCP] Error listing directory: {e}") return f"Error: {str(e)}" + async def ask_agent(self, question: str) -> dict[str, Any]: + """Ask a single question about the codebase and get an answer. + + This tool executes the question using the RAG agent and returns the response + in a structured format suitable for MCP clients. + + Logging is suppressed during execution to prevent token waste in LLM context. + + Args: + question: The question to ask about the codebase + + Returns: + Dictionary with 'output' key containing the answer + """ + import io + from contextlib import redirect_stderr, redirect_stdout + + # Suppress all logging output during agent execution + try: + # Temporarily redirect stdout and stderr to suppress all output + with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()): + # Temporarily disable loguru logging + logger.disable("codebase_rag") + try: + # Run the query using the RAG agent + response = await self.rag_agent.run(question, message_history=[]) + return {"output": response.output} + finally: + # Re-enable logging + logger.enable("codebase_rag") + except Exception: + # Fail silently without logging or printing error details + return { + "output": "There was an error processing your question", + "error": True, + } + def get_tool_schemas(self) -> list[dict[str, Any]]: """Get MCP tool schemas for all registered tools. diff --git a/codebase_rag/services/llm.py b/codebase_rag/services/llm.py index 204b40b2b..74bbebf92 100644 --- a/codebase_rag/services/llm.py +++ b/codebase_rag/services/llm.py @@ -111,6 +111,7 @@ def create_rag_orchestrator(tools: list[Tool]) -> Agent: model=llm, system_prompt=RAG_ORCHESTRATOR_SYSTEM_PROMPT, tools=tools, + retries=3, # Increase retries to handle output validation issues ) except Exception as e: raise LLMGenerationError(f"Failed to initialize RAG Orchestrator: {e}") from e