Update rag_pipeline.py

rag_pipeline.py

@@ -1,57 +1,56 @@
 from sentence_transformers import SentenceTransformer
-import chromadb
-from chromadb.config import Settings
+from chromadb import Client
+from chromadb.utils import embedding_functions
 import uuid
 
 class RAGPipeline:
     """
-    Retrieval-Augmented Generation (RAG) pipeline:
-    - Stores documents in a vector database (ChromaDB)
-    - Generates embeddings using SentenceTransformer
-    - Retrieves relevant documents for queries
+    RAG pipeline using ChromaDB new client API.
     """
 
     def __init__(self, db_dir: str = "./chroma_store"):
-        # Initialize the embedding model
+        # Embedding model
         self.embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-        # Initialize ChromaDB client
-        self.client = chromadb.Client(
-            Settings(chroma_db_impl="duckdb+parquet", persist_directory=db_dir)
-        )
-
-        # Create or get a collection for storing documents
-        self.collection = self.client.get_or_create_collection(name="rag_collection")
+        # ChromaDB new client
+        self.client = Client()  # No Settings() needed in v0.5+
+        self.collection_name = "rag_collection"
+
+        # Check if collection exists; create if not
+        if self.collection_name in [c.name for c in self.client.list_collections()]:
+            self.collection = self.client.get_collection(self.collection_name)
+        else:
+            self.collection = self.client.create_collection(
+                name=self.collection_name,
+                embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                    model_name="sentence-transformers/all-MiniLM-L6-v2"
+                )
+            )
 
     def add_document(self, text: str, doc_id: str = None):
-        """
-        Adds a single document to the vector database.
-        """
         if not doc_id:
             doc_id = str(uuid.uuid4())
-
-        embedding = self.embedder.encode(text).tolist()
-        self.collection.add(documents=[text], ids=[doc_id], embeddings=[embedding])
+        self.collection.add(
+            documents=[text],
+            ids=[doc_id],
+        )
 
     def retrieve(self, query: str, top_k: int = 3):
-        """
-        Retrieves the top_k most relevant documents for a query.
-        """
-        q_embedding = self.embedder.encode(query).tolist()
         results = self.collection.query(
-            query_embeddings=[q_embedding],
+            query_texts=[query],
             n_results=top_k
         )
-
-        # Return list of documents
         return results.get("documents", [[]])[0]
 
     def reset_vector_store(self):
-        """
-        Clears all documents from the collection.
-        """
+        # Delete and recreate the collection
         try:
-            self.client.delete_collection("rag_collection")
+            self.client.delete_collection(self.collection_name)
         except Exception:
             pass
-        self.collection = self.client.get_or_create_collection(name="rag_collection")
+        self.collection = self.client.create_collection(
+            name=self.collection_name,
+            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="sentence-transformers/all-MiniLM-L6-v2"
+            )
+        )