Update app.py

app.py

@@ -27,7 +27,29 @@ import os
 import tempfile
 from datetime import datetime
 import pytz
-
+from langgraph.graph import StateGraph, START, END, Send, add_messages
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import HumanMessage, SystemMessage, AnyMessage
+from pydantic import BaseModel
+from typing import List, Annotated, Any
+import re, operator
+
+
+class MultiAgentState(BaseModel):
+    state: List[str] = []
+    messages: Annotated[list[AnyMessage], add_messages]
+    topic: List[str] = []
+    context: List[str] = []
+    sub_topic_list: List[str] = []
+    sub_topics: Annotated[list[AnyMessage], add_messages]
+    stories: Annotated[list[AnyMessage], add_messages]
+    stories_lst: Annotated[list, operator.add]
+
+class StoryState(BaseModel):
+    retrieved_docs: List[Any] = []
+    stories: Annotated[list[AnyMessage], add_messages]
+    story_topic: str = ""
+    stories_lst: Annotated[list, operator.add]
 
 class DocumentRAG:
     def __init__(self):
@@ -248,10 +270,82 @@ class DocumentRAG:
         except Exception as e:
             return history + [("System", f"Error: {str(e)}")]
 
+    def extract_subtopics(self, messages):
+        text = "\n".join([msg.content for msg in messages])
+        return re.findall(r'- \*\*(.*?)\*\*', text)
+
+    def beginner_topic(self, state: MultiAgentState):
+        prompt = f"What are the beginner-level topics you can learn about {', '.join(state.topic)} in {', '.join(state.context)}?"
+        msg = self.llm.invoke([SystemMessage("Suppose you're a middle grader..."), HumanMessage(prompt)])
+        return {"message": msg, "sub_topics": msg}
+
+    def middle_topic(self, state: MultiAgentState):
+        prompt = f"What are the middle-level topics for {', '.join(state.topic)} in {', '.join(state.context)}? Avoid previous."
+        msg = self.llm.invoke([SystemMessage("Suppose you're a college student..."), HumanMessage(prompt)])
+        return {"message": msg, "sub_topics": msg}
+
+    def advanced_topic(self, state: MultiAgentState):
+        prompt = f"What are the advanced-level topics for {', '.join(state.topic)} in {', '.join(state.context)}? Avoid previous."
+        msg = self.llm.invoke([SystemMessage("Suppose you're a teacher..."), HumanMessage(prompt)])
+        return {"message": msg, "sub_topics": msg}
+
+    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 run_multiagent_storygraph(self, topic: str, context: str):
+        self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)
+
+        # Story subgraph
+        story_graph = StateGraph(StoryState)
+        story_graph.add_node("Retrieve", self.retrieve_docs)
+        story_graph.add_node("Generate", self.generate_story)
+        story_graph.set_entry_point("Retrieve")
+        story_graph.add_edge("Retrieve", "Generate")
+        story_graph.set_finish_point("Generate")
+        story_subgraph = story_graph.compile()
+
+        # Main graph
+        graph = StateGraph(MultiAgentState)
+        graph.add_node("beginner_topic", self.beginner_topic)
+        graph.add_node("middle_topic", self.middle_topic)
+        graph.add_node("advanced_topic", self.advanced_topic)
+        graph.add_node("topic_extractor", self.topic_extractor)
+        graph.add_node("story_generator", story_subgraph)
+
+        graph.add_edge(START, "beginner_topic")
+        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",
+            lambda state: [Send("story_generator", {"story_topic": t}) for t in state.sub_topic_list],
+            ["story_generator"])
+        graph.add_edge("story_generator", END)
+
+        compiled = graph.compile(checkpointer=MemorySaver())
+        result = compiled.invoke({"topic": [topic], "context": [context]})
+        return result
+
+
+
+
 # Initialize RAG system in session state
 if "rag_system" not in st.session_state:
     st.session_state.rag_system = DocumentRAG()
 
+
+
+
 # Sidebar
 with st.sidebar:
     st.title("About")
@@ -337,7 +431,31 @@ if st.session_state.rag_system.qa_chain:
 else:
     st.info("Please process documents first to enable Q&A.")
 
-# Step 4: Generate Podcast
+
+# Step 4: Multi-Agent Story Explorer
+st.subheader("Step 5: Explore Subtopics via Multi-Agent Graph")
+story_topic = st.text_input("Enter main topic:", value="Machine Learning")
+story_context = st.text_input("Enter learning context:", value="Education")
+
+if st.button("Run Story Graph"):
+    with st.spinner("Generating subtopics and stories..."):
+        result = st.session_state.rag_system.run_multiagent_storygraph(topic=story_topic, context=story_context)
+
+        subtopics = result.get("sub_topic_list", [])
+        st.markdown("### 🧠 Extracted Subtopics")
+        for sub in subtopics:
+            st.markdown(f"- {sub}")
+
+        stories = result.get("stories", [])
+        if stories:
+            st.markdown("### 📚 Generated Stories")
+            for i, story in enumerate(stories):
+                st.markdown(f"**Story {i+1}:**")
+                st.markdown(story.content)
+        else:
+            st.warning("No stories were generated.")
+
+# Step 5: Generate Podcast
 st.subheader("Step 4: Generate Podcast")
 st.write("Select Podcast Language:")
 podcast_language_options = ["English", "Hindi", "Spanish", "French", "German", "Chinese", "Japanese"]
@@ -348,6 +466,7 @@ podcast_language = st.radio(
     key="podcast_language"
 )
 
+
 if st.session_state.rag_system.document_summary:
     if st.button("Generate Podcast"):
         with st.spinner("Generating podcast, please wait..."):