diff --git a/calculator/main.py b/calculator/main.py index beb4cec..737148e 100644 --- a/calculator/main.py +++ b/calculator/main.py @@ -2,7 +2,6 @@ import sys from pkg.calculator import Calculator from pkg.render import render - def main(): calculator = Calculator() if len(sys.argv) <= 1: diff --git a/functions/get_file_content.py b/functions/get_file_content.py index 1fc90c0..1f3919a 100644 --- a/functions/get_file_content.py +++ b/functions/get_file_content.py @@ -1,5 +1,21 @@ import os from functions.config import MAX_FILE_CONTENT_CHARS +from google.genai import types + +schema_get_file_content = types.FunctionDeclaration( + name="get_file_content", + description="Reads and returns the contents of a file, truncated if too large, constrained to the working directory.", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "file_path": types.Schema( + type=types.Type.STRING, + description="Path to the file, relative to the working directory.", + ), + }, + required=["file_path"], + ), +) def get_file_content(working_directory, file_path): diff --git a/functions/get_files_info.py b/functions/get_files_info.py index 21f9523..da9dc05 100644 --- a/functions/get_files_info.py +++ b/functions/get_files_info.py @@ -1,4 +1,19 @@ import os +from google.genai import types + +schema_get_files_info = types.FunctionDeclaration( + name="get_files_info", + description="Lists files in the specified directory along with their sizes, constrained to the working directory.", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "directory": types.Schema( + type=types.Type.STRING, + description="The directory to list files from, relative to the working directory. If not provided, lists files in the working directory itself.", + ), + }, + ), +) def get_files_info(working_directory, directory="."): try: diff --git a/functions/run_python_file.py b/functions/run_python_file.py index f85316e..678834f 100644 --- a/functions/run_python_file.py +++ b/functions/run_python_file.py @@ -1,6 +1,26 @@ import os import subprocess +from google.genai import types +schema_run_python_file = types.FunctionDeclaration( + name="run_python_file", + description="Executes a Python file with optional arguments, constrained to the working directory.", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "file_path": types.Schema( + type=types.Type.STRING, + description="Path to the Python file, relative to the working directory.", + ), + "args": types.Schema( + type=types.Type.ARRAY, + description="List of arguments to pass to the Python script.", + items=types.Schema(type=types.Type.STRING), + ), + }, + required=["file_path"], + ), +) def run_python_file(working_directory, file_path, args=[]): try: diff --git a/functions/write_file.py b/functions/write_file.py index be8c43f..5972353 100644 --- a/functions/write_file.py +++ b/functions/write_file.py @@ -1,4 +1,24 @@ import os +from google.genai import types + +schema_write_file = types.FunctionDeclaration( + name="write_file", + description="Writes or overwrites file contents within the working directory.", + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "file_path": types.Schema( + type=types.Type.STRING, + description="Path to the file relative to the working directory.", + ), + "content": types.Schema( + type=types.Type.STRING, + description="The text content to write into the file.", + ), + }, + required=["file_path", "content"], + ), +) def write_file(working_directory, file_path, content): diff --git a/main.py b/main.py index d00f4de..88ad19f 100644 --- a/main.py +++ b/main.py @@ -3,6 +3,39 @@ import sys 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 + +available_functions = types.Tool( + function_declarations=[ + schema_get_files_info, + schema_get_file_content, + schema_run_python_file, + schema_write_file, + ] +) + +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. +""" + +config=types.GenerateContentConfig( + tools=[available_functions], system_instruction=system_prompt +) messages = [] @@ -35,15 +68,22 @@ def main(): response = client.models.generate_content( model="gemini-2.0-flash-001", - contents=messages + contents=messages, + config=config, ) - reply_text = response.text + if response.candidates[0].content.parts: + for part in response.candidates[0].content.parts: + if part.function_call: + function_call_part = part.function_call + print(f"Calling function: {function_call_part.name}({function_call_part.args})") + elif part.text: + reply_text = part.text - # Add AI response to history - add_message("model", reply_text) + # Add AI response to history + add_message("model", reply_text) - print(reply_text) + print(reply_text) if verbose: print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")