forked from CodeGraphContext/CodeGraphContext
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.py
More file actions
288 lines (248 loc) · 10.7 KB
/
main.py
File metadata and controls
288 lines (248 loc) · 10.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
# src/codegraphcontext/cli/main.py
"""
This module defines the command-line interface (CLI) for the CodeGraphContext application.
It uses the Typer library to create a user-friendly and well-documented CLI.
Commands:
- setup: Runs an interactive wizard to configure the Neo4j database connection.
- start: Launches the main MCP server.
- tool: A placeholder for directly calling server tools (for debugging).
- help: Displays help information.
- version: Show the installed version.
"""
import typer
from rich.console import Console
from rich.table import Table
from typing import Optional
import asyncio
import logging
import json
import os
import time
from pathlib import Path
from dotenv import load_dotenv, find_dotenv
from importlib.metadata import version as pkg_version, PackageNotFoundError
from codegraphcontext.server import MCPServer
from .setup_wizard import run_setup_wizard
# Set the log level for the noisy neo4j logger to WARNING to keep the output clean.
logging.getLogger("neo4j").setLevel(logging.WARNING)
# Initialize the Typer app and Rich console for formatted output.
app = typer.Typer(
name="cgc",
help="CodeGraphContext: An MCP server for AI-powered code analysis.",
add_completion=True,
)
console = Console(stderr=True)
# Configure basic logging for the application.
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
def get_version() -> str:
"""
Try to read version from the installed package metadata.
Fallback to a dev version if not installed.
"""
try:
return pkg_version("codegraphcontext") # must match [project].name in pyproject.toml
except PackageNotFoundError:
return "0.0.0 (dev)"
@app.command()
def setup():
"""
Runs the interactive setup wizard to configure the server and database connection.
This helps users set up a local Docker-based Neo4j instance or connect to a remote one.
"""
run_setup_wizard()
def _load_credentials():
"""
Loads Neo4j credentials from various sources into environment variables.
Priority order:
1. Local `mcp.json`
2. Global `~/.codegraphcontext/.env`
3. Any `.env` file found in the directory tree.
"""
# 1. Prefer loading from mcp.json
mcp_file_path = Path.cwd() / "mcp.json"
if mcp_file_path.exists():
try:
with open(mcp_file_path, "r") as f:
mcp_config = json.load(f)
server_env = mcp_config.get("mcpServers", {}).get("CodeGraphContext", {}).get("env", {})
for key, value in server_env.items():
os.environ[key] = value
console.print("[green]Loaded Neo4j credentials from local mcp.json.[/green]")
return
except Exception as e:
console.print(f"[bold red]Error loading mcp.json:[/bold red] {e}")
# 2. Try global .env file
global_env_path = Path.home() / ".codegraphcontext" / ".env"
if global_env_path.exists():
try:
load_dotenv(dotenv_path=global_env_path)
console.print(f"[green]Loaded Neo4j credentials from global .env file: {global_env_path}[/green]")
return
except Exception as e:
console.print(f"[bold red]Error loading global .env file from {global_env_path}:[/bold red] {e}")
# 3. Fallback to any discovered .env
try:
dotenv_path = find_dotenv(usecwd=True, raise_error_if_not_found=False)
if dotenv_path:
load_dotenv(dotenv_path)
console.print(f"[green]Loaded Neo4j credentials from discovered .env file: {dotenv_path}[/green]")
else:
console.print("[yellow]No local mcp.json or .env file found. Credentials may not be set.[/yellow]")
except Exception as e:
console.print(f"[bold red]Error loading .env file:[/bold red] {e}")
@app.command()
def start():
"""
Starts the CodeGraphContext MCP server, which listens for JSON-RPC requests from stdin.
"""
console.print("[bold green]Starting CodeGraphContext Server...[/bold green]")
_load_credentials()
server = None
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
# Initialize and run the main server.
server = MCPServer(loop=loop)
loop.run_until_complete(server.run())
except ValueError as e:
# This typically happens if credentials are still not found after all checks.
console.print(f"[bold red]Configuration Error:[/bold red] {e}")
console.print("Please run `cgc setup` to configure the server.")
except KeyboardInterrupt:
# Handle graceful shutdown on Ctrl+C.
console.print("\n[bold yellow]Server stopped by user.[/bold yellow]")
finally:
# Ensure server and event loop are properly closed.
if server:
server.shutdown()
loop.close()
def _run_tool(tool_name: str, tool_args: dict):
"""Helper function to run a tool and handle the server lifecycle."""
_load_credentials()
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
server = MCPServer(loop=loop)
result = loop.run_until_complete(server.handle_tool_call(tool_name, tool_args))
if isinstance(result, dict) and "job_id" in result:
job_id = result["job_id"]
console.print(f"[green]Successfully started job '{job_id}' for tool '{tool_name}'.[/green]")
console.print(f"Estimated files: {result.get('estimated_files')}, Estimated duration: {result.get('estimated_duration_human')}")
console.print(f"\n[bold yellow]Polling for completion...[/bold yellow]")
while True:
time.sleep(2)
status_result = loop.run_until_complete(server.handle_tool_call("check_job_status", {"job_id": job_id}))
job_status = status_result.get("job", {}).get("status")
processed_files = status_result.get("job", {}).get("processed_files", 0)
total_files = status_result.get("job", {}).get("total_files", 0)
console.print(f"Job status: {job_status} ({processed_files}/{total_files} files)")
if job_status in ["completed", "failed", "cancelled"]:
console.print(json.dumps(status_result, indent=2))
break
else:
console.print(json.dumps(result, indent=2))
except ValueError as e:
console.print(f"[bold red]Configuration Error:[/bold red] {e}")
except Exception as e:
console.print(f"[bold red]An unexpected error occurred:[/bold red] {e}")
finally:
if 'loop' in locals() and loop.is_running():
loop.close()
@app.command()
def index(path: Optional[str] = typer.Argument(None, help="Path to the directory or file to index. Defaults to the current directory.")):
"""
Indexes a directory or file by adding it to the code graph.
If no path is provided, it indexes the current directory.
"""
if path is None:
path = "."
_run_tool("add_code_to_graph", {"path": path})
@app.command()
def delete(path: str = typer.Argument(..., help="Path of the repository to delete from the code graph.")):
"""
Deletes a repository from the code graph.
"""
_run_tool("delete_repository", {"repo_path": path})
@app.command()
def visualize(query: Optional[str] = typer.Argument(None, help="The Cypher query to visualize.")):
"""
Generates a URL to visualize a Cypher query in the Neo4j Browser.
If no query is provided, a default query will be used.
"""
if query is None:
query = "MATCH p=()-->() RETURN p"
_run_tool("visualize_graph_query", {"cypher_query": query})
@app.command(name="list_repos")
def list_repos():
"""
Lists all indexed repositories.
"""
_run_tool("list_indexed_repositories", {})
@app.command(name="add_package")
def add_package(package_name: str = typer.Argument(..., help="Name of the package to add."), language: str = typer.Argument(..., help="Language of the package." )):
"""
Adds a package to the code graph.
"""
_run_tool("add_package_to_graph", {"package_name": package_name, "language": language})
@app.command()
def cypher(query: str = typer.Argument(..., help="The read-only Cypher query to execute.")):
"""
Executes a read-only Cypher query.
"""
_run_tool("execute_cypher_query", {"cypher_query": query})
@app.command(name="list_mcp_tools")
def list_mcp_tools():
"""
Lists all available tools and their descriptions.
"""
_load_credentials()
console.print("[bold green]Available Tools:[/bold green]")
try:
# Instantiate the server to access the tool definitions.
server = MCPServer()
tools = server.tools.values()
table = Table(show_header=True, header_style="bold magenta")
table.add_column("Tool Name", style="dim", width=30)
table.add_column("Description")
for tool in sorted(tools, key=lambda t: t['name']):
table.add_row(tool['name'], tool['description'])
console.print(table)
except ValueError as e:
console.print(f"[bold red]Error loading tools:[/bold red] {e}")
console.print("Please ensure your Neo4j credentials are set up correctly (`cgc setup`), as they are needed to initialize the server.")
except Exception as e:
console.print(f"[bold red]An unexpected error occurred:[/bold red] {e}")
@app.command()
def help(ctx: typer.Context):
"""Show the main help message and exit."""
root_ctx = ctx.parent or ctx
typer.echo(root_ctx.get_help())
@app.command("version")
def version_cmd():
"""Show the application version."""
console.print(f"CodeGraphContext [bold cyan]{get_version()}[/bold cyan]")
@app.callback(invoke_without_command=True)
def main(
ctx: typer.Context,
version_: bool = typer.Option(
None,
"--version",
"-v",
help="Show the application version and exit.",
is_eager=True,
),
):
"""
Main entry point for the cgc CLI application.
If no subcommand is provided, it displays a welcome message with instructions.
"""
if version_:
console.print(f"CodeGraphContext [bold cyan]{get_version()}[/bold cyan]")
raise typer.Exit()
if ctx.invoked_subcommand is None:
console.print("[bold green]👋 Welcome to CodeGraphContext (cgc)![/bold green]\n")
console.print("👉 Run [cyan]cgc setup[/cyan] to configure the server and database.")
console.print("👉 Run [cyan]cgc start[/cyan] to launch the server.")
console.print("👉 Run [cyan]cgc help[/cyan] to see all available commands.\n")
console.print("👉 Run [cyan]cgc --version[/cyan] to check the version.\n")
console.print("👉 Running [green]codegraphcontext [white]works the same as using [green]cgc")