107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from typing import Any, Dict
|
|
|
|
from dotenv import load_dotenv
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
from functions.call_function import available_functions, call_function
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="AI Code Assistant")
|
|
parser.add_argument("user_prompt", type=str, help="Prompt to send to Gemini")
|
|
parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
|
|
args = parser.parse_args()
|
|
|
|
load_dotenv()
|
|
api_key = os.environ.get("GEMINI_API_KEY")
|
|
|
|
if not api_key:
|
|
raise RuntimeError("GEMINI_API_KEY environment variable not set")
|
|
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
messages: list[types.Content] = [
|
|
types.Content(role="user", parts=[types.Part(text=args.user_prompt)])
|
|
]
|
|
|
|
if args.verbose:
|
|
print(f"User prompt: {args.user_prompt}\n")
|
|
|
|
generate_content(client, messages, args.verbose)
|
|
|
|
|
|
def generate_content(
|
|
client: genai.Client, messages: list[types.Content], verbose: bool
|
|
) -> None:
|
|
system_instructions = """
|
|
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 for security reasons.
|
|
"""
|
|
|
|
for _ in range(20):
|
|
model_response = client.models.generate_content(
|
|
model="gemini-2.5-flash",
|
|
contents=messages,
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_instructions, tools=[available_functions]
|
|
),
|
|
)
|
|
|
|
if not model_response.usage_metadata:
|
|
raise RuntimeError("Gemini API response appears to be malformed")
|
|
|
|
if verbose:
|
|
print("Prompt tokens:", model_response.usage_metadata.prompt_token_count)
|
|
print(
|
|
"Response tokens:", model_response.usage_metadata.candidates_token_count
|
|
)
|
|
|
|
if model_response.candidates != None and len(model_response.candidates) != 0:
|
|
for candidate in model_response.candidates:
|
|
if candidate.content != None:
|
|
messages.append(candidate.content)
|
|
|
|
if not model_response.function_calls:
|
|
print("Final Response:")
|
|
print(model_response.text)
|
|
return
|
|
|
|
function_responses: list[types.Part] = []
|
|
|
|
for function_call in model_response.function_calls:
|
|
result = call_function(function_call, verbose)
|
|
|
|
if (
|
|
not result.parts
|
|
or not result.parts[0].function_response
|
|
or not result.parts[0].function_response.response
|
|
):
|
|
raise RuntimeError(f"Empty function response for {function_call.name}")
|
|
|
|
if verbose:
|
|
print(f"-> {result.parts[0].function_response.response}")
|
|
|
|
function_responses.append(result.parts[0])
|
|
|
|
messages.append(types.Content(role="user", parts=function_responses))
|
|
continue
|
|
|
|
print("Model took more than the max iterations to produce final response")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|