adding session support to the client

app.py

@@ -1,5 +1,14 @@
 import gradio as gr
-from task_management import get_task_status, a2a_respond
+from task_management import (
+    get_task_status,
+    get_task_history,
+    send_task,
+    cancel_task,
+    new_session,
+    get_current_session_id,
+    list_sessions,
+    change_session
+)
 from server_management import (
     connect_to_server,
     disconnect_from_server,
@@ -74,15 +83,24 @@ def fetch_tasks() -> list:
         statuses.append(f"{tid}: {status.get('result', {}).get('status', 'unknown')}")
     return statuses
 
+def create_session_and_select() -> list:
+    new_session_id = new_session()
+    sessions = list_sessions()
+    return gr.Dropdown(list(sessions), value=new_session_id)
+
+def load_session_history(session_id: str) -> list:
+    history = get_task_history(session_id)
+    return history
+
 def respond(message: str, history: list) -> tuple[str, list]:
     server_url = get_selected_server()
-    return a2a_respond(message, history, server_url)
+    return "", send_task(server_url, message)
 
 def create_chat_interface():
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column(scale=1):
-                server_url_input = gr.Textbox(label="Server URL", placeholder="Enter server URL")
+                server_url_input = gr.Textbox(label="Server URL", value="https://agent-too-agent.azurewebsites.net/api/")
                 add_server_btn = gr.Button("Add Server")
                 servers_list = gr.Dropdown(
                     label="Available Servers",
@@ -92,8 +110,13 @@ def create_chat_interface():
                     allow_custom_value=True
                 )
                 remove_server_btn = gr.Button("Remove Selected Server")
-                tasks_list = gr.List(label="Executing Tasks")
-                refresh_tasks_btn = gr.Button("Refresh Tasks")
+                sessions_list = gr.Dropdown(
+                    label="Previous Sessions",
+                    choices=[],
+                    value=None,
+                    interactive=True,
+                    allow_custom_value=True)
+                new_session_btn = gr.Button("New Session")
 
                 add_server_btn.click(
                     add_server_and_select,
@@ -119,12 +142,14 @@ def create_chat_interface():
                     outputs=servers_list,
                 )
 
-                refresh_tasks_btn.click(
-                    fetch_tasks,
+                new_session_btn.click(
+                    create_session_and_select,
                     inputs=None,
-                    outputs=tasks_list,
+                    outputs=sessions_list,
                 )
 
+
+
             with gr.Column(scale=2):
                 chatbot = gr.Chatbot(label="Chat", type="messages")
                 chat_input = gr.Textbox(label="Message", placeholder="Type a message...")
@@ -136,6 +161,13 @@ def create_chat_interface():
                     outputs=[chat_input, chatbot],
                     queue=False,
                 )
+            
+                sessions_list.change(
+                    load_session_history,
+                    inputs=sessions_list,
+                    outputs=chatbot,
+                )
+
     return demo
 
 if __name__ == "__main__":

task_management.py

@@ -1,28 +1,91 @@
 import uuid
 import requests
 
-def send_task(server_url, task_id, thread, content):
+# Data class of a session with ID and task history
+class Session:
+    def __init__(self, session_id):
+        self.session_id = session_id
+        #tasks are a list of tuples with role and content
+        self.task_history = []
+
+    def add_task(self, task):
+        self.task_history.append(task)
+
+    def get_history(self):
+        return self.task_history
+
+# Global session tracking
+sessions = {}
+current_session_id = None
+
+def new_session():
+    """ Creates a new session, sets it as current, and returns the session ID. """
+    global current_session_id
+    session_id = str(uuid.uuid4())
+    # add a new Session object to the sessions dictionary
+    sessions[session_id] = Session(session_id)
+    current_session_id = session_id
+    return session_id
+
+def get_current_session_id():
+    """ Returns the current active session ID, or None if none exists. """
+    return current_session_id
+
+def list_sessions():
+    """ Returns a list of all session IDs being tracked. """
+    return list(sessions.keys())
+
+def change_session(session_id):
+    """ Changes the current session to the specified session ID. """
+    global current_session_id
+    if session_id in sessions:
+        current_session_id = session_id
+    else:
+        raise ValueError(f"Session ID {session_id} not found.")
+
+def get_task_history(session_id):
+    """ Retrieves the task history for a given session ID."""
+    return sessions[session_id].get_history()
+
+def send_task(server_url, content):
+    """
+    Sends a message to the A2A server using the current session ID.
+    If no session exists, a new one is created automatically.
+    """
+    global current_session_id
+    if not current_session_id:
+        new_session()
+
+    task_id = str(uuid.uuid4())
+    sessions[current_session_id].add_task({"role": "user", "content": content})
+
     payload = {
         "jsonrpc": "2.0",
         "method": "tasks/send",
         "params": {
-            "sessionId": thread,
-            "message": content
+            "sessionId": current_session_id,
+            "message": content,
+            "historyLength": 5
         },
         "id": task_id
     }
-    #remove trailing slash if present
-    if server_url.endswith("/"):
-        server_url = server_url[:-1]
-    json_endpoint = f"{server_url}/jsonrpc"
+
+    json_endpoint = f"{server_url.rstrip('/')}/jsonrpc"
     response = requests.post(json_endpoint, json=payload)
     response.raise_for_status()
-    return response.json()
+    result = response.json()
 
-def get_task_status(server_url, task_id):
-    import requests
-    import json
+    if "error" in result:
+        error_message = result["error"].get("message", "Unknown error")
+        sessions[current_session_id].append({"role": "assistant", "content": f"Error: {error_message}"})
+        return f"Error: {error_message}"
 
+    agent_reply = result.get("result", {}).get("content", "")
+    if agent_reply:
+        sessions[current_session_id].add_task({"role": "assistant", "content": agent_reply})
+    return sessions[current_session_id].get_history()
+
+def get_task_status(server_url, task_id):
     payload = {
         "jsonrpc": "2.0",
         "method": "tasks/get",
@@ -36,9 +99,6 @@ def get_task_status(server_url, task_id):
     return response.json()
 
 def cancel_task(server_url, task_id):
-    import requests
-    import json
-
     payload = {
         "jsonrpc": "2.0",
         "method": "tasks/cancel",
@@ -50,34 +110,3 @@ def cancel_task(server_url, task_id):
 
     response = requests.post(server_url, json=payload)
     return response.json()
-
-def a2a_respond(user_message, history, server_url):
-    """
-    Sends a plain JSON POST to the server_url with user_input and session_id.
-    Only adds the 'content' field from the response to the message history.
-    """
-    if not server_url:
-        history = history or []
-        history.append({"role": "assistant", "content": "No server selected."})
-        return "", history
-
-    session_id = str(uuid.uuid4())
-
-    history = history or []
-    history.append({"role": "user", "content": user_message})
-
-    try:
-        # Call send_task with the server_url and session_id
-        response = send_task(server_url, session_id, "thread", user_message)
-        if response.get("error"):
-            error_message = response["error"].get("message", "Unknown error")
-            history.append({"role": "assistant", "content": f"Error: {error_message}"})
-            return "", history
-        # get the result.content from the json
-        agent_reply = response.get("result", {}).get("content", "")
-        if agent_reply:
-            history.append({"role": "assistant", "content": agent_reply})
-    except Exception as e:
-        history.append({"role": "assistant", "content": f"Error: {e}"})
-
-    return "", history