working version of using llama stack with multi-turn Qs

docs/zero_to_hero_guide/gmail_agent/functions_prompt.py

@@ -3,6 +3,7 @@ from llama_stack_client.types.tool_param_definition_param import ToolParamDefini
 from llama_stack_client.types import CompletionMessage, ToolResponseMessage
 from llama_stack_client.lib.agents.custom_tool import CustomTool
 from gmagent import *
+import json
 
 class ListEmailsTool(CustomTool):
     """Custom tool for List Emails."""
@@ -49,15 +50,14 @@ class ListEmailsTool(CustomTool):
     async def run_impl(self, query: str, maxResults: int = 100) -> Dict[str, Any]:
         """Query to get a list of emails matching the query."""
         emails = list_emails(query)
-        print(emails)
-        return emails
+        return {"name": self.get_name(), "result": emails}
 
 
-class GetEmailTool(CustomTool):
+class GetEmailDetailTool(CustomTool):
     """Custom tool for Get Email Detail."""
 
     def get_name(self) -> str:
-        return "get_email"
+        return "get_email_detail"
 
     def get_description(self) -> str:
         return "Get detailed info about a specific email"
@@ -65,9 +65,9 @@ class GetEmailTool(CustomTool):
     def get_params_definition(self) -> Dict[str, ToolParamDefinitionParam]:
         return {
             "detail": ToolParamDefinitionParam(
-                param_type="string",
+                param_type="str",
                 description="what detail the user wants to know about - two possible values: body or attachment",
-                required=False
+                required=True
             ),
             "query": ToolParamDefinitionParam(
                 param_type="str",
@@ -95,12 +95,11 @@ class GetEmailTool(CustomTool):
         )
         return [message]
 
-    async def run_impl(self, query: str, maxResults: int = 100) -> Dict[str, Any]:
-        """Query to get a list of emails matching the query."""
-
-        emails = []
-        return emails
+    async def run_impl(self, detail: str, query: str) -> Dict[str, Any]:
+        """Query to get the detail of an email."""
 
+        detail = get_email_detail(detail, query)
+        return {"name": self.get_name(), "result": detail}
 
 
 class SendEmailTool(CustomTool):

docs/zero_to_hero_guide/gmail_agent/gmagent.py

@@ -12,10 +12,10 @@ import pytz
 import base64
 import pickle
 from datetime import datetime, timezone
-import json
 import ollama
 from pypdf import PdfReader
 from pathlib import Path
+from shared import memory
 
 SCOPES = ['https://www.googleapis.com/auth/gmail.readonly', 'https://www.googleapis.com/auth/gmail.compose']
 
@@ -84,9 +84,33 @@ def list_emails(query='', max_results=100):
     
     return emails
 
-def get_email_detail(detail, which=''):
+def get_email_detail(detail, which):
+    message_id = None
+    # pre-processing
+    if 'from ' in which:
+        sender = which.split('from ')[-1]
+        for email in memory['emails']:
+            if email['sender'].find(sender) != -1:
+                message_id = email['message_id']
+                break
+    elif 'subject:' in which:
+        subject = which.split('subject:')[-1]
+        # exact match beats substring
+        for email in memory['emails']:
+            if email['subject'].upper() == subject.upper():
+                message_id = email['message_id']
+                break
+            elif email['subject'].upper().find(subject.upper()) != -1:
+                message_id = email['message_id']
+
+    elif 'id_' in which:
+        message_id = which.split('id_')[-1]
+    else:
+        message_id = memory['emails'][-1]['message_id']
+
+
     if detail == 'body':
-        return get_email_body(which)
+        return get_email_body(message_id)
     elif detail == 'attachment':
         return get_email_attachments(which)
 

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",
-        ]
-
-        for query in queries:
-            print(f"\nQuery: {query}")
-            print("-" * 50)
-
         response = agent.create_turn(
-                messages=[{"role": "user", "content": query}],
+            messages=[{"role": "user", "content": ask}],
             session_id=session_id,
         )
 
         async for log in EventLogger().log(response):
-                log.print()
+            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"
+
+                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: "
+