inference/inference.py

"""
inference.py
------------
Pure LLM inference layer for the Codebase Oracle system.
No HTTP server, no FastAPI — just context + query → response.

Responsibilities:
    1. Accept query type + subtype + function/module name
    2. Retrieve relevant context from vector store + call graph
    3. Build correct prompt via prompts/ modules
    4. Call OpenRouter API
    5. Return clean markdown response string

This module is the single entry point that main.py (FastAPI) calls.

Depends on:
    - retriever.py
    - call_graph.py
    - macro_prompts.py
    - micro_prompts.py
    - cross_module_prompts.py
    - openai (OpenRouter is OpenAI-compatible)
    - rich (logging only)

Install:
    pip install openai rich
"""

import os
from dataclasses import dataclass, field
from dotenv import load_dotenv
from openai import OpenAI
from rich.console import Console

load_dotenv()  # loads OPENROUTER_API_KEY from .env file

from retrieve.retrieve import Retriever, QueryType
from store.call_graph import get_call_graph
from prompt.macro_prompt import (
    get_macro_system_prompt,
    build_macro_user_prompt,
)
from prompt.micro_prompt import (
    get_micro_system_prompt,
    build_micro_user_prompt,
    build_micro_followup_prompt,
)
from prompt.cross_module import (
    get_cross_module_system_prompt,
    build_cross_module_user_prompt,
    build_cross_module_followup_prompt,
    build_call_graph_payload,
)

console = Console()

# ── Constants ─────────────────────────────────────────────────────────────────

OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_MODEL       = "nvidia/nemotron-3-super-120b-a12b:free"
MAX_TOKENS          = 4096
CONTEXT_MAX_CHARS   = 6000    # max chars fed to LLM as retrieved context
N_RESULTS_DEFAULT   = 5       # default number of chunks to retrieve


# ── Request / Response Models ─────────────────────────────────────────────────

@dataclass
class InferenceRequest:
    """
    Unified request model for all query types.

    Fields:
        query_type    : "macro" | "micro" | "cross_module"
        query         : Developer's natural language question
        subtype       : macro subtype or "" for micro/cross_module
        function_name : Target function/method name (micro + cross_module)
        class_name    : Target class name if method (optional)
        module_name   : Target module name (macro module_responsibility)
        followup      : True if this is a follow-up to a previous response
        previous_response: Previous LLM response for follow-up context
    """
    query_type:        str
    query:             str
    subtype:           str        = ""
    function_name:     str        = ""
    class_name:        str        = ""
    module_name:       str        = ""
    followup:          bool       = False
    previous_response: str        = ""


@dataclass
class InferenceResponse:
    """
    Unified response model returned to FastAPI / caller.

    Fields:
        success  : True if inference succeeded
        content  : Markdown response string
        error    : Error message if success is False
        metadata : Optional dict with retrieval stats
    """
    success:  bool
    content:  str        = ""
    error:    str        = ""
    metadata: dict       = field(default_factory=dict)


# ── OpenRouter Client ─────────────────────────────────────────────────────────

def _get_client() -> OpenAI:
    """
    Build and return an OpenAI-compatible client pointed at OpenRouter.
    Reads OPENROUTER_API_KEY from .env file via load_dotenv.

    Raises:
        EnvironmentError: If OPENROUTER_API_KEY is not set.
    """
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "OPENROUTER_API_KEY environment variable is not set. "
            "Add OPENROUTER_API_KEY=your_key to your .env file."
        )
    return OpenAI(
        base_url=OPENROUTER_BASE_URL,
        api_key=api_key,
    )


def _call_llm(system_prompt: str,
              user_prompt:   str,
              model:         str = DEFAULT_MODEL) -> str:
    """
    Make a single LLM call via OpenRouter and return the response text.

    Args:
        system_prompt : System prompt string.
        user_prompt   : User prompt string.
        model         : OpenRouter model identifier.

    Returns:
        Raw response text from LLM.

    Raises:
        Exception: On API errors, propagated to caller.
    """
    client = _get_client()

    response = client.chat.completions.create(
        model=model,
        max_tokens=MAX_TOKENS,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": user_prompt},
        ],
    )

    return response.choices[0].message.content or ""


# ── Inference Engine ──────────────────────────────────────────────────────────

class InferenceEngine:
    """
    Core inference engine for the Codebase Oracle.
    Retrieves context, builds prompts, calls LLM, returns response.

    Single instance shared across all FastAPI requests.
    """

    def __init__(self, model: str = DEFAULT_MODEL):
        self.model     = model
        self.retriever = Retriever()
        console.print(
            f"[green]✔[/green] InferenceEngine ready "
            f"(model: [cyan]{model}[/cyan])\n"
        )

    # ── Public Entry Point ────────────────────────────────────────────────────

    def infer(self, request: InferenceRequest) -> InferenceResponse:
        """
        Main entry point. Routes to the correct handler based on query_type.

        Args:
            request: InferenceRequest with all query parameters.

        Returns:
            InferenceResponse with markdown content or error.
        """
        try:
            if request.query_type == QueryType.MACRO:
                return self._handle_macro(request)

            elif request.query_type == QueryType.MICRO:
                return self._handle_micro(request)

            elif request.query_type == QueryType.CROSS_MODULE:
                return self._handle_cross_module(request)

            else:
                return InferenceResponse(
                    success=False,
                    error=f"Unknown query_type: '{request.query_type}'. "
                          f"Use 'macro', 'micro', or 'cross_module'."
                )

        except EnvironmentError as e:
            return InferenceResponse(success=False, error=str(e))

        except Exception as e:
            console.print(f"[red]❌ Inference error: {e}[/red]")
            return InferenceResponse(
                success=False,
                error=f"Inference failed: {str(e)}"
            )

    # ── Macro Handler ─────────────────────────────────────────────────────────

    def _handle_macro(self, req: InferenceRequest) -> InferenceResponse:
        """Handle macro-level queries (architecture, module, data flow)."""

        if not req.subtype:
            return InferenceResponse(
                success=False,
                error="Macro queries require a subtype: "
                      "'overall_architecture', 'module_responsibility', "
                      "or 'data_flow'."
            )

        # Retrieve class-level chunks — macro uses class collection
        chunks = self.retriever.retrieve(
            query=req.query,
            query_type=QueryType.MACRO,
            n_results=N_RESULTS_DEFAULT,
            filters={"module": {"$eq": req.module_name}}
            if req.module_name and req.subtype == "module_responsibility"
            else None,
        )

        context  = self.retriever.build_context(chunks, CONTEXT_MAX_CHARS)
        system   = get_macro_system_prompt(req.subtype)
        user     = build_macro_user_prompt(
            subtype=req.subtype,
            context=context,
            query=req.query,
            module_name=req.module_name,
        )

        console.print(
            f"[cyan]→ Macro [{req.subtype}][/cyan] "
            f"| {len(chunks)} chunks | {len(context)} chars"
        )

        content = _call_llm(system, user, self.model)

        return InferenceResponse(
            success=True,
            content=content,
            metadata={
                "query_type":    "macro",
                "subtype":       req.subtype,
                "chunks_used":   len(chunks),
                "context_chars": len(context),
            }
        )

    # ── Micro Handler ─────────────────────────────────────────────────────────

    def _handle_micro(self, req: InferenceRequest) -> InferenceResponse:
        """Handle micro-level queries (function definition, body, consideration)."""

        if not req.function_name:
            return InferenceResponse(
                success=False,
                error="Micro queries require a function_name."
            )

        # Follow-up path
        if req.followup and req.previous_response:
            complete_chunks = self.retriever.retrieve_complete_function(
                req.function_name, req.class_name, n_results=1
            )
            context = self.retriever.build_context(complete_chunks, CONTEXT_MAX_CHARS)
            system = get_micro_system_prompt(followup=True)
            user = build_micro_followup_prompt(
                previous_response=req.previous_response,
                followup_query=req.query,
                context=context,
            )
        else:
            # Primary path — retrieve complete function + complete class context
            complete_func_chunks = self.retriever.retrieve_complete_function(
                req.function_name, req.class_name, n_results=1
            )
            complete_class_chunks = (
                self.retriever.retrieve_complete_class(req.class_name, n_results=1)
                if req.class_name else []
            )
            all_chunks = complete_func_chunks + complete_class_chunks
            context = self.retriever.build_context(all_chunks, CONTEXT_MAX_CHARS)
            system = get_micro_system_prompt(followup=False)
            user = build_micro_user_prompt(
                context=context,
                query=req.query,
                function_name=req.function_name,
                class_name=req.class_name,
            )

        console.print(
            f"[cyan]→ Micro[/cyan] "
            f"| fn: {req.function_name} "
            f"| {len(context)} chars"
        )

        content = _call_llm(system, user, self.model)

        return InferenceResponse(
            success=True,
            content=content,
            metadata={
                "query_type": "micro",
                "function_name": req.function_name,
                "class_name": req.class_name,
                "context_chars": len(context),
                "followup": req.followup,
            }
        )

    # ── Cross-Module Handler ──────────────────────────────────────────────────

    def _handle_cross_module(self, req: InferenceRequest) -> InferenceResponse:
        """Handle cross-module queries (dependency + impact analysis)."""

        if not req.function_name:
            return InferenceResponse(
                success=False,
                error="Cross-module queries require a function_name."
            )

        # Primary context — the target function
        primary_chunks = self.retriever.retrieve_by_function(
            req.function_name, req.class_name, n_results=2
        )
        primary_context = self.retriever.build_context(
            primary_chunks, CONTEXT_MAX_CHARS // 2
        )

        # Dependency context — callers and callees
        dep_chunks = self.retriever.retrieve_dependencies(req.function_name)
        dep_context = self.retriever.build_context(
            dep_chunks, CONTEXT_MAX_CHARS // 2
        )

        # Call graph data
        try:
            graph    = get_call_graph()
            calls    = graph.get_calls(req.function_name, req.class_name)
            called_by = graph.get_called_by(req.function_name, req.class_name)
            impact   = graph.get_impact(req.function_name, req.class_name, depth=2)
        except FileNotFoundError:
            calls, called_by, impact = [], [], {}
            console.print(
                "[yellow]⚠ call_graph.json not found — "
                "proceeding without graph data.[/yellow]"
            )

        call_graph_payload = build_call_graph_payload(
            function_name=req.function_name,
            calls=calls,
            called_by=called_by,
            impact=impact,
        )

        # Follow-up path
        if req.followup and req.previous_response:
            system = get_cross_module_system_prompt(followup=True)
            user   = build_cross_module_followup_prompt(
                previous_response=req.previous_response,
                followup_query=req.query,
                primary_context=primary_context,
                dependency_context=dep_context,
                call_graph_data=call_graph_payload,
            )
        else:
            system = get_cross_module_system_prompt(followup=False)
            user   = build_cross_module_user_prompt(
                primary_context=primary_context,
                dependency_context=dep_context,
                call_graph_data=call_graph_payload,
                query=req.query,
                function_name=req.function_name,
                class_name=req.class_name,
            )

        console.print(
            f"[cyan]→ Cross-Module[/cyan] "
            f"| fn: {req.function_name} "
            f"| deps: {len(dep_chunks)} "
            f"| graph nodes: {len(calls) + len(called_by)}"
        )

        content = _call_llm(system, user, self.model)

        return InferenceResponse(
            success=True,
            content=content,
            metadata={
                "query_type":    "cross_module",
                "function_name": req.function_name,
                "calls":         len(calls),
                "called_by":     len(called_by),
                "dep_chunks":    len(dep_chunks),
                "followup":      req.followup,
            }
        )


# ── Singleton Access ──────────────────────────────────────────────────────────

_engine_instance: InferenceEngine | None = None


def get_engine(model: str = DEFAULT_MODEL) -> InferenceEngine:
    """
    Return a singleton InferenceEngine instance.
    Creates it on first call, reuses on subsequent calls.

    Args:
        model: OpenRouter model identifier.

    Returns:
        Shared InferenceEngine instance.
    """
    global _engine_instance
    if _engine_instance is None:
        _engine_instance = InferenceEngine(model)
    return _engine_instance


# ── Entry Point (manual test) ─────────────────────────────────────────────────

if __name__ == "__main__":
    import sys
    from rich.markdown import Markdown

    console.rule("[bold cyan]Inference Engine — Manual Test[/bold cyan]")

    engine = get_engine()

    # Test micro query
    console.rule("[cyan]Test — Micro Query[/cyan]")
    req = InferenceRequest(
        query_type="micro",
        query="What does this function do and how do I use it?",
        function_name=sys.argv[1] if len(sys.argv) > 1 else "parse_file",
        class_name="",
    )

    resp = engine.infer(req)

    if resp.success:
        console.print(Markdown(resp.content))
        console.print(f"\n[dim]Metadata: {resp.metadata}[/dim]")
    else:
        console.print(f"[red]❌ Error: {resp.error}[/red]")