agent.py

import google.generativeai as genai
from typing import List
from prompts import system_prompt
from tools import (
    create_document,
    get_date_and_time,
    news_search,
    read_more,
    scratchpad,
    set_one_time_reminder,
    set_recurring_reminder,
    web_search,
)
from utils import debug, debug_function_call, max_retries, max_steps, model_name

# Define function mappings for Gemini
system_functions = {
    "create_document": create_document,
    "get_date_and_time": get_date_and_time,
    "news_search": news_search,
    "read_more": read_more,
    "scratchpad": scratchpad,
    "set_one_time_reminder": set_one_time_reminder,
    "set_recurring_reminder": set_recurring_reminder,
    "web_search": web_search,
}


def gemini_agent(instruction: str | List[genai.protos.Part]):
    model = genai.GenerativeModel(
        model_name=model_name,
        system_instruction=system_prompt,
        tools=system_functions.values(),
    )
    chat = model.start_chat()
    instruction = f"##Here's the instruction from the CTO\n{instruction}"
    for step in range(max_steps):
        try:
            for _ in range(max_retries):
                response = chat.send_message(instruction)
                fn_results = []
                num_fn_calls = len(
                    [part for part in response.parts if part.function_call]
                )
                if num_fn_calls == 0:
                    yield response.text
                    return

                for part in response.parts:
                    if fn := part.function_call:
                        if debug:
                            yield debug_function_call(fn)
                            callable_fn = system_functions[fn.name]
                            fn_result = callable_fn(**fn.args)
                            fn_results.append((fn.name, fn_result))
                    else:
                        yield part.text

                assert num_fn_calls == len(fn_results), (
                    f"Expected {num_fn_calls} function results, but got {len(fn_results)}: {fn_results}"
                )
                if num_fn_calls > 0:
                    instruction = [
                        genai.protos.Part(
                            function_response=genai.protos.FunctionResponse(
                                name=name, response={"result": result}
                            )
                        )
                        for (name, result) in fn_results
                    ]

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