main.py

import os
from urllib import response
from dotenv import load_dotenv
from google import genai
from google.genai import types
import argparse
import config
from schemas import available_functions
from call_function import call_function

load_dotenv()
api_key = os.environ.get("GEMINI_API_KEY")
system_prompt = config.system_prompt
ai_model = config.ai_model

# If the environment variable wasn't found – i.e., if api_key is None – raise a RuntimeError with a helpful message.
if api_key is None:
    raise RuntimeError("GEMINI_API_KEY environment variable not set.")

client = genai.Client(api_key=api_key)

# test the client by making a simple request to the Gemini API
def chat_with_gemini(chat_messages): 
    response = client.models.generate_content(
    model=ai_model, 
    contents=chat_messages,
    config=types.GenerateContentConfig(system_instruction=system_prompt,tools=[available_functions]),
)   
    return response

def usage_print(response):
    if response.usage_metadata is None:
        print("No usage metadata available.")
        return
    print("\nUSAGE:")
    print("Prompt token count:", response.usage_metadata.prompt_token_count)
    print("Candidates token count:", response.usage_metadata.candidates_token_count)
    print("---------"*10)

def agent_loop(user_prompt, verbose=False, max_iterations=20):
    messages = [
        types.Content(
            role="user",
            parts=[types.Part(text=user_prompt)]
        )
    ]

    function_results = []

    for iteration in range(max_iterations):
        response = chat_with_gemini(messages)

        if not response:
            print("Error: No response from LLM")
            break

        if response.candidates:
            for candidate in response.candidates:
                messages.append(candidate.content)

        if response.function_calls:
            for function_call in response.function_calls:
                if verbose:
                    print(f" - Calling function: {function_call.name}")
                result_content = call_function(function_call, verbose=verbose)
                function_results.append(result_content)
        else:
            print("\nFinal response:")
            print(response.text)
            return

        # --- FIXED: correctly inject function results ---
        if function_results:
            parts_to_add = []
            for result_content in function_results:
                parts_to_add.extend(result_content.parts)
            messages.append(types.Content(role="user", parts=parts_to_add))
            function_results = []

def main():
    parser = argparse.ArgumentParser(description="Chatbot Agent")
    parser.add_argument("user_prompt", type=str, help="User prompt")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
    args = parser.parse_args()

    MAX_ITERATIONS = 20  # safety stop

    # Run the agent loop with the user's prompt
    agent_loop(user_prompt=args.user_prompt, verbose=args.verbose, max_iterations=MAX_ITERATIONS)

if __name__ == "__main__":
    main()