add download button

src/download.py

@@ -0,0 +1,46 @@
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+import json
+
+def format_conversation_for_download(conversation: list) -> str:
+    """Formats the conversation history for download."""
+    formatted_messages = []
+    tool_id_to_name_map = {}
+    for message in conversation:
+        if isinstance(message, AIMessage) and message.tool_calls:
+            for tool_call in message.tool_calls:
+                tool_id_to_name_map[tool_call["id"]] = tool_call["name"]
+
+    for message in conversation:
+        if isinstance(message, HumanMessage):
+            formatted_messages.append({"role": "user", "content": message.content})
+        elif isinstance(message, AIMessage):
+            if not message.content and not message.tool_calls:
+                continue
+            
+            ai_dict = {"role": "assistant"}
+            if message.content:
+                ai_dict["content"] = message.content
+
+            if message.tool_calls:
+                ai_dict["tool_calls"] = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tc["name"],
+                            "arguments": tc["args"],
+                        },
+                    }
+                    for tc in message.tool_calls
+                ]
+            formatted_messages.append(ai_dict)
+        elif isinstance(message, ToolMessage):
+            tool_name = tool_id_to_name_map.get(message.tool_call_id)
+            formatted_messages.append(
+                {
+                    "role": "tool",
+                    "name": tool_name,
+                    "content": message.content,
+                }
+            )
+            
+    return json.dumps(formatted_messages, indent=4) 

src/streamlit_app.py

@@ -1,19 +1,22 @@
+from dotenv import load_dotenv
 import streamlit as st
 from receptionist_agent import get_agent_response
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, messages_to_dict
 import json
 from datetime import datetime
 import time
+from download import format_conversation_for_download
 
-# import langsmith
-# langsmith.debug = True
+#import langsmith
+
+load_dotenv()
+
+#langsmith.debug = True
 
 st.set_page_config(layout="wide")
 
 if "conversation" not in st.session_state:
-    st.session_state.conversation = [
-        AIMessage(content="Thanks for calling, this is Receptionist Agent! Can I help you schedule a service or would you like help with something else?")
-    ]
+    st.session_state.conversation = [AIMessage(content="Thanks for calling, this is Receptionist Agent! Can I help you schedule a service or would you like help with something else?")]
     start_time = datetime.now().strftime("%y%m%d-%H%M%S")
     st.session_state.conversation_file = f"conversations/conversation_{start_time}.json"
 
@@ -31,7 +34,7 @@ def submit_message():
     st.session_state.response_time = f"{end - start:.2f}"
     
     st.session_state.conversation.extend(new_messages)
-    # save_conversation()
+    save_conversation()
 
 st.subheader("Receptionist Agent")
 
@@ -50,6 +53,16 @@ for message in st.session_state.conversation:
         with st.chat_message("tool"):
             st.write(f"Tool `{message.tool_call_id}` output: `{message.content}`")
 
+if len(st.session_state.conversation) > 1:
+    _, col2 = st.columns([0.75, 0.25])
+    with col2:
+        st.download_button(
+            label="Download Conversation",
+            data=format_conversation_for_download(st.session_state.conversation),
+            file_name=st.session_state.conversation_file.replace(".json", "_formatted.json"),
+            mime="application/json",
+        )
+
 if "response_time" in st.session_state:
     st.caption(f"Response time: {st.session_state.response_time}s")