File size: 1,632 Bytes
0914e96 |
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 |
import chromadb
from chromadb.config import Settings
import os
# Path to save the database inside ai-service/data
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
DB_PATH = os.path.join(BASE_DIR, "data", "chroma_db")
class VectorStore:
def __init__(self, collection_name="platform_knowledge"):
"""Initialize persistent ChromaDB client."""
# Folder create agar nahi hai to
os.makedirs(DB_PATH, exist_ok=True)
self.client = chromadb.PersistentClient(path=DB_PATH)
# Create or get collection
self.collection = self.client.get_or_create_collection(name=collection_name)
def add_text(self, text_chunks, metadatas, ids):
"""Text data ko DB mein save karna."""
try:
self.collection.upsert(
documents=text_chunks,
metadatas=metadatas,
ids=ids
)
return True
except Exception as e:
print(f"[RAG Error] Failed to add text: {str(e)}")
return False
def search(self, query, n_results=2):
"""Question ke hisaab se matching data lana."""
try:
results = self.collection.query(
query_texts=[query],
n_results=n_results
)
# Thoda sa formatting taaki clean data mile
if results['documents']:
return results['documents'][0] # Return list of matching texts
return []
except Exception as e:
print(f"[RAG Error] Search failed: {str(e)}")
return [] |