Spaces:
Sleeping
Sleeping
File size: 4,930 Bytes
944f820 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """
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 |