Update app.py

app.py

@@ -50,6 +50,7 @@ class MultiAgentState(BaseModel):
 
 class StoryState(BaseModel):
     retrieved_docs: List[Any] = []
+    reranked_docs: List[str] = []
     stories: Annotated[list[AnyMessage], add_messages]
     story_topic: str = ""
     stories_lst: Annotated[list, operator.add]
@@ -65,7 +66,24 @@ class DocumentRAG:
         self.init_time = datetime.now(pytz.UTC)
         self.embedding_choice = embedding_choice
 
-        # Persistent directory for Chroma to avoid tenant-related errors
+        # Set up appropriate LLM
+        if self.embedding_choice == "Cohere":
+            from langchain_cohere import ChatCohere
+            import cohere
+            self.llm = ChatCohere(
+                model="command-r-plus-08-2024",
+                temperature=0.7,
+                cohere_api_key=os.getenv("COHERE_API_KEY")
+            )
+            self.cohere_client = cohere.Client(os.getenv("COHERE_API_KEY"))
+        else:
+            self.llm = ChatOpenAI(
+                model_name="gpt-4",
+                temperature=0.7,
+                api_key=self.api_key
+            )
+
+        # Persistent directory for Chroma
         self.chroma_persist_dir = "./chroma_storage"
         os.makedirs(self.chroma_persist_dir, exist_ok=True)
 
@@ -324,6 +342,7 @@ class DocumentRAG:
         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}"
@@ -333,12 +352,22 @@ class DocumentRAG:
         if not texts:
             return {"reranked_docs": [], "question": query}
 
-        # Quick fallback: rank by length
-        top_docs = sorted(texts, key=lambda t: -len(t))[:5]
+        if self.embedding_choice == "Cohere":
+            rerank_results = self.cohere_client.rerank(
+                query=query,
+                documents=texts,
+                top_n=5,
+                model="rerank-v3.5"
+            )
+            top_docs = [texts[result.index] for result in rerank_results.results]
+        else:
+            top_docs = sorted(texts, key=lambda t: -len(t))[:5]
+
         return {"reranked_docs": top_docs, "question": query}
 
 
 
+
     def generate_story_node(self, state: StoryState):
         context = "\n\n".join(state.reranked_docs)
         topic = state.story_topic
@@ -361,7 +390,15 @@ class DocumentRAG:
 
 
     def run_multiagent_storygraph(self, topic: str, context: str):
-        self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)
+        if self.embedding_choice == "OpenAI":
+            self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)
+        elif self.embedding_choice == "Cohere":
+            from langchain_cohere import ChatCohere
+            self.llm = ChatCohere(
+                model="command-r-plus-08-2024",
+                temperature=0.7,
+                cohere_api_key=os.getenv("COHERE_API_KEY")
+            )
 
         # Define the story subgraph with reranking
         story_graph = StateGraph(StoryState)
@@ -374,7 +411,7 @@ class DocumentRAG:
         story_graph.set_finish_point("Generate")
         story_subgraph = story_graph.compile()
 
-        # Main graph setup
+        # Define the main graph
         graph = StateGraph(MultiAgentState)
         graph.add_node("beginner_topic", self.beginner_topic)
         graph.add_node("middle_topic", self.middle_topic)
@@ -386,18 +423,20 @@ class DocumentRAG:
         graph.add_edge("beginner_topic", "middle_topic")
         graph.add_edge("middle_topic", "advanced_topic")
         graph.add_edge("advanced_topic", "topic_extractor")
-        graph.add_conditional_edges("topic_extractor",
+        graph.add_conditional_edges(
+            "topic_extractor",
             lambda state: [Send("story_generator", {"story_topic": t}) for t in state.sub_topic_list],
-            ["story_generator"])
+            ["story_generator"]
+        )
         graph.add_edge("story_generator", END)
 
         compiled = graph.compile(checkpointer=MemorySaver())
         thread = {"configurable": {"thread_id": "storygraph-session"}}
 
-        # Initial run to extract subtopics
+        # Initial invocation
         result = compiled.invoke({"topic": [topic], "context": [context]}, thread)
 
-        # Fallback if no subtopics were extracted
+        # Fallback if no subtopics found
         if not result.get("sub_topic_list"):
             fallback_subs = ["Neural Networks", "Reinforcement Learning", "Supervised vs Unsupervised"]
             compiled.update_state(thread, {"sub_topic_list": fallback_subs})
@@ -406,6 +445,7 @@ class DocumentRAG:
         return result
 
 
+
 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)