feat: complete assignment

calculator/README.md

@@ -0,0 +1 @@
+# calculator

calculator/lorem.txt

@@ -0,0 +1 @@
+wait, this isn't lorem ipsum

calculator/main.py

@@ -0,0 +1,28 @@
+import sys
+
+from pkg.calculator import Calculator
+from pkg.render import format_json_output
+
+
+def main() -> None:
+    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)
+        if result is not None:
+            to_print = format_json_output(expression, result)
+            print(to_print)
+        else:
+            print("Error: Expression is empty or contains only whitespace.")
+    except Exception as e:
+        print(f"Error: {e}")
+
+
+if __name__ == "__main__":
+    main()

calculator/pkg/calculator.py

@@ -0,0 +1,62 @@
+from collections.abc import Callable
+
+
+class Calculator:
+    def __init__(self) -> None:
+        self.operators: dict[str, Callable[[float, float], float]] = {
+            "+": lambda a, b: a + b,
+            "-": lambda a, b: a - b,
+            "*": lambda a, b: a * b,
+            "/": lambda a, b: a / b,
+        }
+        self.precedence: dict[str, int] = {
+            "+": 1,
+            "-": 1,
+            "*": 2,
+            "/": 2,
+        }
+
+    def evaluate(self, expression: str) -> float | None:
+        if not expression or expression.isspace():
+            return None
+        tokens = expression.strip().split()
+        return self._evaluate_infix(tokens)
+
+    def _evaluate_infix(self, tokens: list[str]) -> float:
+        values: list[float] = []
+        operators: list[str] = []
+
+        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: list[str], values: list[float]) -> None:
+        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))

calculator/pkg/morelorem.txt

@@ -0,0 +1 @@
+lorem ipsum dolor sit amet

calculator/pkg/render.py

@@ -0,0 +1,14 @@
+import json
+
+
+def format_json_output(expression: str, result: float, indent: int = 2) -> str:
+    if isinstance(result, float) and result.is_integer():
+        result_to_dump = int(result)
+    else:
+        result_to_dump = result
+
+    output_data = {
+        "expression": expression,
+        "result": result_to_dump,
+    }
+    return json.dumps(output_data, indent=indent)

calculator/tests.py

@@ -0,0 +1,48 @@
+import unittest
+
+from pkg.calculator import Calculator
+
+
+class TestCalculator(unittest.TestCase):
+    def setUp(self) -> None:
+        self.calculator = Calculator()
+
+    def test_addition(self) -> None:
+        result = self.calculator.evaluate("3 + 5")
+        self.assertEqual(result, 8)
+
+    def test_subtraction(self) -> None:
+        result = self.calculator.evaluate("10 - 4")
+        self.assertEqual(result, 6)
+
+    def test_multiplication(self) -> None:
+        result = self.calculator.evaluate("3 * 4")
+        self.assertEqual(result, 12)
+
+    def test_division(self) -> None:
+        result = self.calculator.evaluate("10 / 2")
+        self.assertEqual(result, 5)
+
+    def test_nested_expression(self) -> None:
+        result = self.calculator.evaluate("3 * 4 + 5")
+        self.assertEqual(result, 17)
+
+    def test_complex_expression(self) -> None:
+        result = self.calculator.evaluate("2 * 3 - 8 / 2 + 5")
+        self.assertEqual(result, 7)
+
+    def test_empty_expression(self) -> None:
+        result = self.calculator.evaluate("")
+        self.assertIsNone(result)
+
+    def test_invalid_operator(self) -> None:
+        with self.assertRaises(ValueError):
+            self.calculator.evaluate("$ 3 5")
+
+    def test_not_enough_operands(self) -> None:
+        with self.assertRaises(ValueError):
+            self.calculator.evaluate("+ 3")
+
+
+if __name__ == "__main__":
+    unittest.main()

functions/call_function.py

@@ -0,0 +1,62 @@
+from typing import Callable
+
+from google.genai import types
+
+from functions.get_file_content import (get_file_content,
+                                        schema_get_file_content)
+from functions.get_files_info import get_files_info, schema_get_files_info
+from functions.run_python_file import run_python_file, schema_run_python_file
+from functions.write_file import schema_write_file, write_file
+
+available_functions = types.Tool(
+    function_declarations=[
+        schema_get_file_content,
+        schema_get_files_info,
+        schema_run_python_file,
+        schema_write_file,
+    ]
+)
+
+function_map: dict[str, Callable[..., str]] = {
+    "get_file_content": get_file_content,
+    "get_files_info": get_files_info,
+    "run_python_file": run_python_file,
+    "write_file": write_file,
+}
+
+
+def call_function(
+    function_call: types.FunctionCall, verbose: bool = False
+) -> types.Content:
+    if verbose:
+        print(f"Calling function: {function_call.name}({function_call.args})")
+    else:
+        print(f"Calling function: {function_call.name}")
+
+    function_name = function_call.name or ""
+    function = function_map.get(function_name)
+
+    if function is None:
+        return types.Content(
+            role="tool",
+            parts=[
+                types.Part.from_function_response(
+                    name=function_name,
+                    response={"error": f"Unknown function: {function_name}"},
+                )
+            ],
+        )
+
+    args = dict(function_call.args) if function_call.args else {}
+    args["working_directory"] = "./calculator"
+    function_result = function(**args)
+
+    return types.Content(
+        role="tool",
+        parts=[
+            types.Part.from_function_response(
+                name=function_name,
+                response={"result": function_result},
+            )
+        ],
+    )

functions/get_file_content.py

@@ -0,0 +1,47 @@
+import os
+
+from google.genai import types
+
+
+def get_file_content(working_directory: str, file_path: str) -> str:
+    try:
+        wdp = os.path.abspath(working_directory)
+        target = os.path.normpath(os.path.join(wdp, file_path))
+        valid_target_dir = os.path.commonpath([wdp, target]) == wdp
+
+        if not valid_target_dir:
+            return f'Error: Cannot read "{file_path}" as it is outside the permitted working directory'
+
+        if not os.path.isfile(target):
+            return f'Error: File not found or is not a regular file: "{file_path}"'
+
+        MAX_CHARS = 10000
+        content = ""
+
+        with open(target, "r") as f:
+            content += f.read(MAX_CHARS)
+
+            if f.read(1):
+                content += (
+                    f'[...File "{file_path}" truncated at {MAX_CHARS} characters]'
+                )
+
+        return content
+    except Exception as e:
+        return f"Error: {e}"
+
+
+schema_get_file_content = types.FunctionDeclaration(
+    name="get_file_content",
+    description="",
+    parameters=types.Schema(
+        type=types.Type.OBJECT,
+        properties={
+            "file_path": types.Schema(
+                type=types.Type.STRING,
+                description="",
+            ),
+        },
+        required=["file_path"],
+    ),
+)

functions/get_files_info.py

@@ -0,0 +1,44 @@
+import os
+
+from google.genai import types
+
+
+def get_files_info(working_directory, directory: str = ".") -> str:
+    try:
+        wdp = os.path.abspath(working_directory)
+        target = os.path.normpath(os.path.join(wdp, directory))
+        valid_target_dir = os.path.commonpath([wdp, target]) == wdp
+
+        if not valid_target_dir:
+            return f'Error: Cannot list "{directory}" as it is outside the permitted directory'
+
+        if not os.path.isdir(target):
+            return f'Error: "{directory}" is not a directory'
+
+        file_descriptors = []
+
+        for file_name in os.listdir(target):
+            file_path = os.path.normpath(os.path.join(target, file_name))
+            size = os.path.getsize(file_path)
+            is_dir = os.path.isdir(file_path)
+            file_descriptors.append(f"- {file_name}: file_size={size}, is_dir={is_dir}")
+
+        return "\n".join(file_descriptors)
+    except Exception as e:
+        return f"Error: {e}"
+
+
+schema_get_files_info = types.FunctionDeclaration(
+    name="get_files_info",
+    description="Lists files in a specified directory relative to the working directory, providing file size and directory status",
+    parameters=types.Schema(
+        type=types.Type.OBJECT,
+        properties={
+            "directory": types.Schema(
+                type=types.Type.STRING,
+                description="Directory path to list files from, relative to the working directory (default is the working directory itself)",
+            ),
+        },
+        required=["directory"],
+    ),
+)

functions/run_python_file.py

@@ -0,0 +1,66 @@
+import os
+import subprocess
+
+from google.genai import types
+
+
+def run_python_file(
+    working_directory: str, file_path: str, args: list[str] | None = None
+) -> str:
+    try:
+        wdp = os.path.abspath(working_directory)
+        target = os.path.normpath(os.path.join(wdp, file_path))
+        valid_target_dir = os.path.commonpath([wdp, target]) == wdp
+
+        if not valid_target_dir:
+            return f'Error: Cannot execute "{file_path}" as it is outside the permitted working directory'
+
+        if not os.path.isfile(target):
+            return f'Error: "{file_path}" does not exist or is not a regular file'
+
+        if not target.endswith(".py"):
+            return f'Error: "{file_path}" is not a Python file'
+
+        command = ["python", target]
+
+        if args != None:
+            command.extend(args)
+
+        result = subprocess.run(
+            command,
+            cwd=working_directory,
+            text=True,
+            timeout=30_000,
+            capture_output=True,
+        )
+
+        if result.returncode != 0:
+            return f"Process exited with code {result.returncode}"
+
+        if len(result.stdout) == 0 and len(result.stderr) == 0:
+            return "No output produced"
+
+        return f"STDOUT: {result.stdout}\nSTDERR: {result.stderr}"
+    except Exception as e:
+        return f"Error: {e}"
+
+
+schema_run_python_file = types.FunctionDeclaration(
+    name="run_python_file",
+    description="",
+    parameters=types.Schema(
+        type=types.Type.OBJECT,
+        properties={
+            "file_path": types.Schema(
+                type=types.Type.STRING,
+                description="",
+            ),
+            "args": types.Schema(
+                type=types.Type.ARRAY,
+                items=types.Schema(type=types.Type.STRING),
+                description="",
+            ),
+        },
+        required=["file_path"],
+    ),
+)

functions/write_file.py

@@ -0,0 +1,47 @@
+import os
+
+from google.genai import types
+
+
+def write_file(working_directory: str, file_path: str, content: str) -> str:
+    try:
+        wdp = os.path.abspath(working_directory)
+        target = os.path.normpath(os.path.join(wdp, file_path))
+        valid_target_dir = os.path.commonpath([wdp, target]) == wdp
+
+        if not valid_target_dir:
+            return f'Error: Cannot write to "{file_path}" as it is outside the permitted working directory'
+
+        if os.path.isdir(target):
+            return f'Error: Cannot write to "{file_path}" as it is a directory'
+
+        os.makedirs("/".join(target.split("/")[:-1]), exist_ok=True)
+
+        with open(target, "w") as f:
+            f.write(content)
+
+        return (
+            f'Successfully wrote to "{file_path}" ({len(content)} characters written)'
+        )
+    except Exception as e:
+        return f"Error: {e}"
+
+
+schema_write_file = types.FunctionDeclaration(
+    name="write_file",
+    description="",
+    parameters=types.Schema(
+        type=types.Type.OBJECT,
+        properties={
+            "file_path": types.Schema(
+                type=types.Type.STRING,
+                description="",
+            ),
+            "content": types.Schema(
+                type=types.Type.STRING,
+                description="",
+            ),
+        },
+        required=["file_path", "content"],
+    ),
+)

main.py

@@ -1,5 +1,105 @@
-def main():
-    print("Hello from einstein!")
+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__":

test_get_file_content.py

@@ -0,0 +1,23 @@
+from functions.get_file_content import get_file_content
+
+
+def test() -> None:
+    result = get_file_content("calculator", "lorem.txt")
+    print(f"lorem.txt length: {len(result)}")
+    print(f"lorem.txt truncated: {'truncated' in result}")
+
+    result = get_file_content("calculator", "main.py")
+    print(result)
+
+    result = get_file_content("calculator", "pkg/calculator.py")
+    print(result)
+
+    result = get_file_content("calculator", "/bin/cat")
+    print(result)
+
+    result = get_file_content("calculator", "pkg/does_not_exist.py")
+    print(result)
+
+
+if __name__ == "__main__":
+    test()

test_get_files_info.py

@@ -0,0 +1,27 @@
+from functions.get_files_info import get_files_info
+
+
+def test() -> None:
+    result = get_files_info("calculator", ".")
+    print("Result for current directory:")
+    print(result)
+    print("")
+
+    result = get_files_info("calculator", "pkg")
+    print("Result for 'pkg':")
+    print(result)
+    print("")
+
+    result = get_files_info("calculator", "/bin")
+    print("Result for '/bin' directory:")
+    print(result)
+    print("")
+
+    result = get_files_info("calculator", "../")
+    print("Result for '../' directory:")
+    print(result)
+    print("")
+
+
+if __name__ == "__main__":
+    test()

test_run_python_file.py

@@ -0,0 +1,29 @@
+from functions.run_python_file import run_python_file
+
+def test() -> None:
+    result = run_python_file("calculator", "main.py")
+    print(result)
+    print("")
+
+    result = run_python_file("calculator", "main.py", ["3 + 5"])
+    print(result)
+    print("")
+
+    result = run_python_file("calculator", "tests.py")
+    print(result)
+    print("")
+
+    result = run_python_file("calculator", "../main.py")
+    print(result)
+    print("")
+
+    result = run_python_file("calculator", "nonexistent.py")
+    print(result)
+    print("")
+
+    result = run_python_file("calculator", "lorem.txt")
+    print(result)
+    print("")
+
+if __name__ == "__main__":
+    test()

test_write_file.py

@@ -0,0 +1,16 @@
+from functions.write_file import write_file
+
+
+def test() -> None:
+    result = write_file("calculator", "lorem.txt", "wait, this isn't lorem ipsum")
+    print(result)
+
+    result = write_file("calculator", "pkg/morelorem.txt", "lorem ipsum dolor sit amet")
+    print(result)
+
+    result = write_file("calculator", "/tmp/temp.txt", "this should not be allowed")
+    print(result)
+
+
+if __name__ == "__main__":
+    test()