Update app.py

app.py

@@ -331,11 +331,11 @@ class DocumentRAG:
 
 
     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})
+        if not self.document_store:
+            return {"retrieved_docs": [], "question": "No documents processed yet."}
+        
+        retriever = self.document_store.as_retriever(search_kwargs={"k": 20})
+
 
         topic = state.story_topic
         query = f"information about {topic}"
@@ -556,22 +556,26 @@ 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.")
+    if st.session_state.rag_system.document_store is None:
+        st.warning("Please process documents first before running the story graph.")
+    else:
+        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")