Create rag_service.py

rag_service.py

@@ -0,0 +1,146 @@
+import os
+import uuid
+import time
+import shutil
+from base64 import b64decode
+from langchain_community.vectorstores import Chroma
+from langchain.storage import InMemoryStore
+from langchain.schema.document import Document
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain.retrievers.multi_vector import MultiVectorRetriever
+import chromadb
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_groq import ChatGroq
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+
+
+class RAGService:
+    def __init__(self):
+        self.gemini_key = os.getenv("GEMINI_API_KEY")
+        self.groq_key = os.getenv("GROQ_API_KEY")
+
+        # Initialize embeddings
+        self.embeddings = GoogleGenerativeAIEmbeddings(
+            model="models/text-embedding-004",
+            google_api_key=self.gemini_key
+        )
+
+        # Setup ChromaDB
+        self.persist_directory = "/app/chromadb"
+        self.vectorstore = None
+        self.store = None
+        self.retriever = None
+        self.chain_with_sources = None
+
+        self._setup_chromadb()
+        self._setup_retriever()
+        self._setup_chain()
+
+    def _setup_chromadb(self):
+        """Initialize ChromaDB """
+
+
+        self.vectorstore = Chroma(
+            collection_name="multi_modal_rag_new",
+            embedding_function=self.embeddings,
+            persist_directory=self.persist_directory
+        )
+
+        self.store = InMemoryStore()
+
+        print(f"Number of documents in vectorstore: {self.vectorstore._collection.count()}")
+        print("ChromaDB loaded successfully!")
+
+    def _setup_retriever(self):
+        """Setup the MultiVectorRetriever"""
+        self.retriever = MultiVectorRetriever(
+            vectorstore=self.vectorstore,
+            docstore=self.store,
+            id_key="doc_id",
+        )
+
+        # Load data into docstore
+        collection = self.vectorstore._collection
+        all_data = collection.get(include=['metadatas'])
+
+        doc_store_pairs = []
+        for doc_id, metadata in zip(all_data['ids'], all_data['metadatas']):
+            if metadata and 'original_content' in metadata and 'doc_id' in metadata:
+                doc_store_pairs.append((metadata['doc_id'], metadata['original_content']))
+
+        if doc_store_pairs:
+            self.store.mset(doc_store_pairs)
+            print(f"Populated docstore with {len(doc_store_pairs)} documents")
+
+        print(f"Vectorstore count: {self.vectorstore._collection.count()}")
+        print(f"Docstore count: {len(self.store.store)}")
+        print("ChromaDB loaded and ready for querying!")
+
+    def _setup_chain(self):
+        """Setup the RAG chain"""
+        self.chain_with_sources = {
+                                      "context": self.retriever | RunnableLambda(self.parse_docs),
+                                      "question": RunnablePassthrough(),
+                                  } | RunnablePassthrough().assign(
+            response=(
+                    RunnableLambda(self.build_prompt)
+                    | ChatGroq(model="llama-3.1-8b-instant", groq_api_key=self.groq_key)
+                    | StrOutputParser()
+            )
+        )
+
+    def parse_docs(self, docs):
+        """Split base64-encoded images and texts"""
+        b64 = []
+        text = []
+        for doc in docs:
+            try:
+                b64decode(doc)
+                b64.append(doc)
+            except Exception as e:
+                text.append(doc)
+        return {"images": b64, "texts": text}
+
+    def build_prompt(self, kwargs):
+        """Build prompt with context and images"""
+        docs_by_type = kwargs["context"]
+        user_question = kwargs["question"]
+
+        context_text = ""
+        if len(docs_by_type["texts"]) > 0:
+            for text_element in docs_by_type["texts"]:
+                context_text += str(text_element)
+
+        prompt_template = f"""
+        Answer the question based only on the following context, which can include text, tables, and the below image.
+        Context: {context_text}
+        Question: {user_question}
+        """
+
+        prompt_content = [{"type": "text", "text": prompt_template}]
+
+        if len(docs_by_type["images"]) > 0:
+            for image in docs_by_type["images"]:
+                prompt_content.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{image}"},
+                    }
+                )
+
+        return ChatPromptTemplate.from_messages(
+            [
+                HumanMessage(content=prompt_content),
+            ]
+        )
+
+    def ask_question(self, question: str):
+        """Process a question and return response"""
+        response = self.chain_with_sources.invoke(question)
+        return response['response']
+
+
+# Create a global instance
+rag_service = RAGService()