Update app.py

app.py

@@ -65,6 +65,9 @@ class AppConfig:
             """)
             st.stop()
 
+# ------------------------------
+# Chroma Setup
+# ------------------------------
 class ChromaManager:
     def __init__(self, config: AppConfig):
         os.makedirs(config.CHROMA_PATH, exist_ok=True)
@@ -127,7 +130,13 @@ class DocumentProcessor:
         }))
 
 # ------------------------------
-# Workflow
+# Workflow State
+# ------------------------------
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+
+# ------------------------------
+# Workflow Setup
 # ------------------------------
 class AgentWorkflow:
     def __init__(self, chroma: ChromaManager):
@@ -168,9 +177,6 @@ class AgentWorkflow:
         
         self.app = self.workflow.compile()
 
-    class AgentState(TypedDict):
-        messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
-
     def agent(self, state: AgentState):
         try:
             messages = state["messages"]
@@ -259,7 +265,7 @@ def apply_theme():
     </style>
     """, unsafe_allow_html=True)
 
-def main():
+def main(config: AppConfig, chroma: ChromaManager):
     apply_theme()
     
     with st.sidebar:
@@ -276,14 +282,18 @@ def main():
     
     if st.button("Analyze"):
         with st.spinner("Processing..."):
-            workflow = AgentWorkflow(chroma_manager)
-            results = workflow.app.invoke({"messages": [HumanMessage(content=query)]})
-            
-            with st.expander("Processing Details", expanded=True):
-                st.write("### Raw Results", results)
+            try:
+                workflow = AgentWorkflow(chroma)
+                results = workflow.app.invoke({"messages": [HumanMessage(content=query)]})
+                
+                with st.expander("Processing Details", expanded=True):
+                    st.write("### Raw Results", results)
+                    
+                st.success("### Final Answer")
+                st.markdown(results['messages'][-1].content)
                 
-            st.success("### Final Answer")
-            st.markdown(results['messages'][-1].content)
+            except Exception as e:
+                st.error(f"Processing failed: {str(e)}")
 
 # ------------------------------
 # Initialization
@@ -298,7 +308,7 @@ if __name__ == "__main__":
     try:
         config = AppConfig()
         config.validate()
-        chroma_manager = ChromaManager(config)
-        main()
+        chroma = ChromaManager(config)
+        main(config, chroma)
     except Exception as e:
         st.error(f"Initialization failed: {str(e)}")