coderag_api.py

#!/usr/bin/env python
"""Local HTTP API for the CodeRAG React demo."""

import argparse
import json
import mimetypes
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlparse

from coderag import build_file_index
from coderag.vector_db import CodeRagVectorDB


ROOT = Path(__file__).resolve().parent
MAX_REQUEST_BYTES = 2 * 1024 * 1024
STATIC_ROOT = ROOT / "coderag_ui" / "dist"
DEFAULT_SAMPLE_CODE = """from collections import defaultdict


def normalize_user(user):
    email = user["email"].strip().lower()
    return {
        "id": user["id"],
        "email": email,
        "active": bool(user.get("active", True)),
    }


def group_active_users(users):
    groups = defaultdict(list)
    for user in users:
        normalized = normalize_user(user)
        if normalized["active"]:
            domain = normalized["email"].split("@")[-1]
            groups[domain].append(normalized)
    return dict(groups)


def find_user_by_email(users, email):
    target = email.strip().lower()
    for user in users:
        normalized = normalize_user(user)
        if normalized["email"] == target:
            return normalized
    return None
"""


class CodeRagApiHandler(BaseHTTPRequestHandler):
    server_version = "CodeRAGApi/1.0"

    def log_message(self, format_string, *args):
        print(f"{self.address_string()} - {format_string % args}")

    @property
    def vector_db(self) -> CodeRagVectorDB:
        return self.server.vector_db

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.end_headers()
        self.wfile.write(body)

    def do_OPTIONS(self):
        self._send_json(204, {})

    def do_GET(self):
        path = urlparse(self.path).path
        if path == "/api/health":
            self._send_json(200, {"status": "ok"})
            return
        if path == "/api/coderag/stats":
            try:
                self._send_json(200, self.vector_db.stats())
            except Exception as error:
                self._send_json(
                    500,
                    {"error": f"Stats failed: {type(error).__name__}: {error}"},
                )
            return
        self._serve_static(path)

    def do_POST(self):
        path = urlparse(self.path).path
        if path not in {
            "/api/coderag/search",
            "/api/coderag/rebuild",
            "/api/coderag/index-code",
        }:
            self._send_json(404, {"error": "Not found"})
            return

        try:
            payload = self._read_payload()
        except ValueError as error:
            self._send_json(400, {"error": str(error)})
            return

        if path == "/api/coderag/rebuild":
            try:
                if self.server.repository_index.exists():
                    result = self.vector_db.build(
                        repository_index_path=self.server.repository_index,
                        synthetic_questions_path=self.server.synthetic_questions,
                    )
                else:
                    sample_dir = ROOT / ".code_query_repo_cache" / "sample"
                    sample_dir.mkdir(parents=True, exist_ok=True)
                    sample_path = sample_dir / "sample_users.py"
                    sample_path.write_text(DEFAULT_SAMPLE_CODE, encoding="utf-8")
                    result = self.vector_db.build_from_file_index(
                        build_file_index(sample_path)
                    )
            except Exception as error:
                self._send_json(
                    500,
                    {
                        "error": f"Rebuild failed: "
                        f"{type(error).__name__}: {error}"
                    },
                )
                return
            self._send_json(200, result)
            return

        if path == "/api/coderag/index-code":
            code = payload.get("code")
            filename = payload.get("filename", "pasted_code.py")
            if not isinstance(code, str) or not code.strip():
                self._send_json(400, {"error": "'code' must be a non-empty string"})
                return
            if not isinstance(filename, str) or not filename.strip():
                filename = "pasted_code.py"
            if not filename.endswith(".py"):
                filename += ".py"
            safe_name = "".join(
                char if char.isalnum() or char in {"_", "-", "."} else "_"
                for char in filename
            )
            pasted_dir = ROOT / ".code_query_repo_cache" / "pasted"
            pasted_dir.mkdir(parents=True, exist_ok=True)
            pasted_path = pasted_dir / safe_name
            try:
                pasted_path.write_text(code, encoding="utf-8")
                index = build_file_index(
                    pasted_path,
                    propagation_depth=int(payload.get("propagation_depth", 2)),
                    max_callees=int(payload.get("max_callees", 4)),
                    callee_context_tokens=int(
                        payload.get("callee_context_tokens", 192)
                    ),
                )
                result = self.vector_db.build_from_file_index(index)
            except SyntaxError as error:
                self._send_json(
                    400,
                    {
                        "error": (
                            f"Python parse failed at line {error.lineno}: "
                            f"{error.msg}"
                        )
                    },
                )
                return
            except Exception as error:
                self._send_json(
                    500,
                    {
                        "error": f"Indexing failed: "
                        f"{type(error).__name__}: {error}"
                    },
                )
                return
            self._send_json(200, result)
            return

        query = payload.get("query")
        if not isinstance(query, str) or not query.strip():
            self._send_json(400, {"error": "'query' must be a non-empty string"})
            return
        try:
            result = self.vector_db.search(
                query,
                top_k=int(payload.get("top_k", 8)),
                per_source=int(payload.get("per_source", 5)),
                sources=payload.get("sources") or None,
            )
        except Exception as error:
            self._send_json(
                500,
                {"error": f"Search failed: {type(error).__name__}: {error}"},
            )
            return
        self._send_json(200, result)

    def _read_payload(self) -> dict:
        try:
            content_length = int(self.headers.get("Content-Length", "0"))
        except ValueError as error:
            raise ValueError("Invalid Content-Length") from error
        if content_length < 0 or content_length > MAX_REQUEST_BYTES:
            raise ValueError("Request body is too large")
        if content_length == 0:
            return {}
        try:
            return json.loads(self.rfile.read(content_length))
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise ValueError("Request body must be valid JSON") from error

    def _serve_static(self, request_path: str) -> None:
        if not STATIC_ROOT.exists():
            self._send_json(404, {"error": "UI build not found. Run npm run build."})
            return
        relative = request_path.lstrip("/") or "index.html"
        candidate = (STATIC_ROOT / relative).resolve()
        try:
            candidate.relative_to(STATIC_ROOT.resolve())
        except ValueError:
            self._send_json(403, {"error": "Forbidden"})
            return
        if not candidate.exists() or candidate.is_dir():
            candidate = STATIC_ROOT / "index.html"
        content_type = mimetypes.guess_type(candidate.name)[0] or "application/octet-stream"
        body = candidate.read_bytes()
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)


class CodeRagApiServer(ThreadingHTTPServer):
    def __init__(
        self,
        server_address,
        handler_class,
        *,
        vector_db: CodeRagVectorDB,
        repository_index: Path,
        synthetic_questions: Path,
    ):
        super().__init__(server_address, handler_class)
        self.vector_db = vector_db
        self.repository_index = repository_index
        self.synthetic_questions = synthetic_questions


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run the CodeRAG API server.")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8010")))
    parser.add_argument(
        "--database",
        type=Path,
        default=ROOT / ".code_query_repo_cache" / "coderag_vectors.sqlite",
    )
    parser.add_argument(
        "--repository-index",
        type=Path,
        default=ROOT / "coderag_repository_index.json",
    )
    parser.add_argument(
        "--synthetic-questions",
        type=Path,
        default=ROOT / "coderag_synthetic_queries.json",
    )
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Rebuild the SQLite vector database before serving.",
    )
    return parser


def main() -> int:
    args = build_parser().parse_args()
    vector_db = CodeRagVectorDB(args.database)
    vector_db.initialize()
    if args.rebuild or not args.database.exists():
        stats = vector_db.build(
            repository_index_path=args.repository_index,
            synthetic_questions_path=args.synthetic_questions,
        )
        print(f"Built vector database: {json.dumps(stats, indent=2)}", flush=True)
    server = CodeRagApiServer(
        (args.host, args.port),
        CodeRagApiHandler,
        vector_db=vector_db,
        repository_index=args.repository_index,
        synthetic_questions=args.synthetic_questions,
    )
    url = f"http://{args.host}:{args.port}"
    print(f"CodeRAG API running at {url}", flush=True)
    print("Press Ctrl+C to stop.", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopping CodeRAG API.")
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())