-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.py
More file actions
76 lines (65 loc) · 2.78 KB
/
main.py
File metadata and controls
76 lines (65 loc) · 2.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import os
import argparse
from dotenv import load_dotenv
from google import genai
from google.genai import types
from functions.call_function import call_function
from functions.tool_definitions import *
load_dotenv()
api_key = os.environ.get("GEMINI_API_KEY")
if api_key is None:
raise RuntimeError("Please set your GEMINI_API_KEY in .env")
parser = argparse.ArgumentParser(description="Chatbot")
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()
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 for security reasons.
"""
available_functions = types.Tool(
function_declarations=[schema_get_files_info,
schema_get_file_content,
schema_run_python_file,
schema_write_file,
],
)
client = genai.Client(api_key=api_key)
messages = [types.Content(role="user", parts=[types.Part(text=args.user_prompt)])]
job_done = False
iterations = 0
while not job_done:
iterations += 1
response = client.models.generate_content(
model='gemini-2.5-flash',
contents=messages,
config=types.GenerateContentConfig(
tools=[available_functions],
system_instruction=system_prompt
),
)
usage_metadata = response.usage_metadata
if usage_metadata is None:
raise RuntimeError("failed to get a meaninful answer from the api. Try again later")
if args.verbose:
print("User prompt: " + args.user_prompt)
print("Prompt tokens: " + str(usage_metadata.prompt_token_count))
print("Response tokens: " + str(usage_metadata.candidates_token_count))
for candidate in response.candidates:
messages.append(candidate.content)
if response.text is not None:
print("Response: " + response.text)
if response.function_calls is not None:
for function_call_part in response.function_calls:
print(f"Calling function: {function_call_part.name}({function_call_part.args})")
function_call_result = call_function(function_call_part, args.verbose)
messages.append(function_call_result)
if args.verbose:
print(f"-> {function_call_result.parts[0].function_response.response}")
if iterations > 10:
job_done = True