diff --git a/docs/zero_to_hero_guide/gmail_agent/functions_prompt.py b/docs/zero_to_hero_guide/gmail_agent/functions_prompt.py index 11e3abc3a..e9cbfb88a 100644 --- a/docs/zero_to_hero_guide/gmail_agent/functions_prompt.py +++ b/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): diff --git a/docs/zero_to_hero_guide/gmail_agent/gmagent.py b/docs/zero_to_hero_guide/gmail_agent/gmagent.py index 4184ae1a4..45ab78a05 100644 --- a/docs/zero_to_hero_guide/gmail_agent/gmagent.py +++ b/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) diff --git a/docs/zero_to_hero_guide/gmail_agent/main.py b/docs/zero_to_hero_guide/gmail_agent/main.py index f3d557b67..468a42b58 100644 --- a/docs/zero_to_hero_guide/gmail_agent/main.py +++ b/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()