working version of using llama stack with multi-turn Qs

This commit is contained in:
Jeff Tang 2024-12-09 19:05:05 -08:00
parent 7ff39d1857
commit 52937f9c3a
3 changed files with 81 additions and 39 deletions

View file

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

View file

@ -12,10 +12,10 @@ import pytz
import base64 import base64
import pickle import pickle
from datetime import datetime, timezone from datetime import datetime, timezone
import json
import ollama import ollama
from pypdf import PdfReader from pypdf import PdfReader
from pathlib import Path from pathlib import Path
from shared import memory
SCOPES = ['https://www.googleapis.com/auth/gmail.readonly', 'https://www.googleapis.com/auth/gmail.compose'] 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 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': if detail == 'body':
return get_email_body(which) return get_email_body(message_id)
elif detail == 'attachment': elif detail == 'attachment':
return get_email_attachments(which) return get_email_attachments(which)

View file

@ -1,8 +1,8 @@
import argparse import argparse
import gmagent import gmagent
import asyncio import asyncio
from gmagent import * import json
from functions_prompt import * #system_prompt from functions_prompt import *
from llama_stack_client import LlamaStackClient from llama_stack_client import LlamaStackClient
from llama_stack_client.lib.agents.agent import Agent from llama_stack_client.lib.agents.agent import Agent
@ -11,6 +11,8 @@ from llama_stack_client.types.agent_create_params import (
AgentConfig, AgentConfig,
) )
from shared import memory
LLAMA_STACK_API_TOGETHER_URL="https://llama-stack.together.ai" LLAMA_STACK_API_TOGETHER_URL="https://llama-stack.together.ai"
LLAMA31_8B_INSTRUCT = "Llama3.1-8B-Instruct" 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.""" """Create an agent with gmail tool capabilities."""
listEmailsTool = ListEmailsTool() listEmailsTool = ListEmailsTool()
getEmailTool = GetEmailTool() getEmailDetailTool = GetEmailDetailTool()
sendEmailTool = SendEmailTool() sendEmailTool = SendEmailTool()
getPDFSummaryTool = GetPDFSummaryTool() getPDFSummaryTool = GetPDFSummaryTool()
createDraftTool = CreateDraftTool() createDraftTool = CreateDraftTool()
@ -34,7 +36,7 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
}, },
tools=[ tools=[
listEmailsTool.get_tool_definition(), listEmailsTool.get_tool_definition(),
getEmailTool.get_tool_definition(), getEmailDetailTool.get_tool_definition(),
sendEmailTool.get_tool_definition(), sendEmailTool.get_tool_definition(),
getPDFSummaryTool.get_tool_definition(), getPDFSummaryTool.get_tool_definition(),
createDraftTool.get_tool_definition(), createDraftTool.get_tool_definition(),
@ -52,7 +54,7 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
client=client, client=client,
agent_config=agent_config, agent_config=agent_config,
custom_tools=[listEmailsTool, custom_tools=[listEmailsTool,
getEmailTool, getEmailDetailTool,
sendEmailTool, sendEmailTool,
getPDFSummaryTool, getPDFSummaryTool,
createDraftTool, createDraftTool,
@ -62,9 +64,6 @@ async def create_gmail_agent(client: LlamaStackClient) -> Agent:
return agent return agent
async def main(): async def main():
parser = argparse.ArgumentParser(description="Set email address") parser = argparse.ArgumentParser(description="Set email address")
parser.add_argument("--gmail", type=str, required=True, help="Your Gmail 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.") 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_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: while True:
ask = input(agent_response) ask = input(agent_response)
@ -82,30 +83,48 @@ async def main():
print(llama31("bye")) print(llama31("bye"))
break break
print("\n-------------------------\nCalling Llama...") print("\n-------------------------\nCalling Llama...")
# agent(ask)
# agent_response = "Your ask: "
client = LlamaStackClient(base_url=LLAMA_STACK_API_TOGETHER_URL) client = LlamaStackClient(base_url=LLAMA_STACK_API_TOGETHER_URL)
agent = await create_gmail_agent(client) agent = await create_gmail_agent(client)
session_id = agent.create_session("email-session") session_id = agent.create_session("email-session")
queries = [ response = agent.create_turn(
"do i have any emails with attachments?", messages=[{"role": "user", "content": ask}],
"what's the content of the email from LangSmith", session_id=session_id,
] )
for query in queries: async for log in EventLogger().log(response):
print(f"\nQuery: {query}") if log.role == "CustomTool":
print("-" * 50) 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( elif tool_name == "get_email_detail":
messages=[{"role": "user", "content": query}], output = result
session_id=session_id,
) 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()