codegraphcontext 0.1.0__tar.gz → 0.1.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (28) hide show
  1. {codegraphcontext-0.1.0/src/codegraphcontext.egg-info → codegraphcontext-0.1.2}/PKG-INFO +1 -1
  2. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/pyproject.toml +6 -3
  3. codegraphcontext-0.1.2/src/codegraphcontext/cli/__init__.py +1 -0
  4. codegraphcontext-0.1.2/src/codegraphcontext/cli/main.py +66 -0
  5. codegraphcontext-0.1.2/src/codegraphcontext/cli/setup_wizard.py +313 -0
  6. codegraphcontext-0.1.2/src/codegraphcontext/core/__init__.py +1 -0
  7. codegraphcontext-0.1.2/src/codegraphcontext/core/database.py +84 -0
  8. codegraphcontext-0.1.2/src/codegraphcontext/core/jobs.py +102 -0
  9. codegraphcontext-0.1.2/src/codegraphcontext/core/watcher.py +100 -0
  10. codegraphcontext-0.1.2/src/codegraphcontext/tools/__init__.py +0 -0
  11. codegraphcontext-0.1.2/src/codegraphcontext/tools/code_finder.py +568 -0
  12. codegraphcontext-0.1.2/src/codegraphcontext/tools/graph_builder.py +608 -0
  13. codegraphcontext-0.1.2/src/codegraphcontext/tools/import_extractor.py +118 -0
  14. codegraphcontext-0.1.2/src/codegraphcontext/tools/system.py +131 -0
  15. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2/src/codegraphcontext.egg-info}/PKG-INFO +1 -1
  16. codegraphcontext-0.1.2/src/codegraphcontext.egg-info/SOURCES.txt +25 -0
  17. codegraphcontext-0.1.0/src/codegraphcontext.egg-info/SOURCES.txt +0 -13
  18. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/LICENSE +0 -0
  19. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/README.md +0 -0
  20. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/setup.cfg +0 -0
  21. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext/__init__.py +0 -0
  22. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext/__main__.py +0 -0
  23. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext/prompts.py +0 -0
  24. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext/server.py +0 -0
  25. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext.egg-info/dependency_links.txt +0 -0
  26. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext.egg-info/entry_points.txt +0 -0
  27. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext.egg-info/requires.txt +0 -0
  28. {codegraphcontext-0.1.0 → codegraphcontext-0.1.2}/src/codegraphcontext.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: codegraphcontext
3
- Version: 0.1.0
3
+ Version: 0.1.2
4
4
  Summary: An MCP server that indexes local code into a graph database to provide context to AI assistants.
5
5
  Author-email: Shashank Shekhar Singh <shashankshekharsingh1205@gmail.com>
6
6
  License: MIT License
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "codegraphcontext"
3
- version = "0.1.0"
3
+ version = "0.1.2"
4
4
  description = "An MCP server that indexes local code into a graph database to provide context to AI assistants."
5
5
  authors = [{ name = "Shashank Shekhar Singh", email = "shashankshekharsingh1205@gmail.com" }]
6
6
  readme = "README.md"
@@ -39,5 +39,8 @@ dev = [
39
39
  ]
40
40
 
41
41
  [tool.setuptools]
42
- packages = ["codegraphcontext"]
43
- package-dir = { "" = "src" }
42
+ package-dir = { "" = "src" }
43
+
44
+ [tool.setuptools.packages.find]
45
+ where = ["src"]
46
+ include = ["codegraphcontext*"]
@@ -0,0 +1 @@
1
+ # src/codegraphcontext/cli/__init__.py
@@ -0,0 +1,66 @@
1
+ # src/codegraphcontext/cli/main.py
2
+ import typer
3
+ from rich.console import Console
4
+ import asyncio
5
+ import logging
6
+ from codegraphcontext.server import MCPServer
7
+ from .setup_wizard import run_setup_wizard
8
+
9
+ # Set the log level for the noisy neo4j logger to WARNING
10
+ logging.getLogger("neo4j").setLevel(logging.WARNING) # <-- ADD THIS LINE
11
+
12
+ app = typer.Typer(
13
+ name="cgc",
14
+ help="CodeGraphContext: An MCP server for AI-powered code analysis.",
15
+ add_completion=False,
16
+ )
17
+ console = Console()
18
+
19
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
20
+
21
+ @app.command()
22
+ def setup():
23
+ """
24
+ Run the interactive setup wizard to configure the server and database.
25
+ """
26
+ run_setup_wizard()
27
+
28
+ @app.command()
29
+ def start():
30
+ """
31
+ Start the CodeGraphContext MCP server.
32
+ """
33
+ console.print("[bold green]Starting CodeGraphContext Server...[/bold green]")
34
+ server = None
35
+ loop = asyncio.new_event_loop()
36
+ asyncio.set_event_loop(loop)
37
+ try:
38
+ server = MCPServer(loop=loop)
39
+ loop.run_until_complete(server.run())
40
+ except ValueError as e:
41
+ console.print(f"[bold red]Configuration Error:[/bold red] {e}")
42
+ console.print("Please run `cgc setup` to configure the server.")
43
+ except KeyboardInterrupt:
44
+ console.print("\n[bold yellow]Server stopped by user.[/bold yellow]")
45
+ finally:
46
+ if server:
47
+ server.shutdown()
48
+ loop.close()
49
+
50
+
51
+ @app.command()
52
+ def tool(
53
+ name: str = typer.Argument(..., help="The name of the tool to call."),
54
+ args: str = typer.Argument("{}", help="A JSON string of arguments for the tool."),
55
+ ):
56
+ """
57
+ Directly call a server tool (for debugging and scripting).
58
+ """
59
+ # This is a placeholder for a more advanced tool caller that would
60
+ # connect to the running server via a different mechanism (e.g., a socket).
61
+ # For now, it's a conceptual part of the CLI.
62
+ console.print(f"Calling tool [bold cyan]{name}[/bold cyan] with args: {args}")
63
+ console.print("[yellow]Note: This is a placeholder for direct tool invocation.[/yellow]")
64
+
65
+ if __name__ == "__main__":
66
+ app()
@@ -0,0 +1,313 @@
1
+ from InquirerPy import prompt
2
+ from rich.console import Console
3
+ import subprocess
4
+ import platform
5
+ import os
6
+ from pathlib import Path
7
+ import time
8
+ import json
9
+ import sys
10
+
11
+ console = Console()
12
+
13
+ def _generate_mcp_json(creds):
14
+ """Generates and prints the MCP JSON configuration."""
15
+ cgc_path = sys.executable # Path to the python executable in the current venv
16
+
17
+ mcp_config = {
18
+ "mcpServers": {
19
+ "CodeGraphContext": {
20
+ "command": cgc_path,
21
+ "args": ["start"],
22
+ "env": {
23
+ "NEO4J_URI": creds.get("uri", ""),
24
+ "NEO4J_USER": creds.get("username", "neo4j"),
25
+ "NEO4J_PASSWORD": creds.get("password", "")
26
+ },
27
+ "tools": {
28
+ "alwaysAllow": [
29
+ "list_imports", "add_code_to_graph", "add_package_to_graph",
30
+ "check_job_status", "list_jobs", "find_code",
31
+ "analyze_code_relationships", "watch_directory",
32
+ "find_dead_code", "execute_cypher_query"
33
+ ],
34
+ "disabled": False
35
+ },
36
+ "disabled": False,
37
+ "alwaysAllow": []
38
+ }
39
+ }
40
+ }
41
+
42
+ console.print("\n[bold green]Configuration successful![/bold green]")
43
+ console.print("Copy the following JSON and add it to your MCP server configuration file:")
44
+ console.print(json.dumps(mcp_config, indent=2))
45
+
46
+ # Also save to a file for convenience
47
+ mcp_file = Path.cwd() / "mcp.json"
48
+ with open(mcp_file, "w") as f:
49
+ json.dump(mcp_config, f, indent=2)
50
+ console.print(f"\n[cyan]For your convenience, the configuration has also been saved to: {mcp_file}[/cyan]")
51
+
52
+
53
+ def get_project_root() -> Path:
54
+ """Always return the directory where the user runs `cgc` (CWD)."""
55
+ return Path.cwd()
56
+
57
+ def run_command(command, console, shell=False, check=True, input_text=None):
58
+ """
59
+ Runs a command, captures its output, and handles execution.
60
+ Returns the completed process object on success, None on failure.
61
+ """
62
+ cmd_str = command if isinstance(command, str) else ' '.join(command)
63
+ console.print(f"[cyan]$ {cmd_str}[/cyan]")
64
+ try:
65
+ process = subprocess.run(
66
+ command,
67
+ shell=shell,
68
+ check=check,
69
+ capture_output=True, # Always capture to control what gets displayed
70
+ text=True,
71
+ timeout=300,
72
+ input=input_text
73
+ )
74
+ return process
75
+ except subprocess.CalledProcessError as e:
76
+ console.print(f"[bold red]Error executing command:[/bold red] {cmd_str}")
77
+ if e.stdout:
78
+ console.print(f"[red]STDOUT: {e.stdout}[/red]")
79
+ if e.stderr:
80
+ console.print(f"[red]STDERR: {e.stderr}[/red]")
81
+ return None
82
+ except subprocess.TimeoutExpired:
83
+ console.print(f"[bold red]Command timed out:[/bold red] {cmd_str}")
84
+ return None
85
+
86
+ def run_setup_wizard():
87
+ """Guides the user through setting up CodeGraphContext."""
88
+ console.print("[bold cyan]Welcome to the CodeGraphContext Setup Wizard![/bold cyan]")
89
+
90
+ questions = [
91
+ {
92
+ "type": "list",
93
+ "message": "Where is your Neo4j database located?",
94
+ "choices": [
95
+ "Local (Recommended: I'll help you run it on this machine)",
96
+ "Hosted (Connect to a remote database like AuraDB)",
97
+ ],
98
+ "name": "db_location",
99
+ }
100
+ ]
101
+ result = prompt(questions)
102
+ db_location = result.get("db_location")
103
+
104
+ if db_location and "Hosted" in db_location:
105
+ setup_hosted_db()
106
+ elif db_location:
107
+ setup_local_db()
108
+
109
+
110
+ def find_latest_neo4j_creds_file():
111
+ """Finds the latest Neo4j credentials file in the Downloads folder."""
112
+ downloads_path = Path.home() / "Downloads"
113
+ if not downloads_path.exists():
114
+ return None
115
+
116
+ cred_files = list(downloads_path.glob("Neo4j*.txt"))
117
+ if not cred_files:
118
+ return None
119
+
120
+ latest_file = max(cred_files, key=lambda f: f.stat().st_mtime)
121
+ return latest_file
122
+
123
+
124
+
125
+ def setup_hosted_db():
126
+ """Guides user to configure a remote Neo4j instance."""
127
+ questions = [
128
+ {
129
+ "type": "list",
130
+ "message": "How would you like to add your Neo4j credentials?",
131
+ "choices": ["Add credentials from file", "Add credentials manually"],
132
+ "name": "cred_method",
133
+ }
134
+ ]
135
+ result = prompt(questions)
136
+ cred_method = result.get("cred_method")
137
+
138
+ creds = {}
139
+ if cred_method and "file" in cred_method:
140
+ latest_file = find_latest_neo4j_creds_file()
141
+ file_to_parse = None
142
+ if latest_file:
143
+ confirm_questions = [
144
+ {
145
+ "type": "confirm",
146
+ "message": f"Found a credentials file: {latest_file}. Use this file?",
147
+ "name": "use_latest",
148
+ "default": True,
149
+ }
150
+ ]
151
+ if prompt(confirm_questions).get("use_latest"):
152
+ file_to_parse = latest_file
153
+
154
+ if not file_to_parse:
155
+ path_questions = [
156
+ {"type": "input", "message": "Please enter the path to your credentials file:", "name": "cred_file_path"}
157
+ ]
158
+ file_path_str = prompt(path_questions).get("cred_file_path", "")
159
+ file_path = Path(file_path_str.strip())
160
+ if file_path.exists() and file_path.is_file():
161
+ file_to_parse = file_path
162
+ else:
163
+ console.print("[red]❌ The specified file path does not exist or is not a file.[/red]")
164
+ return
165
+
166
+ if file_to_parse:
167
+ try:
168
+ with open(file_to_parse, "r") as f:
169
+ for line in f:
170
+ if "=" in line:
171
+ key, value = line.strip().split("=", 1)
172
+ if key == "NEO4J_URI":
173
+ creds["uri"] = value
174
+ elif key == "NEO4J_USERNAME":
175
+ creds["username"] = value
176
+ elif key == "NEO4J_PASSWORD":
177
+ creds["password"] = value
178
+ except Exception as e:
179
+ console.print(f"[red]❌ Failed to parse credentials file: {e}[/red]")
180
+ return
181
+
182
+ elif cred_method: # Manual entry
183
+ console.print("Please enter your remote Neo4j connection details.")
184
+ questions = [
185
+ {"type": "input", "message": "URI (e.g., neo4j+s://xxxx.databases.neo4j.io):", "name": "uri"},
186
+ {"type": "input", "message": "Username:", "name": "username", "default": "neo4j"},
187
+ {"type": "password", "message": "Password:", "name": "password"},
188
+ ]
189
+ manual_creds = prompt(questions)
190
+ if not manual_creds: return # User cancelled
191
+ creds = manual_creds
192
+
193
+ if creds.get("uri") and creds.get("password"):
194
+ _generate_mcp_json(creds)
195
+ else:
196
+ console.print("[red]❌ Incomplete credentials. Please try again.[/red]")
197
+
198
+ def setup_local_db():
199
+ """Guides user to set up a local Neo4j instance."""
200
+ questions = [
201
+ {
202
+ "type": "list",
203
+ "message": "How would you like to run Neo4j locally?",
204
+ "choices": ["Docker (Easiest)", "Local Binary (Advanced)"],
205
+ "name": "local_method",
206
+ }
207
+ ]
208
+ result = prompt(questions)
209
+ local_method = result.get("local_method")
210
+
211
+ if local_method and "Docker" in local_method:
212
+ setup_docker()
213
+ elif local_method:
214
+ setup_local_binary()
215
+
216
+ def setup_docker():
217
+ """Creates Docker files and runs docker-compose."""
218
+ console.print("This will create a `docker-compose.yml` and `.env` file in your current directory.")
219
+ # Here you would write the file contents
220
+ console.print("[green]docker-compose.yml and .env created.[/green]")
221
+ console.print("Please set your NEO4J_PASSWORD in the .env file.")
222
+ confirm_q = [{"type": "confirm", "message": "Ready to launch Docker containers?", "name": "proceed"}]
223
+ if prompt(confirm_q).get("proceed"):
224
+ try:
225
+ # Using our run_command to handle the subprocess call
226
+ docker_process = run_command(["docker", "compose", "up", "-d"], console, check=True)
227
+ if docker_process:
228
+ console.print("[bold green]Docker containers started successfully![/bold green]")
229
+ except Exception as e:
230
+ console.print(f"[bold red]Failed to start Docker containers:[/bold red] {e}")
231
+
232
+ def setup_local_binary():
233
+ """Automates the installation and configuration of Neo4j on Ubuntu/Debian."""
234
+ os_name = platform.system()
235
+ console.print(f"Detected Operating System: [bold yellow]{os_name}[/bold yellow]")
236
+
237
+ if os_name != "Linux" or not os.path.exists("/etc/debian_version"):
238
+ console.print("[yellow]Automated installer is designed for Debian-based systems (like Ubuntu).[/yellow]")
239
+ console.print(f"For other systems, please follow the manual installation guide: [bold blue]https://neo4j.com/docs/operations-manual/current/installation/[/bold blue]")
240
+ return
241
+
242
+ console.print("[bold]Starting automated Neo4j installation for Ubuntu/Debian.[/bold]")
243
+ console.print("[yellow]This will run several commands with 'sudo'. You will be prompted for your password.[/yellow]")
244
+ confirm_q = [{"type": "confirm", "message": "Do you want to proceed?", "name": "proceed", "default": True}]
245
+ if not prompt(confirm_q).get("proceed"):
246
+ return
247
+
248
+ NEO4J_VERSION = "1:5.21.0"
249
+
250
+ install_commands = [
251
+ ("Creating keyring directory", ["sudo", "mkdir", "-p", "/etc/apt/keyrings"]),
252
+ ("Adding Neo4j GPG key", "wget -qO- https://debian.neo4j.com/neotechnology.gpg.key | sudo gpg --dearmor --yes -o /etc/apt/keyrings/neotechnology.gpg", True),
253
+ ("Adding Neo4j repository", "echo 'deb [signed-by=/etc/apt/keyrings/neotechnology.gpg] https://debian.neo4j.com stable 5' | sudo tee /etc/apt/sources.list.d/neo4j.list > /dev/null", True),
254
+ ("Updating apt sources", ["sudo", "apt-get", "-qq", "update"]),
255
+ (f"Installing Neo4j ({NEO4J_VERSION}) and Cypher Shell", ["sudo", "apt-get", "install", "-qq", "-y", f"neo4j={NEO4J_VERSION}", "cypher-shell"])
256
+ ]
257
+
258
+ for desc, cmd, use_shell in [(c[0], c[1], c[2] if len(c) > 2 else False) for c in install_commands]:
259
+ console.print(f"\n[bold]Step: {desc}...[/bold]")
260
+ if not run_command(cmd, console, shell=use_shell):
261
+ console.print(f"[bold red]Failed on step: {desc}. Aborting installation.[/bold]")
262
+ return
263
+
264
+ console.print("\n[bold green]Neo4j installed successfully![/bold green]")
265
+
266
+ console.print("\n[bold]Please set the initial password for the 'neo4j' user.""")
267
+
268
+ new_password = ""
269
+ while True:
270
+ questions = [
271
+ {"type": "password", "message": "Enter a new password for Neo4j:", "name": "password"},
272
+ {"type": "password", "message": "Confirm the new password:", "name": "password_confirm"},
273
+ ]
274
+ passwords = prompt(questions)
275
+ if not passwords: return # User cancelled
276
+ new_password = passwords.get("password")
277
+ if new_password and new_password == passwords.get("password_confirm"):
278
+ break
279
+ console.print("[red]Passwords do not match or are empty. Please try again.[/red]")
280
+
281
+ console.print("\n[bold]Stopping Neo4j to set the password...""")
282
+ if not run_command(["sudo", "systemctl", "stop", "neo4j"], console):
283
+ console.print("[bold red]Could not stop Neo4j service. Aborting.[/bold red]")
284
+ return
285
+
286
+ console.print("\n[bold]Setting initial password using neo4j-admin...""")
287
+ pw_command = ["sudo", "-u", "neo4j", "neo4j-admin", "dbms", "set-initial-password", new_password]
288
+ if not run_command(pw_command, console, check=True):
289
+ console.print("[bold red]Failed to set the initial password. Please check the logs.[/bold red]")
290
+ run_command(["sudo", "systemctl", "start", "neo4j"], console)
291
+ return
292
+
293
+ console.print("\n[bold]Starting Neo4j service...""")
294
+ if not run_command(["sudo", "systemctl", "start", "neo4j"], console):
295
+ console.print("[bold red]Failed to start Neo4j service after setting password.[/bold red]")
296
+ return
297
+
298
+ console.print("\n[bold]Enabling Neo4j service to start on boot...""")
299
+ if not run_command(["sudo", "systemctl", "enable", "neo4j"], console):
300
+ console.print("[bold yellow]Could not enable Neo4j service. You may need to start it manually after reboot.[/bold yellow]")
301
+
302
+ console.print("[bold green]Password set and service started.[/bold green]")
303
+
304
+ console.print("\n[yellow]Waiting 10 seconds for the database to become available...""")
305
+ time.sleep(10)
306
+
307
+ creds = {
308
+ "uri": "neo4j://localhost:7687",
309
+ "username": "neo4j",
310
+ "password": new_password
311
+ }
312
+ _generate_mcp_json(creds)
313
+ console.print("\n[bold green]All done! Your local Neo4j instance is ready to use.[/bold green]")
@@ -0,0 +1 @@
1
+ # src/codegraphcontext/core/__init__.py
@@ -0,0 +1,84 @@
1
+ # src/codegraphcontext/core/database.py
2
+ import os
3
+ import logging
4
+ import threading
5
+ from typing import Optional
6
+
7
+ from neo4j import GraphDatabase, Driver
8
+
9
+ logger = logging.getLogger(__name__)
10
+
11
+ class DatabaseManager:
12
+ """
13
+ Singleton class to manage Neo4j database connections.
14
+ """
15
+ _instance = None
16
+ _driver: Optional[Driver] = None
17
+ _lock = threading.Lock()
18
+
19
+ def __new__(cls):
20
+ if cls._instance is None:
21
+ with cls._lock:
22
+ if cls._instance is None:
23
+ cls._instance = super(DatabaseManager, cls).__new__(cls)
24
+ return cls._instance
25
+
26
+ def __init__(self):
27
+ if hasattr(self, '_initialized'):
28
+ return
29
+
30
+ self.neo4j_uri = os.getenv('NEO4J_URI')
31
+ self.neo4j_username = os.getenv('NEO4J_USERNAME', 'neo4j')
32
+ self.neo4j_password = os.getenv('NEO4J_PASSWORD')
33
+ self._initialized = True
34
+
35
+ def get_driver(self) -> Driver:
36
+ """Get the Neo4j driver instance"""
37
+ if self._driver is None:
38
+ with self._lock:
39
+ if self._driver is None:
40
+ if not all([self.neo4j_uri, self.neo4j_username, self.neo4j_password]):
41
+ raise ValueError(
42
+ "Neo4j credentials must be set via environment variables:\n"
43
+ "- NEO4J_URI\n"
44
+ "- NEO4J_USERNAME\n"
45
+ "- NEO4J_PASSWORD"
46
+ )
47
+
48
+ logger.info(f"Creating Neo4j driver connection to {self.neo4j_uri}")
49
+ self._driver = GraphDatabase.driver(
50
+ self.neo4j_uri,
51
+ auth=(self.neo4j_username, self.neo4j_password)
52
+ )
53
+ # Test the connection
54
+ try:
55
+ with self._driver.session() as session:
56
+ session.run("RETURN 1").consume()
57
+ logger.info("Neo4j connection established successfully")
58
+ except Exception as e:
59
+ logger.error(f"Failed to connect to Neo4j: {e}")
60
+ if self._driver:
61
+ self._driver.close()
62
+ self._driver = None
63
+ raise
64
+ return self._driver
65
+
66
+ def close_driver(self):
67
+ """Close the Neo4j driver"""
68
+ if self._driver is not None:
69
+ with self._lock:
70
+ if self._driver is not None:
71
+ logger.info("Closing Neo4j driver")
72
+ self._driver.close()
73
+ self._driver = None
74
+
75
+ def is_connected(self) -> bool:
76
+ """Check if connected to database"""
77
+ if self._driver is None:
78
+ return False
79
+ try:
80
+ with self._driver.session() as session:
81
+ session.run("RETURN 1").consume()
82
+ return True
83
+ except Exception:
84
+ return False
@@ -0,0 +1,102 @@
1
+ # src/codegraphcontext/core/jobs.py
2
+ import uuid
3
+ import threading
4
+ from datetime import datetime, timedelta
5
+ from dataclasses import dataclass, asdict
6
+ from enum import Enum
7
+ from typing import Any, Dict, List, Optional
8
+
9
+ class JobStatus(Enum):
10
+ """Job status enumeration"""
11
+ PENDING = "pending"
12
+ RUNNING = "running"
13
+ COMPLETED = "completed"
14
+ FAILED = "failed"
15
+ CANCELLED = "cancelled"
16
+
17
+ @dataclass
18
+ class JobInfo:
19
+ """Data class for job information"""
20
+ job_id: str
21
+ status: JobStatus
22
+ start_time: datetime
23
+ end_time: Optional[datetime] = None
24
+ total_files: int = 0
25
+ processed_files: int = 0
26
+ current_file: Optional[str] = None
27
+ estimated_duration: Optional[float] = None
28
+ actual_duration: Optional[float] = None
29
+ errors: List[str] = None
30
+ result: Optional[Dict[str, Any]] = None
31
+ path: Optional[str] = None
32
+ is_dependency: bool = False
33
+
34
+ def __post_init__(self):
35
+ if self.errors is None:
36
+ self.errors = []
37
+
38
+ @property
39
+ def progress_percentage(self) -> float:
40
+ """Calculate progress percentage"""
41
+ if self.total_files == 0:
42
+ return 0.0
43
+ return (self.processed_files / self.total_files) * 100
44
+
45
+ @property
46
+ def estimated_time_remaining(self) -> Optional[float]:
47
+ """Calculate estimated time remaining"""
48
+ if self.status != JobStatus.RUNNING or self.processed_files == 0:
49
+ return None
50
+ elapsed = (datetime.now() - self.start_time).total_seconds()
51
+ avg_time_per_file = elapsed / self.processed_files
52
+ remaining_files = self.total_files - self.processed_files
53
+ return remaining_files * avg_time_per_file
54
+
55
+ class JobManager:
56
+ """Manager for background jobs"""
57
+ def __init__(self):
58
+ self.jobs: Dict[str, JobInfo] = {}
59
+ self.lock = threading.Lock()
60
+
61
+ def create_job(self, path: str, is_dependency: bool = False) -> str:
62
+ """Create a new job and return job ID"""
63
+ job_id = str(uuid.uuid4())
64
+ with self.lock:
65
+ self.jobs[job_id] = JobInfo(
66
+ job_id=job_id,
67
+ status=JobStatus.PENDING,
68
+ start_time=datetime.now(),
69
+ path=path,
70
+ is_dependency=is_dependency
71
+ )
72
+ return job_id
73
+
74
+ def update_job(self, job_id: str, **kwargs):
75
+ """Update job information"""
76
+ with self.lock:
77
+ if job_id in self.jobs:
78
+ job = self.jobs[job_id]
79
+ for key, value in kwargs.items():
80
+ if hasattr(job, key):
81
+ setattr(job, key, value)
82
+
83
+ def get_job(self, job_id: str) -> Optional[JobInfo]:
84
+ """Get job information"""
85
+ with self.lock:
86
+ return self.jobs.get(job_id)
87
+
88
+ def list_jobs(self) -> List[JobInfo]:
89
+ """List all jobs"""
90
+ with self.lock:
91
+ return list(self.jobs.values())
92
+
93
+ def cleanup_old_jobs(self, max_age_hours: int = 24):
94
+ """Clean up jobs older than specified hours"""
95
+ cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
96
+ with self.lock:
97
+ jobs_to_remove = [
98
+ job_id for job_id, job in self.jobs.items()
99
+ if job.end_time and job.end_time < cutoff_time
100
+ ]
101
+ for job_id in jobs_to_remove:
102
+ del self.jobs[job_id]