working version of using llama stack with multi-turn Qs

docs/zero_to_hero_guide/gmail_agent/main.py

@@ -1,8 +1,8 @@
 import argparse
 import gmagent
 import asyncio
-from gmagent import *
-from functions_prompt import * #system_prompt
+import json
+from functions_prompt import *
 
 from llama_stack_client import LlamaStackClient
 from llama_stack_client.lib.agents.agent import Agent
@@ -11,6 +11,8 @@ from llama_stack_client.types.agent_create_params import (
     AgentConfig,
 )
 
+from shared import memory
+
 LLAMA_STACK_API_TOGETHER_URL="https://llama-stack.together.ai"
 LLAMA31_8B_INSTRUCT = "Llama3.1-8B-Instruct"
 
@@ -18,7 +20,7 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
     """Create an agent with gmail tool capabilities."""
 
     listEmailsTool = ListEmailsTool()
-    getEmailTool = GetEmailTool()
+    getEmailDetailTool = GetEmailDetailTool()
     sendEmailTool = SendEmailTool()
     getPDFSummaryTool = GetPDFSummaryTool()
     createDraftTool = CreateDraftTool()
@@ -34,7 +36,7 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
         },
         tools=[
             listEmailsTool.get_tool_definition(),
-            getEmailTool.get_tool_definition(),
+            getEmailDetailTool.get_tool_definition(),
             sendEmailTool.get_tool_definition(),
             getPDFSummaryTool.get_tool_definition(),
             createDraftTool.get_tool_definition(),
@@ -52,7 +54,7 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
         client=client,
         agent_config=agent_config,
         custom_tools=[listEmailsTool,
-                      getEmailTool,
+                      getEmailDetailTool,
                       sendEmailTool,
                       getPDFSummaryTool,
                       createDraftTool,
@@ -62,9 +64,6 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
     return agent
 
 
-
-
-
 async def main():
     parser = argparse.ArgumentParser(description="Set email address")
     parser.add_argument("--gmail", type=str, required=True, help="Your Gmail address")
@@ -74,7 +73,9 @@ async def main():
 
     greeting = llama31("hello", "Your name is Gmagent, an assistant that can perform all Gmail related tasks for your user.")
     agent_response = f"{greeting}\n\nYour ask: "
-    #agent = Agent(system_prompt)
+
+    # do i have emails with attachment larger than 5mb?
+    # what's the detail of the email with subject this is an interesting paper
 
     while True:
         ask = input(agent_response)
@@ -82,30 +83,48 @@ async def main():
             print(llama31("bye"))
             break
         print("\n-------------------------\nCalling Llama...")
-        # agent(ask)
-        # agent_response = "Your ask: "
-
 
         client = LlamaStackClient(base_url=LLAMA_STACK_API_TOGETHER_URL)
         agent = await create_gmail_agent(client)
         session_id = agent.create_session("email-session")
 
-        queries = [
-            "do i have any emails with attachments?",
-            "what's the content of the email from LangSmith",
-        ]
+        response = agent.create_turn(
+            messages=[{"role": "user", "content": ask}],
+            session_id=session_id,
+        )
 
-        for query in queries:
-            print(f"\nQuery: {query}")
-            print("-" * 50)
+        async for log in EventLogger().log(response):
+            if log.role == "CustomTool":
+                tool_name = json.loads(log.content)['name']
+                result = json.loads(log.content)['result']
+                if tool_name == 'list_emails':
+                    # post processing
+                    memory['emails'] = result
+                    num = len(result)
+                    if num == 0:
+                        output = "I couldn't find any such emails. What else would you like to do?"
+                    elif num <= 5:
+                        output = f"I found {num} email{'s' if num > 1 else ''} matching your query:\n"
+                        for i, email in enumerate(result, start=1):
+                            output += f"{i}. From: {email['sender']}, Subject: {email['subject']}, Received on: {email['received_time']}\n"
+                    else:
+                        output = f"I found {num} emails matching your query. Here are the first 5 emails:\n"
+                        for i in range(1, 6):
+                            output += f"{i}. From: {result[i - 1]['sender']}, Subject: {result[i - 1]['subject']}, Received on: {result[i - 1]['received_time']}\n"
 
-            response = agent.create_turn(
-                messages=[{"role": "user", "content": query}],
-                session_id=session_id,
-            )
+                elif tool_name == "get_email_detail":
+                    output = result
+
+                print(f"\n-------------------------\n\nGmagent: {output}\n")
+            elif log.role == "inference":
+                print("Llama returned: ", end="")
+            else:
+                print(log, end="")
+
+
+
+        agent_response = "\n\nYour ask: "
 
-            async for log in EventLogger().log(response):
-                log.print()