wright 0.1.0__py3-none-any.whl

api/routes/chat.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+
+from fastapi import APIRouter, Depends, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from api.auth import verify_api_key
+from api.quota import check_quota
+from api.rate_limit import limiter
+
+router = APIRouter(prefix="/chat", tags=["chat"], dependencies=[Depends(verify_api_key)])
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+
+class ChatRequest(BaseModel):
+    question: str
+    repo_root: str
+    conversation_history: list[ChatMessage] = []
+
+
+@router.post("")
+@limiter.limit("40/minute")
+async def chat(request: ChatRequest, http_request: Request) -> StreamingResponse:
+    """
+    Handles an incoming chat POST request by retrieving relevant code context via hybrid retrieval and streaming AI-generated responses back to the client as server-sent events.
+
+    This async endpoint initializes the LLM gateway (using Anthropic), a VoyageEmbedder, a ChromaStore vector store, and a HybridRetriever to fetch up to 5 semantically and dependency-relevant code contexts for the user's question. Conversation history is trimmed to the last 20 messages to stay within context limits. The response is streamed as server-sent events, emitting 'token' events for incremental text, 'citations' events for referenced source files, 'followups' events for suggested follow-up questions, and a terminal '[DONE]' sentinel. If embeddings or the vector store are unavailable (e.g., missing API keys or unindexed repository), retrieval is skipped gracefully and the LLM answers from the question alone.
+
+    Args:
+        request (ChatRequest): The chat request object containing the user's question, the repository root path used to locate the Chroma vector store, and the prior conversation history as a list of role/content message pairs.
+
+    Returns:
+        StreamingResponse: A Starlette StreamingResponse with media type 'text/event-stream' that emits JSON-encoded server-sent events of types 'token' (incremental LLM text), 'citations' (referenced source files), and 'followups' (suggested follow-up questions), terminated by a '[DONE]' message.
+
+    Example:
+        ```
+        response = await chat(ChatRequest(question='How does authentication work?', repo_root='/path/to/repo', conversation_history=[]))
+        ```
+    """
+    from api.embedder import get_gateway
+    from api.routes.repos import ensure_repo_local
+    from api.usage_store import record_event
+
+    api_key = http_request.headers.get("X-Wright-API-Key", "")
+
+    # Gate: chat_messages_per_month == 0 on free → 403; >0 enforces monthly cap
+    quota = check_quota(api_key, "chat_message", raise_on_blocked=True)
+    await ensure_repo_local(request.repo_root)
+    gateway = get_gateway()
+    history = [{"role": m.role, "content": m.content} for m in request.conversation_history]
+
+    # Retrieval is best-effort — if embeddings are unavailable (no Voyage/OpenAI
+    # key, or repo not yet indexed) we fall back to empty context and Claude
+    # answers from the question alone.
+    contexts = []
+    try:
+        from api.chroma_cache import get as get_chroma
+        from api.embedder import get_embedder
+        from api.routes.repos import get_vector_store
+        from core.parser.dep_graph import DependencyGraph
+        from core.retrieval.hybrid_retriever import HybridRetriever
+
+        embedder = get_embedder()
+        chroma_path = os.getenv("CHROMA_PATH", os.path.join(request.repo_root, ".wright", "chroma"))
+        chroma = get_vector_store(request.repo_root, get_chroma(chroma_path, request.repo_root))
+        dep_graph = DependencyGraph()
+        retriever = HybridRetriever(chroma, dep_graph, embedder)
+        contexts = retriever.retrieve_for_query(request.question, n=5)
+    except Exception:
+        pass  # no embeddings available — Claude answers without code context
+
+    # Trim history to last 20 messages to stay within context limits
+    if len(history) > 20:
+        history = history[-20:]
+
+    async def event_stream():
+        tokens_used = 0
+        from core.llm.gateway import LLMGateway as _LLMGateway
+
+        model_used = _LLMGateway.PRIMARY_MODEL
+        try:
+            if quota.warning:
+                yield f"data: {json.dumps({'type': 'quota_warning', 'used': quota.used, 'limit': quota.limit, 'pct': quota.pct, 'upgrade_url': quota.upgrade_url})}\n\n"
+            async for kind, payload in gateway.chat_stream(request.question, contexts, history):
+                if kind == "token":
+                    yield f"data: {json.dumps({'type': 'token', 'content': payload})}\n\n"
+                elif kind == "citations":
+                    yield f"data: {json.dumps({'type': 'citations', 'files': payload})}\n\n"
+                elif kind == "followups":
+                    yield f"data: {json.dumps({'type': 'followups', 'questions': payload})}\n\n"
+                elif kind == "model":
+                    model_used = payload
+                elif kind == "usage":
+                    tokens_used = payload
+            yield "data: [DONE]\n\n"
+        finally:
+            # Runs on normal completion AND client disconnect — ensures event is always recorded
+            asyncio.create_task(
+                asyncio.to_thread(
+                    record_event,
+                    api_key,
+                    "chat_message",
+                    tokens=tokens_used,
+                    repo_name=os.path.basename(request.repo_root),
+                    model=model_used,
+                    is_fallback=(model_used != _LLMGateway.PRIMARY_MODEL),
+                    conversation_turns=len(history),
+                    context_chunks=len(contexts),
+                )
+            )
+
+    return StreamingResponse(event_stream(), media_type="text/event-stream")

api/routes/coverage.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import asyncio
+import os
+
+from fastapi import APIRouter, Depends, Query, Request
+from pydantic import BaseModel
+
+from api.auth import verify_api_key
+
+router = APIRouter(prefix="/coverage", tags=["coverage"], dependencies=[Depends(verify_api_key)])
+
+
+class UndocumentedFunction(BaseModel):
+    function_name: str
+    file_path: str
+    line: int
+
+
+class CoverageResponse(BaseModel):
+    overall_pct: float
+    total: int
+    documented: int
+    undocumented: list[UndocumentedFunction]
+    by_file: dict[str, float]
+    by_folder: dict[str, float]
+
+
+@router.get("", response_model=CoverageResponse)
+async def get_coverage(
+    repo_root: str = Query(..., description="Repository root path"),
+    http_request: Request = None,
+) -> CoverageResponse:
+    """
+    Scans all Python files under a repository root and returns documentation coverage statistics at the overall, per-file, and per-folder levels.
+
+    Loads repository configuration via load_config to determine exclude paths, then uses CodeParser to parse every Python file under repo_root. For each parsed file, named functions are counted and checked for existing docstrings to compute total and documented function counts. Undocumented functions are collected with their file path and line number. Coverage percentages are calculated overall, per file, and per folder. If an HTTP request object is provided and contains an X-Wright-API-Key header, the scan event is recorded in the usage store.
+
+    Args:
+        repo_root (str): Absolute or relative path to the repository root directory to scan for Python files and analyze documentation coverage.
+        http_request (Request): The incoming FastAPI/Starlette HTTP request object used to extract the X-Wright-API-Key header for usage tracking. Defaults to None if not provided.
+
+    Returns:
+        CoverageResponse: A CoverageResponse containing: overall_pct (float, overall documentation coverage percentage), total (int, total function count), documented (int, count of documented functions), undocumented (list of UndocumentedFunction with function_name, file_path, and line), by_file (dict mapping file path to coverage percentage), and by_folder (dict mapping folder path to coverage percentage).
+
+    Example:
+        ```
+        coverage = await get_coverage(repo_root="/home/user/projects/my_python_app")
+        ```
+
+    Complexity: O(n*m) time where n is the number of Python files and m is the average number of functions per file; O(n*m) space for storing parsed function data across all files.
+    """
+    from api.routes.repos import ensure_repo_local
+    from core.config import load_config
+    from core.parser.tree_sitter_parser import CodeParser
+
+    await ensure_repo_local(repo_root)
+
+    loop = asyncio.get_event_loop()
+    config = await loop.run_in_executor(None, load_config, repo_root)
+    parser = CodeParser()
+    parsed_files = await loop.run_in_executor(
+        None, parser.parse_directory, repo_root, config.exclude
+    )
+
+    total = 0
+    documented = 0
+    undoc_list: list[UndocumentedFunction] = []
+    by_file: dict[str, float] = {}
+    by_folder: dict[str, tuple[int, int]] = {}
+
+    for pf in parsed_files:
+        funcs = [f for f in pf.functions if f.name != "<anonymous>"]
+        for cls in pf.classes:
+            funcs += [f for f in cls.methods if f.name != "<anonymous>"]
+        file_total = len(funcs)
+        file_doc = sum(1 for f in funcs if f.existing_docstring)
+        total += file_total
+        documented += file_doc
+
+        if file_total > 0:
+            by_file[pf.path] = file_doc / file_total * 100
+
+        folder = os.path.dirname(pf.path)
+        t, d = by_folder.get(folder, (0, 0))
+        by_folder[folder] = (t + file_total, d + file_doc)
+
+        for func in funcs:
+            if not func.existing_docstring:
+                undoc_list.append(
+                    UndocumentedFunction(
+                        function_name=func.name,
+                        file_path=pf.path,
+                        line=func.start_line + 1,
+                    )
+                )
+
+    overall_pct = (documented / total * 100) if total else 100.0
+    folder_pct = {folder: (d / t * 100 if t > 0 else 100.0) for folder, (t, d) in by_folder.items()}
+
+    if http_request:
+        from api.usage_store import record_event
+
+        asyncio.create_task(
+            asyncio.to_thread(
+                record_event,
+                http_request.headers.get("X-Wright-API-Key", ""),
+                "coverage_scans",
+                repo_name=os.path.basename(repo_root),
+            )
+        )
+
+    return CoverageResponse(
+        overall_pct=overall_pct,
+        total=total,
+        documented=documented,
+        undocumented=undoc_list,
+        by_file=by_file,
+        by_folder=folder_pct,
+    )