main.py

import os
import sys

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

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

# Map function names to actual Python implementations
FUNCTION_MAP = {
    "get_file_content": get_file_content,
    "get_files_info": get_files_info,
    "run_python_file": run_python_file,
    "write_file": write_file,
}


def call_function(function_call_part, verbose=False):
    """Call the Python function corresponding to the function_call from the model"""
    function_name = function_call_part.name
    function_args = function_call_part.args or {}

    # Always run in the calculator working directory
    function_args["working_directory"] = "./calculator"

    if verbose:
        print(f"Calling function: {function_name}({function_args})")
    else:
        print(f" - Calling function: {function_name}")

    func = FUNCTION_MAP.get(function_name)
    if not func:
        return types.Content(
            role="tool",
            parts=[
                types.Part.from_function_response(
                    name=function_name,
                    response={"error": f"Unknown function: {function_name}"},
                )
            ],
        )

    try:
        result = func(**function_args)
    except Exception as e:
        result = {"error": str(e)}

    return types.Content(
        role="tool",
        parts=[
            types.Part.from_function_response(
                name=function_name, response={"result": result}
            )
        ],
    )


# Available functions for the agent
available_functions = types.Tool(
    function_declarations=[
        schema_get_files_info,
        schema_get_file_content,
        schema_run_python_file,
        schema_write_file,
    ]
)


def main():
    load_dotenv()
    verbose = "--verbose" in sys.argv
    args = [arg for arg in sys.argv[1:] if not arg.startswith("--")]

    if not args:
        print("AI Code Assistant")
        print('\nUsage: python main.py "your prompt here" [--verbose]')
        sys.exit(1)

    user_prompt = " ".join(args)
    if verbose:
        print(f"User prompt: {user_prompt}\n")

    api_key = os.environ.get("GEMINI_API_KEY")
    client = genai.Client(api_key=api_key)

    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
    - Read file contents
    - 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.
    """

    # Initialize the conversation
    messages = [types.Content(role="user", parts=[types.Part(text=user_prompt)])]

    generate_content(client, messages, system_prompt, verbose)


def generate_content(client, messages, system_prompt, verbose):
    MAX_ITERATIONS = 20

    for iteration in range(MAX_ITERATIONS):
        try:
            response = client.models.generate_content(
                model="gemini-2.0-flash-001",
                contents=messages,
                config=types.GenerateContentConfig(
                    tools=[available_functions], system_instruction=system_prompt
                ),
            )

            if verbose:
                print(f"\n=== Iteration {iteration + 1} ===")
                print("Prompt tokens:", response.usage_metadata.prompt_token_count)
                print(
                    "Response tokens:", response.usage_metadata.candidates_token_count
                )

            done = True

            for candidate in response.candidates:
                for part in candidate.content.parts:
                    if part.function_call:
                        done = False
                        function_call_part = part.function_call
                        function_call_result = call_function(
                            function_call_part, verbose
                        )

                        # Append the function result as a user message
                        messages.append(
                            types.Content(
                                role="user",
                                parts=[
                                    types.Part.from_function_response(
                                        name=function_call_part.name,
                                        response=function_call_result.parts[
                                            0
                                        ].function_response.response,
                                    )
                                ],
                            )
                        )

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

            # Append candidate content for context
            for candidate in response.candidates:
                messages.append(candidate.content)

            # If no function calls remain, print final text and break
            if done and response.text:
                print("\nFinal Response:")
                print(response.text)
                break

        except Exception as e:
            print(f"Error during iteration {iteration + 1}: {e}")
            break


if __name__ == "__main__":
    main()