Demo app - email assistant

.gitignore

@@ -239,3 +239,8 @@ zipalign*
 *.sln
 *.sw?
 
+# demo files
+demo/email_assistant/credentials.json
+demo/email_assistant/token.json
+
+yarn.lock

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.

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.

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"]}}

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"] = ["[email protected]", "[email protected]"]
+    message["From"] = "[email protected]"
+    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)

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,
+}
+
+
+
+

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}")
+
+
+