retrieve/retrieve.py

"""
retriever.py
------------
Retrieves relevant chunks from ChromaDB based on query type.

Query type determines which collection to search:
    - Macro / Cross-Module  → class_chunks
    - Micro                 → function_chunks

Provides filtered retrieval by module, class, or function name
for precise context fetching.

Depends on:
    - embedder.py  (query_chunks, get_chroma_client, load_embedding_model)
    - rich         (terminal output)
"""

from dataclasses import dataclass
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich import box

from ingest.embed import (
    query_chunks,
    load_embedding_model,
    CLASS_COLLECTION,
    FUNCTION_COLLECTION,
    COMPLETE_COLLECTION
)

console = Console()


# ── Query Types ───────────────────────────────────────────────────────────────

class QueryType:
    MACRO        = "macro"
    MICRO        = "micro"
    CROSS_MODULE = "cross_module"


# ── Result Model ──────────────────────────────────────────────────────────────

@dataclass
class RetrievedChunk:
    text:       str
    metadata:   dict
    distance:   float
    collection: str

    @property
    def name(self) -> str:
        return self.metadata.get("name", "unknown")

    @property
    def module(self) -> str:
        return self.metadata.get("module", "unknown")

    @property
    def file(self) -> str:
        return self.metadata.get("file", "unknown")

    @property
    def class_name(self) -> str:
        return self.metadata.get("class_name", "")

    @property
    def chunk_type(self) -> str:
        return self.metadata.get("type", "unknown")

    @property
    def relevance_score(self) -> float:
        """Convert distance to 0-1 relevance score (lower distance = higher relevance)."""
        return round(1 / (1 + self.distance), 4)


# ── Core Retriever ────────────────────────────────────────────────────────────

class Retriever:
    """
    Unified retriever for the Codebase Oracle system.
    Maintains a single embedding model and ChromaDB client across queries.
    """

    def __init__(self):
        self.model = load_embedding_model()
        console.print("[green]✔[/green] Retriever ready.\n")

    # ── Public API ────────────────────────────────────────────────────────────

    def retrieve(self,
                 query:      str,
                 query_type: str,
                 n_results:  int = 5,
                 filters:    dict | None = None) -> list[RetrievedChunk]:
        """
        Main retrieval entry point. Routes to the correct collection
        based on query_type.

        Args:
            query:      Natural language query string.
            query_type: One of QueryType.MACRO / MICRO / CROSS_MODULE.
            n_results:  Number of chunks to retrieve.
            filters:    Optional metadata filters (e.g. filter by module).

        Returns:
            List of RetrievedChunk objects sorted by relevance.
        """
        collection = self._route_collection(query_type)

        raw = query_chunks(
            query=query,
            collection_name=collection,
            model=self.model,
            n_results=n_results,
            filters=filters,
        )

        chunks = [
            RetrievedChunk(
                text=r["text"],
                metadata=r["metadata"],
                distance=r["distance"],
                collection=collection,
            )
            for r in raw
        ]

        return chunks

    def retrieve_by_class(self,
                          class_name: str,
                          n_results:  int = 1) -> list[RetrievedChunk]:
        """
        Retrieve class-level chunk by exact class name.
        Used by micro agent to ground function queries with class context.

        Args:
            class_name: Exact class name string.
            n_results:  Usually 1 — we want the specific class.

        Returns:
            List of RetrievedChunk from class_chunks collection.
        """
        return self.retrieve(
            query=f"class {class_name}",
            query_type=QueryType.MACRO,
            n_results=n_results,
            filters={"name": {"$eq": class_name}},
        )

    def retrieve_by_function(self,
                             function_name: str,
                             class_name:    str | None = None,
                             n_results:     int = 3) -> list[RetrievedChunk]:
        """
        Retrieve function-level chunks by function name.
        Optionally filter by class name for method disambiguation.

        Args:
            function_name: Exact function/method name.
            class_name:    Optional class name to narrow results.
            n_results:     Number of results.

        Returns:
            List of RetrievedChunk from function_chunks collection.
        """
        filters = {"name": {"$eq": function_name}}
        if class_name:
            filters = {
                "$and": [
                    {"name":       {"$eq": function_name}},
                    {"class_name": {"$eq": class_name}},
                ]
            }

        return self.retrieve(
            query=f"function {function_name}",
            query_type=QueryType.MICRO,
            n_results=n_results,
            filters=filters,
        )

    def retrieve_by_module(self,
                           module_name: str,
                           query:       str,
                           query_type:  str = QueryType.MACRO,
                           n_results:   int = 5) -> list[RetrievedChunk]:
        """
        Retrieve chunks scoped to a specific module.

        Args:
            module_name: Top-level module/package name.
            query:       Natural language query within that module.
            query_type:  MACRO or MICRO.
            n_results:   Number of results.

        Returns:
            List of RetrievedChunk filtered to the given module.
        """
        return self.retrieve(
            query=query,
            query_type=query_type,
            n_results=n_results,
            filters={"module": {"$eq": module_name}},
        )

    def retrieve_dependencies(self, function_name: str) -> list[RetrievedChunk]:
        """
        Retrieve all functions that the given function calls.
        Used by cross-module agent for dependency/impact analysis.

        Args:
            function_name: Name of the function to trace dependencies for.

        Returns:
            List of RetrievedChunk for each called function found in index.
        """
        # First get the function itself to extract its call list
        source_chunks = self.retrieve_by_function(function_name, n_results=1)

        if not source_chunks:
            console.print(f"[yellow]⚠ Function '{function_name}' not found in index.[/yellow]")
            return []

        import json
        calls = json.loads(source_chunks[0].metadata.get("calls", "[]"))

        if not calls:
            return []

        # Retrieve each called function from the index
        dep_chunks = []
        seen = set()
        for call in calls:
            # Strip object prefix if present (e.g. "self.calculate" → "calculate")
            name = call.split(".")[-1]
            if name in seen:
                continue
            seen.add(name)

            results = self.retrieve_by_function(name, n_results=1)
            dep_chunks.extend(results)

        return dep_chunks

    def build_context(self, chunks: list[RetrievedChunk],
                      max_chars: int = 6000) -> str:
        """
        Concatenate retrieved chunks into a single context string
        for the LLM prompt. Truncates at max_chars to stay within
        context window limits.

        Args:
            chunks:    Retrieved chunks to concatenate.
            max_chars: Maximum total character length of context.

        Returns:
            A single string ready to inject into an LLM prompt.
        """
        parts = []
        total = 0

        for i, chunk in enumerate(chunks, 1):
            section = (
                f"--- Chunk {i} ({chunk.chunk_type}: {chunk.name}) ---\n"
                f"{chunk.text}\n"
            )
            if total + len(section) > max_chars:
                parts.append("\n[Context truncated — token limit reached]")
                break
            parts.append(section)
            total += len(section)

        return "\n".join(parts)

    def retrieve_complete_function(self,
                                   function_name: str,
                                   class_name: str | None = None,
                                   n_results: int = 1) -> list[RetrievedChunk]:
        """
        Retrieve complete function chunk including full source code.
        Used for micro queries requiring edge case or usage analysis.
        """
        filters = {
            "$and": [
                {"type": {"$eq": "complete_function"}},
                {"name": {"$eq": function_name}},
            ]
        }
        if class_name:
            filters = {
                "$and": [
                    {"type": {"$eq": "complete_function"}},
                    {"name": {"$eq": function_name}},
                    {"class_name": {"$eq": class_name}},
                ]
            }

        raw = query_chunks(
            query=f"function {function_name}",
            collection_name=COMPLETE_COLLECTION,
            model=self.model,
            n_results=n_results,
            filters=filters,
        )

        return [
            RetrievedChunk(
                text=r["text"],
                metadata=r["metadata"],
                distance=r["distance"],
                collection=COMPLETE_COLLECTION,
            )
            for r in raw
        ]

    def retrieve_complete_class(self,
                                class_name: str,
                                n_results: int = 1) -> list[RetrievedChunk]:
        """
        Retrieve complete class chunk including full source code.
        Used for class-level deep queries.
        """
        filters = {
            "$and": [
                {"type": {"$eq": "complete_class"}},
                {"name": {"$eq": class_name}},
            ]
        }

        raw = query_chunks(
            query=f"class {class_name}",
            collection_name=COMPLETE_COLLECTION,
            model=self.model,
            n_results=n_results,
            filters=filters,
        )

        return [
            RetrievedChunk(
                text=r["text"],
                metadata=r["metadata"],
                distance=r["distance"],
                collection=COMPLETE_COLLECTION,
            )
            for r in raw
        ]

    def retrieve_file(self,
                      file_name: str,
                      n_results: int = 1) -> list[RetrievedChunk]:
        """
        Retrieve complete file chunk.
        Used for file-wide queries.
        """
        filters = {
            "$and": [
                {"type": {"$eq": "file"}},
                {"name": {"$eq": file_name}},
            ]
        }

        raw = query_chunks(
            query=f"file {file_name}",
            collection_name=COMPLETE_COLLECTION,
            model=self.model,
            n_results=n_results,
            filters=filters,
        )

        return [
            RetrievedChunk(
                text=r["text"],
                metadata=r["metadata"],
                distance=r["distance"],
                collection=COMPLETE_COLLECTION,
            )
            for r in raw
        ]
    # ── Internal ──────────────────────────────────────────────────────────────

    def _route_collection(self, query_type: str) -> str:
        """Map query type to the appropriate ChromaDB collection."""
        if query_type == QueryType.MICRO:
            return FUNCTION_COLLECTION
        return CLASS_COLLECTION   # MACRO and CROSS_MODULE both use class chunks


# ── Rich Renderers ────────────────────────────────────────────────────────────

def render_chunks(chunks: list[RetrievedChunk], title: str = "Retrieved Chunks") -> None:
    """Render retrieved chunks as a rich table."""
    if not chunks:
        console.print("[yellow]No chunks retrieved.[/yellow]")
        return

    table = Table(
        "Rank", "Type", "Name", "Class", "Module", "File", "Relevance",
        box=box.ROUNDED,
        header_style="bold bright_cyan",
        border_style="cyan",
        show_lines=True,
        title=title,
    )

    for i, chunk in enumerate(chunks, 1):
        score = chunk.relevance_score
        score_style = (
            "bold green" if score > 0.7
            else "yellow" if score > 0.4
            else "red"
        )
        table.add_row(
            str(i),
            chunk.chunk_type,
            f"[bold]{chunk.name}[/bold]",
            chunk.class_name or "[dim]—[/dim]",
            chunk.module,
            chunk.file,
            f"[{score_style}]{score}[/{score_style}]",
        )

    console.print(table)


def render_chunk_detail(chunk: RetrievedChunk) -> None:
    """Render the full text of a single chunk in a panel."""
    header = Text()
    header.append(chunk.chunk_type.upper(), style="bold cyan")
    header.append(f"  {chunk.name}", style="bold white")
    if chunk.class_name:
        header.append(f"  in {chunk.class_name}", style="dim")

    console.print(Panel(
        chunk.text,
        title=str(header),
        border_style="cyan",
        padding=(1, 2),
    ))


def render_context(context: str) -> None:
    """Render the assembled LLM context string."""
    console.print(Panel(
        context,
        title="[bold cyan]Assembled LLM Context[/bold cyan]",
        border_style="dim cyan",
        padding=(1, 2),
    ))


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

if __name__ == "__main__":
    import sys

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

    retriever = Retriever()

    # Test 1 — Macro query
    console.rule("[cyan]Test 1 — Macro Query[/cyan]")
    query = "what classes handle file parsing and walking"
    chunks = retriever.retrieve(query, QueryType.MACRO, n_results=3)
    render_chunks(chunks, title=f'Macro: "{query}"')

    # Test 2 — Micro query
    console.rule("[cyan]Test 2 — Micro Query[/cyan]")
    query = "function that parses a single file and extracts classes"
    chunks = retriever.retrieve(query, QueryType.MICRO, n_results=3)
    render_chunks(chunks, title=f'Micro: "{query}"')
    if chunks:
        render_chunk_detail(chunks[0])

    # Test 3 — Cross module dependencies
    console.rule("[cyan]Test 3 — Dependency Retrieval[/cyan]")
    func_name = sys.argv[1] if len(sys.argv) > 1 else "parse_file"
    dep_chunks = retriever.retrieve_dependencies(func_name)
    render_chunks(dep_chunks, title=f'Dependencies of: {func_name}')

    # Test 4 — Build context string
    console.rule("[cyan]Test 4 — Context Assembly[/cyan]")
    all_chunks = retriever.retrieve("class structure and responsibilities",
                                    QueryType.MACRO, n_results=4)
    context = retriever.build_context(all_chunks)
    console.print(f"[green]✔[/green] Context assembled: [cyan]{len(context)}[/cyan] chars")