tools.py

import os
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.date import DateTrigger
from apscheduler.triggers.cron import CronTrigger
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
import google.generativeai as genai
import uuid
from typing import Dict, List
from datetime import datetime
from prompts import execute_task_prompt
from utils import debug, debug_function_call, max_retries, model_name, send_email

from duckduckgo_search import DDGS
from trafilatura import fetch_url, extract


# Initialize BackgroundScheduler
scheduler = BackgroundScheduler(
    {
        "apscheduler.jobstores.default": {
            "type": "sqlalchemy",
            "url": "sqlite:///jobs.sqlite",
        }
    }
)
scheduler.start()

notes = []


# functions to browse web
def web_search(query: str, max_results: int) -> List:
    """Use this function to perform web searches.

    You should use this function only to perform regular web searches,
        and use `news_search` instead to search for news stories because this function
        does not return news results.

    Args:
        query (str): The query to search for.
        max_results (int): Maximum results to return. Defaults to 10. You can set it to a lower value but not to a higher value.

    Returns:
        list - A list of dictionaries where each dictionary have title, href and body keys. Pass `href` to `get_content_from_url` to read the full main content at that URL.
    """
    max_results = 10 if max_results > 10 else max_results
    for _ in range(max_retries):
        try:
            results = DDGS().text(query, max_results=max_results)
            for result in results:
                result["excerpt"] = result.pop("body")

            return results
        except Exception:
            pass

    # if all attempts fail, return an empty list
    return []


def news_search(query: str, max_results: int) -> List:
    """Use this function to search for recent news stories.

    Always use this function to search for news stories about a given topic and not `web_search`.

    Args:
        query (str): The query to search for.
        max_results (int): Maximum results to return. Defaults to 10. You can set it to a lower value but not to a higher value.

    Returns:
        list - A list of dictionaries where each dictionary have title, href and body keys. Pass `href` to `get_content_from_url` to read the full main content at that URL.
    """
    max_results = 10 if max_results > 10 else max_results
    for _ in range(max_retries):
        try:
            return DDGS().news(query, max_results=max_results)
        except Exception:
            pass

    # if all attempts fail, return an empty list
    return []


def read_more(url: str) -> str:
    """Use this function to get markdown-formatted content from a given URL.

    It is useful to read the content of a url after a web search result because search results are presented only with short excerpts.
        After reading a web page with this function, always remember to organize your ideas with `scratchpad` before continuing with the next steps.

    Args:
        url (str): URL to fetch content from.

    Returns:
        str - markdown-formatted content with metadata
    """
    downloaded = fetch_url(url)
    return extract(downloaded, output_format="markdown", with_metadata=True)


def scratchpad(note: str) -> str:
    """Use this function to take notes for yourself.

    When working on research tasks, always remember to take temporary notes about your important findings or next steps of the task that will help you deliver the best result for this task.

    Args:
        note (str): A free-form note that you want to save.

    Returns:
        str - A short message about success or failure.
    """
    notes.append(note)
    return "Ok, noted. Continue"


def create_document(name: str, content: str) -> str:
    """Use this function to save your deliverables in  a document.

    Remember that all the notes you put in the scratchpad will be lost after the task is complete, and you can deliver only the document you create with this function.

    Args:
        name(str): A short yet informatory title for this document.
        content (str): The full content you want to put in this document.

    Return:
        str - A message indicating success or failure.
    """
    os.makedirs("docs", exist_ok=True)
    with open(f"./docs/{name}.md", "w", encoding="utf8") as f:
        f.write(content)


# Function to get current date and time
def get_date_and_time() -> Dict:
    now = datetime.now()
    return {"date": now.strftime("%Y-%m-%d"), "time": now.strftime("%H:%M:%S")}


# Function to set a one-time reminder
def set_one_time_reminder(
    reminder_name: str,
    firing_time: str,
    task_instructions: str,
) -> str:
    """Use this function to set one-time reminders, e.g., that will executed only once at a certain time.

    Args:
        reminder_name: (str): A short descriptive name of the task, briefly summarizing what it is about.
        firing_time (str): A future date and time in isoformat indicating when this reminder should be fired. You will get a message at this time.
            It should be always a future datetime, so call `get_date_and_time` first to learn the current date and time.
        task_instructions(str): a free-form text with any relevant instructions that will help you remember and perform the task at `firing_time`.
            Use it wisely and put any useful information to it because you will have access to only this piece of information
            in addition to the tools to perform the task.
            This information can include notes for yourself regarding what tools to use in order to perform this task.
            It should include only the instructions, not the result of the task itself.
            Remember that you will always execute the task at `firing_time` if you set a reminder.

    Returns
        str - a status message indicating success or failure. You should retry in case of a failure.
    """
    reminder_id = str(uuid.uuid4())
    scheduled_at = datetime.now().isoformat()
    try:
        trigger = DateTrigger(run_date=datetime.fromisoformat(firing_time))
        scheduler.add_job(
            execute_reminder,
            trigger=trigger,
            args=[reminder_id, reminder_name, scheduled_at, task_instructions],
            id=reminder_id,
            name=reminder_name,
            replace_existing=True,
        )
        return f"Reminder `{reminder_name}` successfully scheduled"
    except Exception as e:
        return f"could not schedule reminder `{reminder_name}` due to the following error: {e}"


# Function to set a recurring reminder
def set_recurring_reminder(
    reminder_name: str,
    crontab_expression: str,
    task_instructions: str,
) -> str:
    """Use this function to set recurring reminders, e.g., that will be repeatedly executed according to a crontab expression.

    Args:
        reminder_name: (str): A short descriptive name of the task, briefly summarizing what it is about.
        crontab_expression (str): A standard crontab expression indicating when this reminder should be fired. You will get a message at any matching time according to this expression.
        task_instructions(str): a free-form text with any relevant instructions that will help you remember and perform the task at any matching time of `crontab_expression`.
            Use it wisely and put any useful information to it because you will have access to only this piece of information
            in addition to the tools to perform the task.
            This information can include notes for yourself regarding what tools to use in order to perform this task.
            It should include only the instructions, not the result of the task itself.
            Remember that you will always execute the task at any matching time of `crontab_expression` if you set a reminder.


    Returns
        str - a status message indicating success or failure. You should retry in case of a failure.
    """
    reminder_id = str(uuid.uuid4())
    scheduled_at = datetime.now().isoformat()
    try:
        trigger = CronTrigger.from_crontab(expr=crontab_expression)
        scheduler.add_job(
            execute_reminder,
            trigger=trigger,
            args=[reminder_id, reminder_name, scheduled_at, task_instructions],
            id=reminder_id,
            name=reminder_name,
            replace_existing=True,
        )
        return f"Reminder `{reminder_name}` successfully scheduled"
    except Exception as e:
        return f"could not schedule reminder `{reminder_name}` due to the following error: {e}"


def send_notification(title: str, text: str) -> str:
    """Use this function to send a notification to your manager.

    You can use it to remind them about something or submit the result of a task or a deliverable, e.g., a research report or a summary etc.

    Args:
        title: A short title for this notification.
        text (str): Full text that will be sent to your manager.

    Returns:
        str - indicating success or failure. Retry in case of a failure because your manager never wants to miss a notification.
    """
    return send_email(title, text)


execute_task_functions = {
    "create_document": create_document,
    "get_date_and_time": get_date_and_time,
    "news_search": news_search,
    "read_more": read_more,
    "scratchpad": scratchpad,
    "send_notification": send_notification,
    "web_search": web_search,
}


# Function to execute a reminder task
def execute_reminder(
    reminder_id: str, reminder_name: str, scheduled_at: str, task_instructions: str
):
    firing_time = datetime.now().isoformat()
    compiled_prompt = execute_task_prompt.format(
        reminder_name=reminder_name,
        scheduled_at=scheduled_at,
        firing_time=firing_time,
        task_instructions=task_instructions,
    )
    model = genai.GenerativeModel(
        model_name=model_name, tools=execute_task_functions.values()
    )
    try:
        for _ in range(max_retries):
            chat = model.start_chat(enable_automatic_function_calling=True)
            response = chat.send_message(compiled_prompt)
            if debug:
                for content in chat.history:
                    for part in content.parts:
                        if fn := part.function_call:
                            debug_function_call(fn)

            print(response.text)
            break
    except genai.types.generation_types.StopCandidateException:
        print("retrying...")