Update app.py

app.py

@@ -37,6 +37,7 @@ from typing import List, Annotated, Any
 import re, operator
 
 
+
 class MultiAgentState(BaseModel):
     state: List[str] = []
     messages: Annotated[list[AnyMessage], add_messages]
@@ -54,22 +55,38 @@ class StoryState(BaseModel):
     stories_lst: Annotated[list, operator.add]
 
 class DocumentRAG:
-    def __init__(self):
+    def __init__(self, embedding_choice="OpenAI"):
         self.document_store = None
         self.qa_chain = None
         self.document_summary = ""
         self.chat_history = []
         self.last_processed_time = None
-        self.api_key = os.getenv("OPENAI_API_KEY")  # Fetch the API key from environment variable
+        self.api_key = os.getenv("OPENAI_API_KEY")  
         self.init_time = datetime.now(pytz.UTC)
-
-        if not self.api_key:
-            raise ValueError("API Key not found. Make sure to set the 'OPENAI_API_KEY' environment variable.")
+        self.embedding_choice = embedding_choice
 
         # Persistent directory for Chroma to avoid tenant-related errors
         self.chroma_persist_dir = "./chroma_storage"
         os.makedirs(self.chroma_persist_dir, exist_ok=True)
 
+
+    def _get_embedding_model(self):
+        if self.embedding_choice == "OpenAI":
+            return OpenAIEmbeddings(api_key=self.api_key)
+        else:
+            from langchain.embeddings import CohereEmbeddings
+            return CohereEmbeddings(
+                model="embed-multilingual-light-v3.0",
+                cohere_api_key=os.getenv("COHERE_API_KEY")
+            )
+
+        
+
+        if not self.api_key:
+            raise ValueError("API Key not found. Make sure to set the 'OPENAI_API_KEY' environment variable.")
+
+
+
     def process_documents(self, uploaded_files):
         """Process uploaded files by saving them temporarily and extracting content."""
         if not self.api_key:
@@ -118,7 +135,7 @@ class DocumentRAG:
             self.document_text = " ".join([doc.page_content for doc in documents])  # Store for later use
 
             # Create embeddings and initialize retrieval chain
-            embeddings = OpenAIEmbeddings(api_key=self.api_key)
+            embeddings = self._get_embedding_model()
             self.document_store = Chroma.from_documents(
                 documents,
                 embeddings,
@@ -294,29 +311,53 @@ class DocumentRAG:
     def topic_extractor(self, state: MultiAgentState):
         return {"sub_topic_list": self.extract_subtopics(state.sub_topics)}
 
-    def retrieve_docs(self, state: StoryState):
-        retriever = self.document_store.as_retriever(search_kwargs={"k": 20})
-        docs = retriever.get_relevant_documents(f"information about {state.story_topic}")
-        return {"retrieved_docs": docs}
 
-    def generate_story(self, state: StoryState):
-        context = "\n\n".join([doc.page_content for doc in state.retrieved_docs[:5]])
-        prompt = f"""You're a witty science storyteller. Create a short, child-friendly story that explains **{state.story_topic}** based on this:\n\n{context}"""
-        msg = self.llm.invoke([SystemMessage("Use humor. Be clear."), HumanMessage(prompt)])
-        return {"stories": msg}
+    def retrieve_node(self, state: StoryState):
+        embedding = self._get_embedding_model()
+        retriever = Chroma(
+            persist_directory=self.chroma_persist_dir,
+            embedding_function=embedding
+        ).as_retriever(search_kwargs={"k": 20})
 
-    def rerank_docs(self, state: StoryState):
         topic = state.story_topic
+        query = f"information about {topic}"
+        docs = retriever.get_relevant_documents(query)
+        return {"retrieved_docs": docs, "question": query}
+
+    def rerank_node(self, state: StoryState):
+        topic = state.story_topic
+        query = f"Rerank documents based on how well they explain the topic {topic}"
         docs = state.retrieved_docs
         texts = [doc.page_content for doc in docs]
 
-        # Fallback: return top 5 if no reranker available
         if not texts:
-            return {"reranked_docs": []}
+            return {"reranked_docs": [], "question": query}
+
+        # Quick fallback: rank by length
+        top_docs = sorted(texts, key=lambda t: -len(t))[:5]
+        return {"reranked_docs": top_docs, "question": query}
+
+
 
-        # Quick ranking by doc length (or use a real reranker if you have access)
-        ranked = sorted(texts, key=lambda t: -len(t))[:5]
-        return {"reranked_docs": ranked}
+    def generate_story_node(self, state: StoryState):
+        context = "\n\n".join(state.reranked_docs)
+        topic = state.story_topic
+
+        system_message = f"""
+        Suppose you're a brilliant science storyteller.
+        You write stories that help middle schoolers understand complex science topics with fun and clarity.
+        Add subtle humor and make it engaging.
+        """
+        prompt = f"""
+        Use the following context to write a fun and simple story explaining **{topic}** to a middle schooler:\n
+        Context:\n{context}\n\n
+        Story:
+        """
+
+        msg = self.llm.invoke([SystemMessage(system_message), HumanMessage(prompt)])
+        return {"stories": msg}
+
+    
 
 
     def run_multiagent_storygraph(self, topic: str, context: str):
@@ -324,9 +365,9 @@ class DocumentRAG:
 
         # Define the story subgraph with reranking
         story_graph = StateGraph(StoryState)
-        story_graph.add_node("Retrieve", self.retrieve_docs)
-        story_graph.add_node("Rerank", self.rerank_docs)  # Add rerank step
-        story_graph.add_node("Generate", self.generate_story)
+        story_graph.add_node("Retrieve", self.retrieve_node)
+        story_graph.add_node("Rerank", self.rerank_node)
+        story_graph.add_node("Generate", self.generate_story_node)
         story_graph.set_entry_point("Retrieve")
         story_graph.add_edge("Retrieve", "Rerank")
         story_graph.add_edge("Rerank", "Generate")
@@ -365,13 +406,9 @@ class DocumentRAG:
         return result
 
 
-
-
-
-# Initialize RAG system in session state
-if "rag_system" not in st.session_state:
-    st.session_state.rag_system = DocumentRAG()
-
+if "rag_system" not in st.session_state or st.session_state.embedding_model != embedding_choice:
+    st.session_state.embedding_model = embedding_choice
+    st.session_state.rag_system = DocumentRAG(embedding_choice=embedding_choice)