Update app.py

app.py

@@ -18,10 +18,13 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 st.set_page_config(page_title="IT Support System (RAG)", layout="centered")
 
-# Initialize session memory
+# Initialize session state for chat and logs
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
+if "workflow_logs" not in st.session_state:
+    st.session_state.workflow_logs = []
+
 # Knowledge Base Setup
 kb_path = os.path.join(os.path.dirname(__file__), 'kb.json')
 with open(kb_path, encoding='utf-8') as f:
@@ -61,10 +64,9 @@ def escalate_ticket(query: str, analysis: str = "") -> str:
     logger.info(f"Escalating issue with ticket {ticket_id}: {description}")
     return f"Escalated issue. Created ticket {ticket_id}. A support technician will contact you shortly."
 
-# LLM Configuration (local/open-source LLM API)
 llm_config = {
     "config_list": [{
-        "model": "llama3",  # or mistral, gemma, etc.
+        "model": "llama3",
         "base_url": os.getenv("LLM_API_BASE", "http://localhost:11434/v1"),
         "api_key": "NULL",
     }],
@@ -72,7 +74,6 @@ llm_config = {
     "temperature": 0.5,
 }
 
-# Agent Definitions
 master_agent = AssistantAgent(
     name="Master",
     llm_config=llm_config,
@@ -165,17 +166,14 @@ Always include the ticket ID and expected follow-up timeframe.
     function_map={"escalate_ticket": escalate_ticket}
 )
 
-def handle_it_query(query: str, show_logs: bool = False) -> str:
+def handle_it_query(query: str) -> (str, dict):
     query = query.strip()
     if not query:
-        return "Please enter an IT question or issue."
+        return "Please enter an IT question or issue.", {}
 
     workflow_logs = {"query": query}
-
     try:
-        master_proxy = UserProxyAgent(
-            name="MasterProxy", human_input_mode="NEVER", code_execution_config=False
-        )
+        master_proxy = UserProxyAgent(name="MasterProxy", human_input_mode="NEVER", code_execution_config=False)
         master_prompt = f"User query: '{query}'. First, determine if this is an IT-related issue."
         master_proxy.initiate_chat(master_agent, message=master_prompt, max_turns=1)
         initial_assessment = master_proxy.chat_messages[master_agent][-1]["content"]
@@ -183,27 +181,21 @@ def handle_it_query(query: str, show_logs: bool = False) -> str:
         workflow_logs["initial_assessment"] = initial_assessment
 
         if "NOT IT-RELATED" in initial_assessment.upper():
-            return initial_assessment
+            return initial_assessment, workflow_logs
 
-        plan_proxy = UserProxyAgent(
-            name="PlanningProxy", human_input_mode="NEVER", code_execution_config=False
-        )
+        plan_proxy = UserProxyAgent(name="PlanningProxy", human_input_mode="NEVER", code_execution_config=False)
         plan_proxy.initiate_chat(planning_agent, message=query, max_turns=1)
         planning_output = plan_proxy.chat_messages[planning_agent][-1]["content"]
         logger.info(f"Planning Agent Response: {planning_output}")
         workflow_logs["planning"] = planning_output
 
-        analysis_proxy = UserProxyAgent(
-            name="AnalysisProxy", human_input_mode="NEVER", code_execution_config=False
-        )
+        analysis_proxy = UserProxyAgent(name="AnalysisProxy", human_input_mode="NEVER", code_execution_config=False)
         analysis_proxy.initiate_chat(analysis_agent, message=planning_output, max_turns=1)
         analysis_output = analysis_proxy.chat_messages[analysis_agent][-1]["content"]
         logger.info(f"Analysis Agent Response: {analysis_output}")
         workflow_logs["analysis"] = analysis_output
 
-        res_proxy = UserProxyAgent(
-            name="ResolutionProxy", human_input_mode="NEVER", code_execution_config=False
-        )
+        res_proxy = UserProxyAgent(name="ResolutionProxy", human_input_mode="NEVER", code_execution_config=False)
         resolution_input = f"User Query: {query}\n\nPlanning: {planning_output}\n\nAnalysis: {analysis_output}"
         res_proxy.initiate_chat(resolution_agent, message=resolution_input, max_turns=1)
         resolution_output = res_proxy.chat_messages[resolution_agent][-1]["content"]
@@ -212,18 +204,14 @@ def handle_it_query(query: str, show_logs: bool = False) -> str:
 
         escalation_output = None
         if "ESCALATION NEEDED" in resolution_output.upper():
-            esc_proxy = UserProxyAgent(
-                name="EscalationProxy", human_input_mode="NEVER", code_execution_config=False
-            )
+            esc_proxy = UserProxyAgent(name="EscalationProxy", human_input_mode="NEVER", code_execution_config=False)
             escalation_input = f"Original Query: {query}\n\nAnalysis: {analysis_output}\n\nResolution Attempt: {resolution_output}"
             esc_proxy.initiate_chat(escalation_agent, message=escalation_input, max_turns=1)
             escalation_output = esc_proxy.chat_messages[escalation_agent][-1]["content"]
             logger.info(f"Escalation Agent Response: {escalation_output}")
             workflow_logs["escalation"] = escalation_output
 
-        final_master_proxy = UserProxyAgent(
-            name="FinalMasterProxy", human_input_mode="NEVER", code_execution_config=False
-        )
+        final_master_proxy = UserProxyAgent(name="FinalMasterProxy", human_input_mode="NEVER", code_execution_config=False)
 
         if escalation_output:
             final_prompt = (
@@ -248,25 +236,11 @@ def handle_it_query(query: str, show_logs: bool = False) -> str:
         logger.info(f"Final Master Agent Response: {final_response}")
         workflow_logs["final_response"] = final_response
 
-        # Store chat history with or without logs based on user preference
-        if show_logs:
-            st.session_state.chat_history.append({
-                "user": query,
-                "assistant": final_response,
-                "workflow_logs": workflow_logs
-            })
-        else:
-            st.session_state.chat_history.append({
-                "user": query,
-                "assistant": final_response,
-            })
-
-        return final_response
+        return final_response, workflow_logs
 
     except Exception as e:
         logger.error(f"Error in workflow: {e}", exc_info=True)
-        return f"An error occurred during processing: {str(e)}\n\nPlease try rephrasing your question."
-
+        return f"An error occurred during processing: {str(e)}\n\nPlease try rephrasing your question.", {}
 
 
 st.title("AI Help Desk")
@@ -277,21 +251,31 @@ with st.form(key="query_form", clear_on_submit=True):
     show_logs = st.checkbox("Show workflow details", value=False)
     submitted = st.form_submit_button("Submit")
 
-    if submitted:
-        if not user_input:
-            st.error("Please type a message before submitting.")
-        else:
-            with st.spinner("Processing your request through our agent workflow..."):
-                response = handle_it_query(user_input, show_logs=show_logs)
-
-            st.markdown("### Response")
-            st.write(response)
-
-            if show_logs and st.session_state.chat_history:
-                st.markdown("---")
-                st.markdown("### Workflow Logs")
-                logs = st.session_state.chat_history[-1].get("workflow_logs", {})
-                for step, content in logs.items():
-                    st.markdown(f"**{step.capitalize().replace('_', ' ')}:**")
-                    st.text(content)
-                    st.markdown("---")
+if submitted:
+    if not user_input:
+        st.error("Please type a message before submitting.")
+    else:
+        with st.spinner("Processing your request through our agent workflow..."):
+            response, logs = handle_it_query(user_input)
+
+        # Append to session history to maintain conversation flow
+        st.session_state.chat_history.append({"user": user_input, "assistant": response})
+        if show_logs:
+            st.session_state.workflow_logs.append(logs)
+
+# Display conversation history
+st.markdown("## Conversation")
+for chat in st.session_state.chat_history:
+    st.markdown(f"**User:** {chat['user']}")
+    st.markdown(f"**Assistant:** {chat['assistant']}")
+    st.markdown("---")
+
+# Display logs if requested
+if show_logs and st.session_state.workflow_logs:
+    st.markdown("## Workflow Logs")
+    for i, log in enumerate(st.session_state.workflow_logs):
+        st.markdown(f"### Query {i+1} Logs")
+        for step, content in log.items():
+            st.markdown(f"**{step.capitalize().replace('_', ' ')}:**")
+            st.text(content)
+            st.markdown("---")