Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import json
 import random
 import logging
-from typing import List, Dict, Any
+from typing import List
 import streamlit as st
 from dotenv import load_dotenv
 import autogen
@@ -49,7 +49,6 @@ def retrieve_docs(query: str) -> str:
     hits = retriever.get_relevant_documents(query)
     if not hits:
         return "No relevant documentation found."
-
     results = []
     for doc in hits:
         question = doc.metadata.get('question', 'FAQ')
@@ -62,7 +61,7 @@ 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 (for open-source model served locally or from Hugging Face)
+# LLM Configuration (local/open-source LLM API)
 llm_config = {
     "config_list": [{
         "model": "llama3",  # or mistral, gemma, etc.
@@ -166,8 +165,7 @@ Always include the ticket ID and expected follow-up timeframe.
     function_map={"escalate_ticket": escalate_ticket}
 )
 
-
-def handle_it_query(query: str) -> str:
+def handle_it_query(query: str, show_logs: bool = False) -> str:
     query = query.strip()
     if not query:
         return "Please enter an IT question or issue."
@@ -181,6 +179,7 @@ def handle_it_query(query: str) -> str:
         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"]
+        logger.info(f"Master Agent Response: {initial_assessment}")
         workflow_logs["initial_assessment"] = initial_assessment
 
         if "NOT IT-RELATED" in initial_assessment.upper():
@@ -191,6 +190,7 @@ def handle_it_query(query: str) -> str:
         )
         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(
@@ -198,6 +198,7 @@ def handle_it_query(query: str) -> str:
         )
         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(
@@ -206,6 +207,7 @@ def handle_it_query(query: str) -> str:
         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"]
+        logger.info(f"Resolution Agent Response: {resolution_output}")
         workflow_logs["resolution"] = resolution_output
 
         escalation_output = None
@@ -216,6 +218,7 @@ def handle_it_query(query: str) -> str:
             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(
@@ -242,13 +245,21 @@ def handle_it_query(query: str) -> str:
 
         final_master_proxy.initiate_chat(master_agent, message=final_prompt, max_turns=1)
         final_response = final_master_proxy.chat_messages[master_agent][-1]["content"]
+        logger.info(f"Final Master Agent Response: {final_response}")
         workflow_logs["final_response"] = final_response
 
-        st.session_state.chat_history.append({
-            "user": query,
-            "assistant": final_response,
-            "workflow_logs": workflow_logs
-        })
+        # 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
 
@@ -257,6 +268,7 @@ def handle_it_query(query: str) -> str:
         return f"An error occurred during processing: {str(e)}\n\nPlease try rephrasing your question."
 
 
+
 st.title("AI Help Desk")
 st.write("Ask any IT support question and our multi-agent system will assist you.")
 
@@ -270,6 +282,16 @@ with st.form(key="query_form", clear_on_submit=True):
             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)
-
-            st.markdown("### Response")
+                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("---")