main.py

import os
import argparse
from dotenv import load_dotenv
from google import genai
from google.genai import types

from functions.get_files_info import schema_get_files_info, get_files_info
from functions.get_file_content import schema_get_file_content, get_file_content
from functions.write_file import schema_write_file, write_file
from functions.run_python_file import schema_run_python_file, run_python_file


def call_function(
    function_call_part: types.FunctionCall, verbose=False
) -> types.Content:
    if verbose:
        print(f"Calling function: {function_call_part.name}({function_call_part.args})")
    else:
        print(f" - Calling function: {function_call_part.name}")

    lut = {
        "get_file_content": get_file_content,
        "get_files_info": get_files_info,
        "write_file": write_file,
        "run_python_file": run_python_file,
    }

    function_name = function_call_part.name
    assert function_name is not None
    if function_call_part.name not in lut:
        return types.Content(
            role="tool",
            parts=[
                types.Part.from_function_response(
                    name=function_name,
                    response={"error": f"Unknown function: {function_name}"},
                )
            ],
        )
    if function_call_part.args is None:
        return types.Content(
            role="tool",
            parts=[
                types.Part.from_function_response(
                    name=function_name,
                    response={"error": f"Missing args for: {function_name}"},
                )
            ],
        )

    result = lut[function_name](
        working_directory="./calculator/", **(function_call_part.args)
    )
    return types.Content(
        role="tool",
        parts=[
            types.Part.from_function_response(
                name=function_name,
                response={"result": result},
            )
        ],
    )


def generate_content(client, messages, content_config, verbose):
    response = client.models.generate_content(
        model="gemini-2.0-flash-001", contents=messages, config=content_config
    )
    function_results = []
    if response.candidates:
        for candidate in response.candidates:
            content = candidate.content
            if content:
                messages.append(
                    types.Content(role=content.role, parts=content.parts)
                )

    if response.usage_metadata is None:
        raise Exception("Did not get response usage meta data.")
    if response.function_calls:
        for function_call_part in response.function_calls:
            result = call_function(function_call_part, verbose)
            if not result.parts or not result.parts[0].function_response:
                raise Exception(
                    f"Error: expected a result from the function {function_call_part.name}"
                )

            function_results.append(result.parts[0])
            if verbose:
                print(f"-> {result.parts[0].function_response.response}")

            messages.append(types.Content(role="user", parts=function_results))
    if verbose:
        print(f"User prompt: {content_config.system_instruction}")
        print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
        print(f"Response tokens: {response.usage_metadata.candidates_token_count}")

    if not response.function_calls:
        print(f"Final response:\n{response.text}")
        return True
    
    return False


def main():
    parser = argparse.ArgumentParser(prog="agent", description="My own custom agent.")
    parser.add_argument("prompt")
    parser.add_argument("-v", "--verbose", action="store_true")

    args = parser.parse_args()
    user_prompt = args.prompt
    verbose = args.verbose

    messages = [
        types.Content(role="user", parts=[types.Part(text=user_prompt)]),
    ]
    system_prompt = """
You are a helpful AI coding agent.

When a user asks a question or makes a request, make a function call plan. You can perform the following operations:

- List files and directories. Don't ask for confirmation if this is the function you think you need to execute.
- Read file contents. Don't ask for confirmation if this is the function you think you need to execute.
- Execute Python files with optional arguments
- Write or overwrite files

All paths you provide should be relative to the working directory. You do not need to specify the working directory in your function calls as it is automatically injected for security reasons.
"""
    available_functions = types.Tool(
        function_declarations=[
            schema_get_files_info,
            schema_get_file_content,
            schema_write_file,
            schema_run_python_file,
        ]
    )

    load_dotenv()
    api_key = os.environ.get("GEMINI_API_KEY")
    client = genai.Client(api_key=api_key)
    content_config = types.GenerateContentConfig(
        tools=[available_functions], system_instruction=system_prompt
    )

    for _ in range(20):
        try:
            done = generate_content(client, messages, content_config, verbose)
            if done:
                break
        except Exception as e:
            print(f"Feedback loop iteration got exception: {e}")


if __name__ == "__main__":
    main()