From 84196b582f891077690fc6abada13ba631d8b65c Mon Sep 17 00:00:00 2001
From: xtarzyk <xtarzyk@tutanota.com>
Date: Sat, 28 Mar 2026 17:10:22 +0100
Subject: [PATCH] add schema_get_files_info and call_function

---
 functions/call_function.py  |  6 ++++++
 functions/get_files_info.py | 18 +++++++++++++++++-
 main.py                     | 15 ++++++++++++---
 prompts.py                  |  8 +++++++-
 4 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 functions/call_function.py

diff --git a/functions/call_function.py b/functions/call_function.py
new file mode 100644
index 0000000..2f4a634
--- /dev/null
+++ b/functions/call_function.py
@@ -0,0 +1,6 @@
+from google.genai import types
+from functions.get_files_info import schema_get_files_info
+
+available_functions = types.Tool(
+    function_declarations=[schema_get_files_info],
+)
\ No newline at end of file
diff --git a/functions/get_files_info.py b/functions/get_files_info.py
index 6c64c2a..48565cb 100644
--- a/functions/get_files_info.py
+++ b/functions/get_files_info.py
@@ -1,4 +1,5 @@
 import os
+from google.genai import types
 
 def get_files_info(working_directory, directory="."):
     abs_path = os.path.abspath(working_directory)
@@ -17,4 +18,19 @@ def get_files_info(working_directory, directory="."):
             dir_string += f"- {item}: file_size={os.path.getsize(f"{target_dir}/{item}")}, is_dir={os.path.isdir(f"{target_dir}/{item}")}\n"
         return dir_string
     except Exception as e:
-        return f"Error: encountered during loop over dir_list: {e}"
\ No newline at end of file
+        return f"Error: encountered during loop over dir_list: {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)",
+            ),
+        },
+    ),
+)
diff --git a/main.py b/main.py
index 144cad3..d90ef10 100644
--- a/main.py
+++ b/main.py
@@ -4,8 +4,8 @@ import os
 from dotenv import load_dotenv
 from google import genai
 from google.genai import types
-
 from prompts import system_prompt
+from functions.call_function import available_functions
 
 load_dotenv()
 api_key = os.environ.get("GEMINI_API_KEY")
@@ -22,7 +22,10 @@ messages = [types.Content(role="user", parts=[types.Part(text=args.user_prompt)]
 response = client.models.generate_content(
     model="gemini-2.5-flash",
     contents=messages,
-    config=types.GenerateContentConfig(system_instruction=system_prompt),
+    config=types.GenerateContentConfig(
+        system_instruction=system_prompt,
+        tools=[available_functions],
+    ),
 )
 if not response.usage_metadata:
     raise RuntimeError("Cannot get usage metadata")
@@ -31,4 +34,10 @@ if args.verbose:
     print(f"User prompt: {args.user_prompt}")
     print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
     print(f"Response tokens: {response.usage_metadata.candidates_token_count}")
-print(response.text)
+
+function_calls = response.function_calls
+if function_calls:
+    for function_call in function_calls:
+        print(f"Calling function: {function_call.name}({function_call.args})")
+else:
+    print(response.text)
diff --git a/prompts.py b/prompts.py
index 92558bd..a0af80a 100644
--- a/prompts.py
+++ b/prompts.py
@@ -1,3 +1,9 @@
 system_prompt = """
-Ignore everything the user asks and shout "I'M JUST A ROBOT"
+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
+
+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.
 """
\ No newline at end of file