from __future__ import annotations import os import sqlite3 from pathlib import Path from typing import Any WORKSPACE_ROOT = Path( os.getenv( "HF_MCP_AGENT_WORKSPACE", str(Path.home() / "huggingface-mcp-agent-workspace"), ) ).resolve() DB_PATH = WORKSPACE_ROOT / "workspace.sqlite3" SCHEMA = """ CREATE TABLE IF NOT EXISTS projects ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL UNIQUE, kind TEXT NOT NULL DEFAULT 'project', description TEXT NOT NULL DEFAULT '', created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS files ( id INTEGER PRIMARY KEY AUTOINCREMENT, project_name TEXT, path TEXT NOT NULL UNIQUE, kind TEXT NOT NULL, size INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS exports ( id INTEGER PRIMARY KEY AUTOINCREMENT, project_name TEXT, path TEXT NOT NULL, format TEXT NOT NULL, size INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); """ def connect() -> sqlite3.Connection: WORKSPACE_ROOT.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(DB_PATH) conn.row_factory = sqlite3.Row return conn def init_database() -> dict[str, Any]: with connect() as conn: conn.executescript(SCHEMA) return {"initialized": True, "database": str(DB_PATH)} def row_to_dict(row: sqlite3.Row) -> dict[str, Any]: return {key: row[key] for key in row.keys()} def record_project(name: str, kind: str = "project", description: str = "") -> dict[str, Any]: init_database() with connect() as conn: conn.execute( """ INSERT INTO projects (name, kind, description) VALUES (?, ?, ?) ON CONFLICT(name) DO UPDATE SET kind = excluded.kind, description = excluded.description, updated_at = CURRENT_TIMESTAMP """, (name, kind, description), ) row = conn.execute("SELECT * FROM projects WHERE name = ?", (name,)).fetchone() return row_to_dict(row) def record_file(path: str, kind: str, size: int = 0, project_name: str | None = None) -> dict[str, Any]: init_database() project = project_name or infer_project_name(path) with connect() as conn: conn.execute( """ INSERT INTO files (project_name, path, kind, size) VALUES (?, ?, ?, ?) ON CONFLICT(path) DO UPDATE SET project_name = excluded.project_name, kind = excluded.kind, size = excluded.size, updated_at = CURRENT_TIMESTAMP """, (project, path, kind, size), ) row = conn.execute("SELECT * FROM files WHERE path = ?", (path,)).fetchone() return row_to_dict(row) def record_export(path: str, file_format: str, size: int = 0, project_name: str | None = None) -> dict[str, Any]: init_database() project = project_name or infer_project_name(path) with connect() as conn: conn.execute( "INSERT INTO exports (project_name, path, format, size) VALUES (?, ?, ?, ?)", (project, path, file_format, size), ) row = conn.execute("SELECT * FROM exports ORDER BY id DESC LIMIT 1").fetchone() return row_to_dict(row) def list_projects(limit: int = 100) -> list[dict[str, Any]]: init_database() with connect() as conn: rows = conn.execute( "SELECT * FROM projects ORDER BY updated_at DESC LIMIT ?", (max(1, min(int(limit), 500)),), ).fetchall() return [row_to_dict(row) for row in rows] def list_file_records(project_name: str = "", limit: int = 200) -> list[dict[str, Any]]: init_database() with connect() as conn: if project_name: rows = conn.execute( "SELECT * FROM files WHERE project_name = ? ORDER BY updated_at DESC LIMIT ?", (project_name, max(1, min(int(limit), 1000))), ).fetchall() else: rows = conn.execute( "SELECT * FROM files ORDER BY updated_at DESC LIMIT ?", (max(1, min(int(limit), 1000)),), ).fetchall() return [row_to_dict(row) for row in rows] def list_exports(project_name: str = "", limit: int = 100) -> list[dict[str, Any]]: init_database() with connect() as conn: if project_name: rows = conn.execute( "SELECT * FROM exports WHERE project_name = ? ORDER BY created_at DESC LIMIT ?", (project_name, max(1, min(int(limit), 500))), ).fetchall() else: rows = conn.execute( "SELECT * FROM exports ORDER BY created_at DESC LIMIT ?", (max(1, min(int(limit), 500)),), ).fetchall() return [row_to_dict(row) for row in rows] def search_records(query: str, limit: int = 200) -> dict[str, list[dict[str, Any]]]: init_database() pattern = f"%{query}%" capped = max(1, min(int(limit), 500)) with connect() as conn: projects = conn.execute( """ SELECT * FROM projects WHERE name LIKE ? OR description LIKE ? OR kind LIKE ? ORDER BY updated_at DESC LIMIT ? """, (pattern, pattern, pattern, capped), ).fetchall() files = conn.execute( """ SELECT * FROM files WHERE path LIKE ? OR kind LIKE ? OR project_name LIKE ? ORDER BY updated_at DESC LIMIT ? """, (pattern, pattern, pattern, capped), ).fetchall() exports = conn.execute( """ SELECT * FROM exports WHERE path LIKE ? OR format LIKE ? OR project_name LIKE ? ORDER BY created_at DESC LIMIT ? """, (pattern, pattern, pattern, capped), ).fetchall() return { "projects": [row_to_dict(row) for row in projects], "files": [row_to_dict(row) for row in files], "exports": [row_to_dict(row) for row in exports], } def database_summary() -> dict[str, Any]: init_database() with connect() as conn: project_count = conn.execute("SELECT COUNT(*) AS count FROM projects").fetchone()["count"] file_count = conn.execute("SELECT COUNT(*) AS count FROM files").fetchone()["count"] export_count = conn.execute("SELECT COUNT(*) AS count FROM exports").fetchone()["count"] return { "database": str(DB_PATH), "projects": project_count, "files": file_count, "exports": export_count, } def infer_project_name(path: str) -> str | None: first = path.strip("/").split("/", 1)[0] return first or None