Update app.py

app.py

@@ -305,6 +305,20 @@ class DocumentRAG:
         msg = self.llm.invoke([SystemMessage("Use humor. Be clear."), HumanMessage(prompt)])
         return {"stories": msg}
 
+    def rerank_docs(self, state: StoryState):
+        topic = state.story_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": []}
+
+        # 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 run_multiagent_storygraph(self, topic: str, context: str):
         self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)