Compare commits
10 Commits
b2d99d4de2
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| e889f15b06 | |||
| ee42034333 | |||
| c3e2ebada0 | |||
| 3e18a04d55 | |||
| fb424dfa19 | |||
| eaf3635284 | |||
| 9bf18ffe68 | |||
| 81d318b728 | |||
| 9682b2f4f4 | |||
| d1b535400e |
2
.gitignore
vendored
2
.gitignore
vendored
@@ -1 +1,3 @@
|
||||
.env
|
||||
__pycache__
|
||||
**/__pycache__/
|
||||
1
calculator/README.md
Normal file
1
calculator/README.md
Normal file
@@ -0,0 +1 @@
|
||||
# calculator
|
||||
1
calculator/lorem.txt
Normal file
1
calculator/lorem.txt
Normal file
@@ -0,0 +1 @@
|
||||
wait, this isn't lorem ipsum
|
||||
23
calculator/main.py
Normal file
23
calculator/main.py
Normal file
@@ -0,0 +1,23 @@
|
||||
import sys
|
||||
from pkg.calculator import Calculator
|
||||
from pkg.render import render
|
||||
|
||||
def main():
|
||||
calculator = Calculator()
|
||||
if len(sys.argv) <= 1:
|
||||
print("Calculator App")
|
||||
print('Usage: python main.py "<expression>"')
|
||||
print('Example: python main.py "3 + 5"')
|
||||
return
|
||||
|
||||
expression = " ".join(sys.argv[1:])
|
||||
try:
|
||||
result = calculator.evaluate(expression)
|
||||
to_print = render(expression, result)
|
||||
print(to_print)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
59
calculator/pkg/calculator.py
Normal file
59
calculator/pkg/calculator.py
Normal file
@@ -0,0 +1,59 @@
|
||||
class Calculator:
|
||||
def __init__(self):
|
||||
self.operators = {
|
||||
"+": lambda a, b: a + b,
|
||||
"-": lambda a, b: a - b,
|
||||
"*": lambda a, b: a * b,
|
||||
"/": lambda a, b: a / b,
|
||||
}
|
||||
self.precedence = {
|
||||
"+": 1,
|
||||
"-": 1,
|
||||
"*": 2,
|
||||
"/": 2,
|
||||
}
|
||||
|
||||
def evaluate(self, expression):
|
||||
if not expression or expression.isspace():
|
||||
return None
|
||||
tokens = expression.strip().split()
|
||||
return self._evaluate_infix(tokens)
|
||||
|
||||
def _evaluate_infix(self, tokens):
|
||||
values = []
|
||||
operators = []
|
||||
|
||||
for token in tokens:
|
||||
if token in self.operators:
|
||||
while (
|
||||
operators
|
||||
and operators[-1] in self.operators
|
||||
and self.precedence[operators[-1]] >= self.precedence[token]
|
||||
):
|
||||
self._apply_operator(operators, values)
|
||||
operators.append(token)
|
||||
else:
|
||||
try:
|
||||
values.append(float(token))
|
||||
except ValueError:
|
||||
raise ValueError(f"invalid token: {token}")
|
||||
|
||||
while operators:
|
||||
self._apply_operator(operators, values)
|
||||
|
||||
if len(values) != 1:
|
||||
raise ValueError("invalid expression")
|
||||
|
||||
return values[0]
|
||||
|
||||
def _apply_operator(self, operators, values):
|
||||
if not operators:
|
||||
return
|
||||
|
||||
operator = operators.pop()
|
||||
if len(values) < 2:
|
||||
raise ValueError(f"not enough operands for operator {operator}")
|
||||
|
||||
b = values.pop()
|
||||
a = values.pop()
|
||||
values.append(self.operators[operator](a, b))
|
||||
21
calculator/pkg/render.py
Normal file
21
calculator/pkg/render.py
Normal file
@@ -0,0 +1,21 @@
|
||||
def render(expression, result):
|
||||
if isinstance(result, float) and result.is_integer():
|
||||
result_str = str(int(result))
|
||||
else:
|
||||
result_str = str(result)
|
||||
|
||||
box_width = max(len(expression), len(result_str)) + 4
|
||||
|
||||
box = []
|
||||
box.append("┌" + "─" * box_width + "┐")
|
||||
box.append(
|
||||
"│" + " " * 2 + expression + " " * (box_width - len(expression) - 2) + "│"
|
||||
)
|
||||
box.append("│" + " " * box_width + "│")
|
||||
box.append("│" + " " * 2 + "=" + " " * (box_width - 3) + "│")
|
||||
box.append("│" + " " * box_width + "│")
|
||||
box.append(
|
||||
"│" + " " * 2 + result_str + " " * (box_width - len(result_str) - 2) + "│"
|
||||
)
|
||||
box.append("└" + "─" * box_width + "┘")
|
||||
return "\n".join(box)
|
||||
1
calculator/test_output.txt
Normal file
1
calculator/test_output.txt
Normal file
@@ -0,0 +1 @@
|
||||
hello world
|
||||
47
calculator/tests.py
Normal file
47
calculator/tests.py
Normal file
@@ -0,0 +1,47 @@
|
||||
import unittest
|
||||
from pkg.calculator import Calculator
|
||||
|
||||
|
||||
class TestCalculator(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.calculator = Calculator()
|
||||
|
||||
def test_addition(self):
|
||||
result = self.calculator.evaluate("3 + 5")
|
||||
self.assertEqual(result, 8)
|
||||
|
||||
def test_subtraction(self):
|
||||
result = self.calculator.evaluate("10 - 4")
|
||||
self.assertEqual(result, 6)
|
||||
|
||||
def test_multiplication(self):
|
||||
result = self.calculator.evaluate("3 * 4")
|
||||
self.assertEqual(result, 12)
|
||||
|
||||
def test_division(self):
|
||||
result = self.calculator.evaluate("10 / 2")
|
||||
self.assertEqual(result, 5)
|
||||
|
||||
def test_nested_expression(self):
|
||||
result = self.calculator.evaluate("3 * 4 + 5")
|
||||
self.assertEqual(result, 17)
|
||||
|
||||
def test_complex_expression(self):
|
||||
result = self.calculator.evaluate("2 * 3 - 8 / 2 + 5")
|
||||
self.assertEqual(result, 7)
|
||||
|
||||
def test_empty_expression(self):
|
||||
result = self.calculator.evaluate("")
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_invalid_operator(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.calculator.evaluate("$ 3 5")
|
||||
|
||||
def test_not_enough_operands(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.calculator.evaluate("+ 3")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
57
functions/call_function.py
Normal file
57
functions/call_function.py
Normal file
@@ -0,0 +1,57 @@
|
||||
import os
|
||||
from google.genai import types
|
||||
|
||||
from functions.get_files_info import get_files_info
|
||||
from functions.get_file_content import get_file_content
|
||||
from functions.write_file import write_file
|
||||
from functions.run_python_file import run_python_file
|
||||
|
||||
|
||||
def call_function(function_call_part, verbose=False):
|
||||
|
||||
function_name = function_call_part.name
|
||||
args = dict(function_call_part.args or {})
|
||||
|
||||
args["working_directory"] = "./calculator"
|
||||
|
||||
function_map = {
|
||||
"get_files_info": get_files_info,
|
||||
"get_file_content": get_file_content,
|
||||
"write_file": write_file,
|
||||
"run_python_file": run_python_file,
|
||||
}
|
||||
|
||||
if verbose:
|
||||
print(f"Calling function: {function_name}({args})")
|
||||
else:
|
||||
print(f" - Calling function: {function_name}")
|
||||
|
||||
if function_name not in function_map:
|
||||
return types.Content(
|
||||
role="tool",
|
||||
parts=[
|
||||
types.Part.from_function_response(
|
||||
name=function_name,
|
||||
response={"error": f"Unknown function: {function_name}"},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
function_result = function_map[function_name](**args)
|
||||
except Exception as e:
|
||||
function_result = f"Error executing function: {e}"
|
||||
|
||||
tool_response = types.Content(
|
||||
role="tool",
|
||||
parts=[
|
||||
types.Part.from_function_response(
|
||||
name=function_name, response={"result": function_result}
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
if verbose:
|
||||
print(f"-> {tool_response.parts[0].function_response.response}")
|
||||
|
||||
return tool_response
|
||||
1
functions/config.py
Normal file
1
functions/config.py
Normal file
@@ -0,0 +1 @@
|
||||
MAX_FILE_CONTENT_CHARS = 10000
|
||||
46
functions/get_file_content.py
Normal file
46
functions/get_file_content.py
Normal file
@@ -0,0 +1,46 @@
|
||||
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):
|
||||
try:
|
||||
full_path = os.path.join(working_directory, file_path)
|
||||
|
||||
abs_working_directory = os.path.abspath(working_directory)
|
||||
abs_target_path = os.path.abspath(full_path)
|
||||
|
||||
if not abs_target_path.startswith(abs_working_directory):
|
||||
return (
|
||||
f'Error: Cannot read "{file_path}" as it is outside the permitted working directory'
|
||||
)
|
||||
|
||||
if not os.path.isfile(abs_target_path):
|
||||
return f'Error: File not found or is not a regular file: "{file_path}"'
|
||||
|
||||
with open(abs_target_path, "r", encoding="utf-8", errors="replace") as f:
|
||||
content = f.read()
|
||||
|
||||
if len(content) > MAX_FILE_CONTENT_CHARS:
|
||||
truncated = content[:MAX_FILE_CONTENT_CHARS]
|
||||
return truncated + f'\n[...File "{file_path}" truncated at {MAX_FILE_CONTENT_CHARS} characters]'
|
||||
else:
|
||||
return content
|
||||
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
48
functions/get_files_info.py
Normal file
48
functions/get_files_info.py
Normal file
@@ -0,0 +1,48 @@
|
||||
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:
|
||||
full_path = os.path.join(working_directory, directory)
|
||||
|
||||
abs_working_directory = os.path.abspath(working_directory)
|
||||
abs_target_path = os.path.abspath(full_path)
|
||||
|
||||
if not abs_target_path.startswith(abs_working_directory):
|
||||
return (
|
||||
f'Error: Cannot list "{directory}" as it is outside the permitted working directory'
|
||||
)
|
||||
|
||||
if not os.path.isdir(abs_target_path):
|
||||
return f'Error: "{directory}" is not a directory'
|
||||
|
||||
results = []
|
||||
for entry in os.listdir(abs_target_path):
|
||||
entry_path = os.path.join(abs_target_path, entry)
|
||||
try:
|
||||
is_dir = os.path.isdir(entry_path)
|
||||
size = os.path.getsize(entry_path)
|
||||
results.append(
|
||||
f"- {entry}: file_size={size} bytes, is_dir={str(is_dir)}"
|
||||
)
|
||||
except Exception as e:
|
||||
results.append(f"- {entry}: Error accessing file info ({e})")
|
||||
|
||||
return "\n".join(results)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
66
functions/run_python_file.py
Normal file
66
functions/run_python_file.py
Normal file
@@ -0,0 +1,66 @@
|
||||
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:
|
||||
full_path = os.path.join(working_directory, file_path)
|
||||
abs_working_directory = os.path.abspath(working_directory)
|
||||
abs_target_path = os.path.abspath(full_path)
|
||||
|
||||
if not abs_target_path.startswith(abs_working_directory):
|
||||
return (
|
||||
f'Error: Cannot execute "{file_path}" as it is outside the permitted working directory'
|
||||
)
|
||||
|
||||
if not os.path.exists(abs_target_path):
|
||||
return f'Error: File "{file_path}" not found.'
|
||||
|
||||
if not file_path.endswith(".py"):
|
||||
return f'Error: "{file_path}" is not a Python file.'
|
||||
|
||||
completed = subprocess.run(
|
||||
["python3", abs_target_path] + args,
|
||||
cwd=abs_working_directory,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
result_parts = []
|
||||
if completed.stdout.strip():
|
||||
result_parts.append("STDOUT:\n" + completed.stdout.strip())
|
||||
if completed.stderr.strip():
|
||||
result_parts.append("STDERR:\n" + completed.stderr.strip())
|
||||
if completed.returncode != 0:
|
||||
result_parts.append(f"Process exited with code {completed.returncode}")
|
||||
|
||||
if not result_parts:
|
||||
return "No output produced."
|
||||
|
||||
return "\n".join(result_parts)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return "Error: Process timed out after 30 seconds."
|
||||
except Exception as e:
|
||||
return f"Error: executing Python file: {e}"
|
||||
47
functions/write_file.py
Normal file
47
functions/write_file.py
Normal file
@@ -0,0 +1,47 @@
|
||||
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):
|
||||
try:
|
||||
full_path = os.path.join(working_directory, file_path)
|
||||
abs_working_directory = os.path.abspath(working_directory)
|
||||
abs_target_path = os.path.abspath(full_path)
|
||||
|
||||
if not abs_target_path.startswith(abs_working_directory):
|
||||
return (
|
||||
f'Error: Cannot write to "{file_path}" as it is outside the permitted working directory'
|
||||
)
|
||||
|
||||
parent_dir = os.path.dirname(abs_target_path)
|
||||
if parent_dir and not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir, exist_ok=True)
|
||||
|
||||
with open(abs_target_path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
return (
|
||||
f'Successfully wrote to "{file_path}" ({len(content)} characters written)'
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
92
main.py
92
main.py
@@ -2,6 +2,37 @@ import os
|
||||
import sys
|
||||
from dotenv import load_dotenv
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from functions.call_function import call_function
|
||||
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.
|
||||
"""
|
||||
|
||||
load_dotenv()
|
||||
api_key = os.environ.get("GEMINI_API_KEY")
|
||||
@@ -10,20 +41,63 @@ client = genai.Client(api_key=api_key)
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 2:
|
||||
print("Error: No prompt provided.\nUsage: uv run main.py \"<your prompt here>\"")
|
||||
print("Error: No prompt provided.\nUsage: uv run main.py \"<your prompt here>\" [--verbose]")
|
||||
sys.exit(1)
|
||||
|
||||
user_prompt = sys.argv[1]
|
||||
args = sys.argv[1:]
|
||||
|
||||
response = client.models.generate_content(
|
||||
model="gemini-2.0-flash-001",
|
||||
contents=user_prompt
|
||||
)
|
||||
verbose = False
|
||||
if "--verbose" in args:
|
||||
verbose = True
|
||||
args.remove("--verbose")
|
||||
|
||||
print(response.text)
|
||||
user_prompt = " ".join(args)
|
||||
|
||||
print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
|
||||
print(f"Response tokens: {response.usage_metadata.candidates_token_count}")
|
||||
if verbose:
|
||||
print(f'User prompt: "{user_prompt}"')
|
||||
|
||||
messages = [
|
||||
types.Content(role="user", parts=[types.Part.from_text(text=user_prompt)]),
|
||||
]
|
||||
|
||||
for iteration in range(20):
|
||||
response = client.models.generate_content(
|
||||
model="gemini-2.0-flash-001",
|
||||
contents=messages,
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[available_functions],
|
||||
system_instruction=system_prompt,
|
||||
),
|
||||
)
|
||||
|
||||
if not response.candidates:
|
||||
print("No response, stopping.")
|
||||
break
|
||||
|
||||
# String representation of final text (if the model is done)
|
||||
candidate = response.candidates[0]
|
||||
messages.append(candidate.content)
|
||||
|
||||
has_function_call = False
|
||||
final_texts = []
|
||||
|
||||
for part in candidate.content.parts:
|
||||
if part.function_call:
|
||||
has_function_call = True
|
||||
function_result = call_function(part.function_call, verbose=verbose)
|
||||
messages.append(function_result)
|
||||
elif part.text:
|
||||
final_texts.append(part.text)
|
||||
|
||||
# Only finish if there was NO tool call in this iteration
|
||||
if not has_function_call and final_texts:
|
||||
print("Final response:")
|
||||
print("\n".join(final_texts))
|
||||
break
|
||||
|
||||
if verbose:
|
||||
print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
|
||||
print(f"Response tokens: {response.usage_metadata.candidates_token_count}")
|
||||
|
||||
|
||||
|
||||
|
||||
29
tests.py
Normal file
29
tests.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from functions.run_python_file import run_python_file
|
||||
|
||||
|
||||
def run_tests():
|
||||
cases = [
|
||||
("calculator", "main.py", []),
|
||||
("calculator", "main.py", ["3 + 5"]),
|
||||
("calculator", "tests.py", []),
|
||||
("calculator", "../main.py", []),
|
||||
("calculator", "nonexistent.py", []),
|
||||
]
|
||||
|
||||
for working_dir, file_path, args in cases:
|
||||
if args:
|
||||
print(f'run_python_file("{working_dir}", "{file_path}", {args}):')
|
||||
else:
|
||||
print(f'run_python_file("{working_dir}", "{file_path}"):')
|
||||
|
||||
result = run_python_file(working_dir, file_path, args)
|
||||
|
||||
print("Result:")
|
||||
for line in result.splitlines():
|
||||
print(" " + line)
|
||||
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_tests()
|
||||
Reference in New Issue
Block a user