File size: 1,705 Bytes
abe7564 db774b8 abe7564 db774b8 6d33ce0 | 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 | from typing import Optional
from pinecone import Pinecone, ServerlessSpec
from src.models import Chunk
from src.config import PINECONE_API_KEY, PINECONE_INDEX_NAME
from src.embedder import embedding_dimension
_pc: Optional[Pinecone] = None
_index = None
def _get_pc() -> Pinecone:
global _pc
if _pc is None:
_pc = Pinecone(api_key=PINECONE_API_KEY)
return _pc
def _ensure_index():
global _index
if _index is not None:
return _index
pc = _get_pc()
dim = embedding_dimension()
existing = [i.name for i in pc.list_indexes()]
if PINECONE_INDEX_NAME not in existing:
pc.create_index(
name=PINECONE_INDEX_NAME,
dimension=dim,
metric="cosine",
spec=ServerlessSpec(cloud="aws", region="us-east-1"),
)
_index = pc.Index(PINECONE_INDEX_NAME)
return _index
def upsert_embeddings(chunks: list[Chunk], vectors: list[list[float]]):
index = _ensure_index()
records = []
for chunk, vec in zip(chunks, vectors):
records.append({
"id": chunk.id,
"values": vec,
"metadata": {"text": chunk.text[:1000], "source": chunk.source},
})
batch_size = 100
for i in range(0, len(records), batch_size):
index.upsert(records[i:i + batch_size])
def semantic_search(query_vector: list[float], top_k: int = 5) -> list[dict]:
index = _ensure_index()
result = index.query(vector=query_vector, top_k=top_k, include_metadata=True)
return [{"id": m.id, "score": m.score, "text": m.metadata.get("text", "")}
for m in result.matches]
def close():
global _pc, _index
_index = None
_pc = None
|