main.py

import os 
import sys
import config
from dotenv import load_dotenv
from google import genai
from google.genai import types

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

from functions.call_function import call_function

def main():
    
    
    #@ Check if the correct number of arguments is provided when running the file
    if len(sys.argv) < 2:
        print("Usage: python main.py <prompt goes here>")
        sys.exit(1)
        
    #@ Get the prompt from the command line argument
    user_prompt = sys.argv[1]
    
    #@ Getting our system prompt
    system_prompt = config.SYSTEM_PROMPT
    
    ## Check for flags
    verbose = False
    
    if len(sys.argv) > 2:
        if sys.argv[2] == "--verbose":
            verbose = True
            
            
            
    #---
    
    #! Getting access to the API
    load_dotenv()
    api_key = os.environ.get("GEMINI_API_KEY")
    client = genai.Client(api_key=api_key)
    

    #@ Setting up message history along with roles
    messages = [
        types.Content(role="user", parts=[types.Part(text=user_prompt)])
    ]
    
    #@ create a list of available functions to our LLM
    available_functions = types.Tool(
        function_declarations=[
            schema_get_files_info,
            schema_get_file_content,
            schema_write_file,
            schema_run_python_file
        ]
    )
    
        
    # ---

    #$ Looping so that we can talk to the model
    max_iterations = config.MAX_ITERATIONS
    
    
    for i in range(max_iterations):
        try:
            # ! Getting a response from the API -> CALL THE AI
            response = client.models.generate_content(
                model="gemini-2.0-flash-001", 
                contents=messages, 
                config=types.GenerateContentConfig(system_instruction=system_prompt, tools=[available_functions]),
            )
        
            prompt_tokens = response.usage_metadata.prompt_token_count
            response_tokens = response.usage_metadata.candidates_token_count
            
            # ---
            # @ Checking response's .candidates property
            # -> Candidates are like the AI's thinking process (the possible responses that it has generated)
            #  within the candidates are the function calls -> so we know what functions that the AI wants to call
            candidates = response.candidates
            
            if candidates is not None:
                for item in candidates:
                    candidate_content = item.content
                    if candidate_content is not None:
                        messages.append(candidate_content)
                        
                        
            #@ printing and executing function calls (specified by the LLM)
            function_call = response.function_calls
            if function_call is not None:
                for item in function_call:            
                    try:
                        function_call_result = call_function(item, verbose)
                        function_response = function_call_result.parts[0].function_response.response
                        
                        #@ Add reply to messages list
                        messages.append(types.Content(role="user", parts=function_call_result.parts))
                        print("=================")
                        print("=================")
                        
                        print(f"Added function result to messages: {function_call_result.parts}")

                        if not function_response:
                            raise Exception("Fatal function call exception. See function_result for details.")
                        
                        if verbose:
                            print(f"-> {function_response}")

                    except Exception as e:
                        
                        function_call_error =f'Error occurred while calling function: {item.name}. {e}'
                        messages.append(types.Content(role="user", parts=[types.Part(text=function_call_error)]))
                        continue;
                    
            if function_call is None and response.text:
                #@ Print the response
                print("=================")
                print("FINAL RESPONSE FROM LLM:")
                print("=================")
                print(response.text)
                break
                    
        except Exception as e:
            print(f'Error occurred while talking to the LLM. {e}')
                
        
    # ---
    
    
    if verbose:
        print(f"User prompt: {user_prompt}");
        print(f"Prompt tokens: {prompt_tokens}")
        print(f"Response tokens: {response_tokens}")
        

                
                
    
    
   


if __name__ == "__main__":
    main()