File size: 4,604 Bytes
22e29e9 d939c6f 22e29e9 d939c6f 22e29e9 d939c6f 22e29e9 d939c6f 22e29e9 d939c6f 739b192 22e29e9 739b192 d939c6f 22e29e9 739b192 22e29e9 d939c6f 22e29e9 d939c6f 22e29e9 d939c6f 22e29e9 d939c6f | 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 | import uuid
from typing import List, Dict
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams, PointStruct
from sentence_transformers import SentenceTransformer, CrossEncoder
from core import config
class CRMVectorStore:
def __init__(self):
if not config.QDRANT_URL or not config.QDRANT_API_KEY:
raise ValueError("QDRANT_URL and QDRANT_API_KEY must be set in .env")
self.client = QdrantClient(url=config.QDRANT_URL, api_key=config.QDRANT_API_KEY)
self.collection_name = config.QDRANT_COLLECTION_NAME
# --- Safety Check to prevent accidental crossover with GoNidhi ---
if self.collection_name == "cattle_vectors_spatial" or "crm" not in self.collection_name.lower():
raise ValueError(f"CRITICAL ERROR: Refusing to connect to collection '{self.collection_name}'. "
"This collection name belongs to GoNidhi or lacks the 'crm' prefix. "
"Safety bounds triggered to prevent data corruption.")
# -----------------------------------------------------------------
# Initialize local SentenceTransformers embedding model
print(f"Loading local embedding model: {config.RAG_EMBEDDING_MODEL}")
self.encoder = SentenceTransformer(config.RAG_EMBEDDING_MODEL)
print(f"Loading local reranker model: {config.RAG_RERANKER_MODEL}")
self.reranker = CrossEncoder(config.RAG_RERANKER_MODEL)
# bge-base-en-v1.5 has 768 dimensions
self.vector_size = self.encoder.get_embedding_dimension()
self._init_collection()
def _init_collection(self):
try:
if not self.client.collection_exists(self.collection_name):
print(f"Collection '{self.collection_name}' not found. Creating...")
self.client.create_collection(
collection_name=self.collection_name,
vectors_config=VectorParams(size=self.vector_size, distance=Distance.COSINE),
)
else:
print(f"Collection '{self.collection_name}' found and ready.")
except Exception as e:
print(f"Warning during collection init: {e}")
def add_texts(self, texts: List[str], metadata: List[Dict] = None):
"""Encodes texts into vectors and upserts them into Qdrant."""
if not texts:
return
if not metadata:
metadata = [{} for _ in texts]
# Generate embeddings locally
print("Generating embeddings...")
embeddings = self.encoder.encode(texts, show_progress_bar=True)
# Construct points
points = []
for i, (text, vec) in enumerate(zip(texts, embeddings)):
point_id = str(uuid.uuid4())
payload = metadata[i]
payload["text"] = text
points.append(PointStruct(id=point_id, vector=vec.tolist(), payload=payload))
# Upsert in batches to avoid overwhelming the network
batch_size = 100
for i in range(0, len(points), batch_size):
batch = points[i:i + batch_size]
self.client.upsert(
collection_name=self.collection_name,
points=batch
)
print(f"Successfully upserted {len(points)} chunks into Qdrant.")
def search(self, query: str, top_k: int = 5) -> List[Dict]:
"""Searches Qdrant for the top k most similar chunks and reranks them."""
query_vector = self.encoder.encode(query).tolist()
retrieve_k = getattr(config, 'RAG_RETRIEVE_K', 15)
# Stage 1: Dense Retrieval
response = self.client.query_points(
collection_name=self.collection_name,
query=query_vector,
limit=retrieve_k
)
hits = response.points
if not hits:
return []
# Stage 2: Cross-Encoder Reranking
documents = [hit.payload["text"] for hit in hits]
sentence_pairs = [[query, doc] for doc in documents]
scores = self.reranker.predict(sentence_pairs)
results = []
for i, hit in enumerate(hits):
results.append({
"score": float(scores[i]),
"payload": hit.payload
})
# Sort by reranker score descending
results = sorted(results, key=lambda x: x["score"], reverse=True)
return results[:top_k]
|