store/vector_store.py

"""
vector_store.py
---------------
Pinecone-backed vector store interface for the Codebase Oracle system.
Reads from the same Pinecone index used by embed.py via namespaces.

Collections (as Pinecone namespaces):
    - class_chunks    : one chunk per class (macro / cross-module queries)
    - function_chunks : one chunk per function/method (micro queries)

Depends on:
    - pinecone
    - ingest.embed (get_pinecone_index)
    - rich
"""

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 get_pinecone_index, CLASS_COLLECTION, FUNCTION_COLLECTION

console = Console()


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

@dataclass
class ChunkResult:
    """Represents a single retrieved chunk."""
    id:       str
    text:     str
    metadata: dict
    distance: float | None = None

    @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 chunk_type(self) -> str:
        return self.metadata.get("type", "unknown")

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

    @property
    def relevance(self) -> float:
        if self.distance is None:
            return 0.0
        return round(1 / (1 + self.distance), 4)


# ── VectorStore ───────────────────────────────────────────────────────────────

class VectorStore:
    """
    Pinecone-backed interface for stats and tree queries.
    Reuses the same index as embed.py — no duplicate client.
    """

    def __init__(self):
        self._index = get_pinecone_index()
        console.print("[green]✔[/green] VectorStore ready (Pinecone)\n")

    def _count(self, namespace: str) -> int:
        """Return approximate vector count in a namespace."""
        stats = self._index.describe_index_stats()
        return stats["namespaces"].get(namespace, {}).get("vector_count", 0)

    def stats(self) -> dict:
        class_count = self._count(CLASS_COLLECTION)
        func_count  = self._count(FUNCTION_COLLECTION)
        return {
            "class_chunks":    class_count,
            "function_chunks": func_count,
            "total":           class_count + func_count,
        }

    def is_indexed(self) -> bool:
        s = self.stats()
        return s["total"] > 0

    def get_all(self, namespace: str, limit: int = 10) -> list[ChunkResult]:
        """
        Fetch chunks from a namespace without a query vector.
        Pinecone does not support scan — we use a zero vector as proxy.
        """
        from config.config import EMBEDDING_DIM
        zero_vector = [0.0] * EMBEDDING_DIM

        results = self._index.query(
            vector=zero_vector,
            top_k=limit,
            namespace=namespace,
            include_metadata=True,
        )

        output = []
        for match in results["matches"]:
            meta = dict(match["metadata"])
            text = meta.pop("text", "")
            output.append(ChunkResult(
                id=match["id"],
                text=text,
                metadata=meta,
                distance=1 - match["score"],
            ))
        return output

    def render_stats(self) -> None:
        s = self.stats()
        table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
        table.add_column(style="dim")
        table.add_column(style="bold white")
        table.add_row("Class chunks",    str(s["class_chunks"]))
        table.add_row("Function chunks", str(s["function_chunks"]))
        table.add_row("Total chunks",    str(s["total"]))
        table.add_row(
            "Status",
            "[bold green]✔ Indexed[/bold green]"
            if self.is_indexed()
            else "[bold red]✘ Not indexed[/bold red]"
        )
        console.print(Panel(
            table,
            title="[bold cyan]VectorStore Stats[/bold cyan]",
            border_style="cyan",
        ))


# ── Singleton ─────────────────────────────────────────────────────────────────

_store_instance: VectorStore | None = None


def get_vector_store() -> VectorStore:
    global _store_instance
    if _store_instance is None:
        _store_instance = VectorStore()
    return _store_instance