diff --git a/.gitignore b/.gitignore index ceb23c3..ded2946 100644 --- a/.gitignore +++ b/.gitignore @@ -239,3 +239,8 @@ zipalign* *.sln *.sw? +# demo files +demo/email_assistant/credentials.json +demo/email_assistant/token.json + +yarn.lock \ No newline at end of file diff --git a/README.md b/README.md index 2af096a..2992195 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,8 @@ All models are enhanced from the top open-source LLMs with further post-training Try out the models immediately without downloading anything in Our [Huggingface Spaces]([https://huggingface.co/spaces/sanjay920/rubra-v0.1-dev](https://huggingface.co/spaces/sanjay920/rubra-v0.1-function-calling))! It's free and requires no login. +For more examples, please check out the `demo` directory. + ## Run Rubra Models Locally Check out our [documentation](https://docs.rubra.ai/category/serving--inferencing) to learn how to run Rubra models locally. diff --git a/demo/email_assistant/README.md b/demo/email_assistant/README.md new file mode 100644 index 0000000..3dc2ab9 --- /dev/null +++ b/demo/email_assistant/README.md @@ -0,0 +1,45 @@ +## Rubra model demo + +This demo will walk you through an example that how you can connect rubra tool-call model to your gmail mailbox, and let the ai assistant helps you take care of your emails. + +*In this demo, the assistant is granted privileges only to read your emails and change the status of an email from unread to read.* + +### Prerequisites: +- Python 3.10.7 or greater, with the pip package management tool +- A Google Cloud project. +- Your Google account with Gmail enabled. + +### Get Started + +**1.Start a Rubra Model server:** +Use either [tools.cpp](https://github.com/rubra-ai/tools.cpp?tab=readme-ov-file#toolscpp-quickstart) or [vLLM](https://github.com/rubra-ai/vllm?tab=readme-ov-file#rubra-vllm-quickstart) to serve a Rubra model. + +**2.Enable Gmail API and setup authentication:** +A few things to config to allow the AI assistant to connect to your gmail emails thru Gmail API: +- In the Google Cloud console, [enable the Gmail API](https://console.cloud.google.com/flows/enableapi?apiid=gmail.googleapis.com). +- [Configure the OAuth consent screen](https://developers.google.com/gmail/api/quickstart/python#configure_the_oauth_consent_screen): For User type select Internal, if you can't then simply select external. +- [Authorize credentials for a desktop application](https://developers.google.com/gmail/api/quickstart/python#authorize_credentials_for_a_desktop_application): Don't forget to download `credentials.json` to the `demo` dir or where you'd like to run the code. + +Reference: https://developers.google.com/gmail/api/quickstart/python#set_up_your_environment + +**3.Pip install and Run the python script:** +```python +pip install -r requirements.txt +``` +and then: +```python +python main_email_assistant.py +``` + +The user prompt in this script: +``` +Process my last 5 emails. get the label for all of them, then change the emails with a `daily` label to `read` status. +``` +If everything goes well, the AI assistant will look at the latest 5 emails and mark some of them as `read`. + +### What's next? +In the demo, the assistant is granted privileges only to: +- list and read emails +- change the status of emails from `unread` to `read`. + +You can definitely enhance its capabilities by introducing more tools/functions, such as moving emails to different folders/inboxes, drafting and sending emails, etc. \ No newline at end of file diff --git a/demo/email_assistant/credentials_example.json b/demo/email_assistant/credentials_example.json new file mode 100644 index 0000000..a9132d1 --- /dev/null +++ b/demo/email_assistant/credentials_example.json @@ -0,0 +1 @@ +{"installed":{"client_id":"xxx.apps.googleusercontent.com","project_id":"xxx","auth_uri":"https://accounts.google.com/o/oauth2/auth","token_uri":"https://oauth2.googleapis.com/token","auth_provider_x509_cert_url":"https://www.googleapis.com/oauth2/v1/certs","client_secret":"xxx","redirect_uris":["http://localhost"]}} \ No newline at end of file diff --git a/demo/email_assistant/email_operations.py b/demo/email_assistant/email_operations.py new file mode 100644 index 0000000..e30b4e5 --- /dev/null +++ b/demo/email_assistant/email_operations.py @@ -0,0 +1,172 @@ +''' +https://developers.google.com/gmail/api/quickstart/python +''' + +import os.path +import base64 +from email.message import EmailMessage +from google.auth.transport.requests import Request +from google.oauth2.credentials import Credentials +from google_auth_oauthlib.flow import InstalledAppFlow +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError + +# If modifying these scopes, delete the file token.json. +SCOPES = ["https://mail.google.com/"] +# SCOPES = ["https://www.googleapis.com/auth/gmail.compose", "https://www.googleapis.com/auth/gmail.readonly"] + +def auth(): + """Shows basic usage of the Gmail API. + Lists the user's Gmail labels. + """ + creds = None + # The file token.json stores the user's access and refresh tokens, and is + # created automatically when the authorization flow completes for the first + # time. + if os.path.exists("token.json"): + creds = Credentials.from_authorized_user_file("token.json", SCOPES) + # If there are no (valid) credentials available, let the user log in. + if not creds or not creds.valid: + if creds and creds.expired and creds.refresh_token: + creds.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file( + "credentials.json", SCOPES + ) + creds = flow.run_local_server(port=0) + # Save the credentials for the next run + with open("token.json", "w") as token: + token.write(creds.to_json()) + return creds + +def decode_base64(data): + decoded_bytes = base64.urlsafe_b64decode(data) + decoded_str = decoded_bytes.decode('utf-8') + return decoded_str + + +import google.auth +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError + +def gmail_send_message(): + """Create and insert a draft email. + Print the returned draft's message and id. + Returns: Draft object, including draft id and message meta data. + + Load pre-authorized user credentials from the environment. + TODO(developer) - See https://developers.google.com/identity + for guides on implementing OAuth2 for the application. + """ + + try: + # create gmail api client + service = build("gmail", "v1", credentials=auth()) + + message = EmailMessage() + + message.set_content("This is automated draft mail") + + message["To"] = ["yingbei@acorn.io", "tybalex@gmail.com"] + message["From"] = "tybalex@gmail.com" + message["Subject"] = "Automated draft" + + # encoded message + encoded_message = base64.urlsafe_b64encode(message.as_bytes()).decode() + + create_message = {"raw": encoded_message} + # pylint: disable=E1101 + send_message = ( + service.users() + .messages() + .send(userId="me", body=create_message) + .execute() + ) + print(f'Message Id: {send_message["id"]}') + except HttpError as error: + print(f"An error occurred: {error}") + send_message = None + return send_message + + +def mark_as_read( email_id): + service = build("gmail", "v1", credentials=auth()) + service.users().messages().modify(userId='me', id=email_id, body={'removeLabelIds': ['UNREAD']}).execute() + + + +def read_message(email_id): + service = build("gmail", "v1", credentials=auth()) + msg = service.users().messages().get(userId='me', id=email_id).execute() + + # Extract the parts + email_data = msg + headers = email_data['payload']['headers'] + header_dict = {header['name']: header['value'] for header in headers} + + # Print the extracted information + title = header_dict.get('Subject', 'No Subject') + sender = header_dict.get('From', 'No Sender') + receiver = header_dict.get('To', 'No Receiver') + date = header_dict.get("Date", "No Date Received") + + content_text = "" + try: + if "parts" not in email_data['payload']: + parts = [] + else: + parts = email_data['payload']['parts'] + decoded_parts = {} + + for part in parts: + mime_type = part['mimeType'] + encoded_data = part['body']['data'] + decoded_content = decode_base64(encoded_data) + decoded_parts[mime_type] = decoded_content + + # Extract necessary information + content_text = decoded_parts.get('text/plain', 'No Plain Text Content') + except Exception as e: + print(e) + + return { + "id" : email_id, + "title": title, + "sender": sender, + "receiver": receiver, + "date": date, + "content_text": content_text + } + + +def list_messages(n=5, date = None): + try: + # create gmail api client + service = build("gmail", "v1", credentials=auth()) + results = ( + service.users() + .messages() + .list(userId="me", labelIds=["UNREAD"]) + .execute() + ) + messages = results.get('messages',[]) + res = [] + if not messages: + print('No new messages.') + else: + for i, message in enumerate(messages): + if i >= n: + break + res.append(message["id"]) + + print(res) + return res + + + except HttpError as error: + print(f"An error occurred: {error}") + +if __name__ == "__main__": +# gmail_send_message() + res = list_messages() + print(res) \ No newline at end of file diff --git a/demo/email_assistant/email_tools.py b/demo/email_assistant/email_tools.py new file mode 100644 index 0000000..93463e7 --- /dev/null +++ b/demo/email_assistant/email_tools.py @@ -0,0 +1,141 @@ +from email_operations import mark_as_read, list_messages, read_message +from run_chat_completion import run_chat +import json + +default_system_prompt = "You are a helpful assistant." + +def run_agent(user_query, functions, system_prompt=default_system_prompt): + print(f"User query: {user_query}") + res, msgs = run_chat(user_query=user_query, functions=functions, system_prompt=system_prompt) + while res.message.tool_calls: + tool_calls = [] + func_output =[] + for tool_call in res.message.tool_calls: + + func_name,func_args = tool_call.function.name, tool_call.function.arguments + print(f"\n=====calling function : {func_name}, with args: {func_args}") + tool_calls.append( { + "id": tool_call.id, + "function": {"name": func_name, + "arguments": func_args}, + "type": "function", + }) + func_args = json.loads(func_args) + func_to_run = tool_call_mapping[func_name] + observation = func_to_run(**func_args) + # print(f"Observation: {observation}") + func_output.append([tool_call.id, func_name, str(observation)]) + msgs.append({"role": "assistant", "tool_calls": tool_calls}) + for id,func_name, o in func_output: + msgs.append({ + "role": "tool", + "name": func_name, + "content": o, + "tool_call_id": id + }) + res, msgs = run_chat(user_query=user_query,functions=functions, msgs=msgs) + final_res = res.message.content + return final_res + + +main_functions = [ + { + "type": "function", + "function": { + "name": "list_unread_emails", + "description": "List all unread emails in the mailbox", + "parameters": { + "type": "object", + "properties": { + "n": { + "type": "integer", + "description": "the number of emails to return, default = 5" + }, + "date": { + "type": "string", + "description": "list unread email for a specific date, in yyyy-mm-dd format, default is None. Useful when user want emails for a certain day" + }, + + }, + "required": [ + + ] + } + } + }, + { + "type": "function", + "function": { + "name": "change_email_to_read", + "description": "change the status of an email to `read`", + "parameters": { + "type": "object", + "properties": { + "email_id": { + "type": "string", + "description": "the id of the unread email to be marked as read." + } + }, + "required": [ + "email_id" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "label_email", + "description": "read an email and label it with one of the three label: work, daily, important", + "parameters": { + "type": "object", + "properties": { + "email_id": { + "type": "string", + "description": "the id of the email to process." + } + }, + "required": [ + "email_id" + ] + } + } + }, +] + + +def label_message(email_id) -> str: + """This is a rule based example to label emails. It's also possible to use LLM's help to do so. + + Args: + email_id (_type_): _description_ + Return: + one of the three label: [work, daily, important] + """ + msg_detail = read_message(email_id) + print(msg_detail["title"]) + print(msg_detail["date"]) + print(msg_detail["sender"]) + print(msg_detail["receiver"]) + + # Now do some rule based stuff or use LLM or some model to label the email + label = "daily" + if "@acorn.io" in msg_detail["sender"]: + label = "work" + # some arbitrary keyword rule based stuff + elif "REMINDER" in msg_detail["title"] or "important" in msg_detail["content_text"]: + label = "important" + + print(label) + return f"Label: {label}" + + +tool_call_mapping = { + "list_unread_emails": list_messages, + "change_email_to_read": mark_as_read, + "label_email": label_message, +} + + + + diff --git a/demo/email_assistant/main_email_assistant.py b/demo/email_assistant/main_email_assistant.py new file mode 100644 index 0000000..aefdefa --- /dev/null +++ b/demo/email_assistant/main_email_assistant.py @@ -0,0 +1,11 @@ + +from email_tools import run_agent, main_functions + +msgs = [] +user_query = "Process my last 5 emails. get the label for all of them, then change the emails with a `daily` label to `read` status." + +final_res = run_agent(user_query, main_functions) +print(f"Final AI Response: {final_res}") + + + \ No newline at end of file diff --git a/demo/email_assistant/requirements.txt b/demo/email_assistant/requirements.txt new file mode 100644 index 0000000..254454b --- /dev/null +++ b/demo/email_assistant/requirements.txt @@ -0,0 +1,4 @@ +openai +google-api-python-client +google-auth-httplib2 +google-auth-oauthlib \ No newline at end of file diff --git a/demo/email_assistant/run_chat_completion.py b/demo/email_assistant/run_chat_completion.py new file mode 100644 index 0000000..fd4d0fb --- /dev/null +++ b/demo/email_assistant/run_chat_completion.py @@ -0,0 +1,25 @@ +import openai +openai.api_key = "sk-" +openai.base_url = "http://localhost:1234/v1/" + +model = "gpt-4o" + + +def run_chat(user_query, system_prompt = "", functions=[], msgs = []): + if not msgs or len(msgs) == 0: + msgs = [{"role": "system", "content":system_prompt} ,{"role": "user", "content": user_query}] + else: + msgs.append({"role": "user", "content": user_query}) + try: + completion = openai.chat.completions.create( + model=model, + temperature=0.1, + messages=msgs, + tools=functions, + tool_choice="auto", + stream=False, + ) + res = completion.choices[0] + return res, msgs + except Exception as e: + print(f"Error : {e}") \ No newline at end of file