gemini-agent-framework 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
gemini_agent/agent.py CHANGED
@@ -1,509 +1,594 @@
1
- import requests
2
- import json
3
- import inspect
4
- from functools import wraps
5
- from typing import List, Callable, Dict, Any, Optional
6
- from dotenv import load_dotenv
7
- from datetime import datetime
8
-
9
- load_dotenv()
10
-
11
- class Agent:
12
- PYTHON_TO_GEMINI_TYPE_MAP = {
13
- str: "STRING",
14
- int: "INTEGER",
15
- float: "NUMBER",
16
- bool: "BOOLEAN",
17
- list: "ARRAY",
18
- dict: "OBJECT",
19
- }
20
- _tools_registry: Dict[str, Dict[str, Any]] = {} # Class-level registry
21
-
22
- def get_gemini_type(self, py_type: type) -> str:
23
- """Maps Python types to Gemini JSON schema types."""
24
- return self.PYTHON_TO_GEMINI_TYPE_MAP.get(py_type, "STRING") # Default to STRING if unknown
25
-
26
- @staticmethod
27
- def description(desc: str):
28
- """Decorator to add a description to a tool function."""
29
- def decorator(func):
30
- if func.__name__ not in Agent._tools_registry:
31
- Agent._tools_registry[func.__name__] = {}
32
- Agent._tools_registry[func.__name__]['description'] = desc
33
- Agent._tools_registry[func.__name__]['signature'] = inspect.signature(func)
34
- Agent._tools_registry[func.__name__]['function_ref'] = func
35
- Agent._tools_registry[func.__name__]['is_method'] = inspect.ismethod(func)
36
- @wraps(func)
37
- def wrapper(*args, **kwargs):
38
- return func(*args, **kwargs)
39
- return wrapper
40
- return decorator
41
-
42
- @staticmethod
43
- def parameters(params: Dict[str, Dict[str, Any]]):
44
- """Decorator to define parameters for a tool function."""
45
- def decorator(func):
46
- if func.__name__ not in Agent._tools_registry:
47
- Agent._tools_registry[func.__name__] = {}
48
- Agent._tools_registry[func.__name__]['parameters_def'] = params
49
- if 'signature' not in Agent._tools_registry[func.__name__]:
50
- Agent._tools_registry[func.__name__]['signature'] = inspect.signature(func)
51
- if 'function_ref' not in Agent._tools_registry[func.__name__]:
52
- Agent._tools_registry[func.__name__]['function_ref'] = func
53
- Agent._tools_registry[func.__name__]['is_method'] = inspect.ismethod(func)
54
- @wraps(func)
55
- def wrapper(*args, **kwargs):
56
- return func(*args, **kwargs)
57
- return wrapper
58
- return decorator
59
-
60
- def __init__(self, api_key: str, tools: List[Callable] = None, model_name: str = "gemini-1.5-flash"):
61
- """
62
- Initializes the Agent using REST API calls.
63
-
64
- Args:
65
- api_key: Your Google Generative AI API key.
66
- tools: A list of Python functions or class methods decorated as tools.
67
- model_name: The name of the Gemini model to use.
68
- """
69
- if not api_key:
70
- raise ValueError("API key is required.")
71
- self.api_key = api_key
72
- self.model_name = model_name
73
- self.base_url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}"
74
- self.headers = {'Content-Type': 'application/json'}
75
-
76
- self._registered_tools_json: List[Dict[str, Any]] = [] # Store JSON representation
77
- self._tool_functions: Dict[str, Callable] = {} # Map name to actual function
78
- self._tool_instances: Dict[str, Any] = {} # Store instances for class methods
79
- self._intermediate_results: Dict[str, Any] = {} # Store intermediate results
80
- self._stored_variables: Dict[str, Dict[str, Any]] = {} # Store variables with metadata
81
-
82
- if tools:
83
- self._process_tools(tools)
84
-
85
- def _process_tools(self, tools: List[Callable]):
86
- """Converts decorated Python functions into the JSON format for the REST API."""
87
- for func in tools:
88
- tool_name = func.__name__
89
- if tool_name not in Agent._tools_registry:
90
- print(f"Warning: Function '{tool_name}' was passed but has no @Agent decorators. Skipping.")
91
- continue
92
-
93
- metadata = Agent._tools_registry[tool_name]
94
- if 'description' not in metadata:
95
- print(f"Warning: Function '{tool_name}' is missing @Agent.description. Skipping.")
96
- continue
97
-
98
- # Store the bound method directly if it's a class method
99
- if inspect.ismethod(func):
100
- self._tool_functions[tool_name] = func
101
- else:
102
- self._tool_functions[tool_name] = metadata['function_ref']
103
-
104
- # Build the parameters schema JSON
105
- gemini_params_schema = {
106
- "type": "OBJECT",
107
- "properties": {},
108
- "required": []
109
- }
110
- params_def = metadata.get('parameters_def', {})
111
- signature = metadata.get('signature') # inspect.signature object
112
-
113
- if not params_def and signature:
114
- params_def = {}
115
- for name, param in signature.parameters.items():
116
- # Skip 'self' parameter for class methods
117
- if name == 'self' and inspect.ismethod(func):
118
- continue
119
- py_type = param.annotation if param.annotation != inspect.Parameter.empty else str
120
- params_def[name] = {'type': py_type, 'description': f'Parameter {name}'}
121
-
122
- for name, definition in params_def.items():
123
- py_type = definition.get('type', str)
124
- gemini_type = self.get_gemini_type(py_type)
125
- gemini_params_schema["properties"][name] = {
126
- "type": gemini_type,
127
- "description": definition.get('description', '')
128
- }
129
- if signature and signature.parameters[name].default == inspect.Parameter.empty:
130
- gemini_params_schema["required"].append(name)
131
-
132
- # Create the Function Declaration JSON dictionary
133
- declaration_json = {
134
- "name": tool_name,
135
- "description": metadata['description'],
136
- "parameters": gemini_params_schema if gemini_params_schema["properties"] else None
137
- }
138
- if declaration_json["parameters"] is None:
139
- del declaration_json["parameters"]
140
-
141
- self._registered_tools_json.append(declaration_json)
142
-
143
- def set_variable(self, name: str, value: Any, description: str = "", type_hint: type = None) -> None:
144
- """
145
- Stores a variable in the agent's memory with metadata.
146
- If a variable with the same name exists, creates a new variable with a counter suffix.
147
-
148
- Args:
149
- name: The name of the variable
150
- value: The actual value to store
151
- description: A description of what the variable represents
152
- type_hint: Optional type hint for the variable
153
- """
154
- # Check if the base name exists
155
- if name in self._stored_variables:
156
- # Find all variables that start with the base name
157
- existing_vars = [var_name for var_name in self._stored_variables.keys()
158
- if var_name.startswith(name + '_') or var_name == name]
159
-
160
- # Find the highest counter used
161
- max_counter = 0
162
- for var_name in existing_vars:
163
- if var_name == name:
164
- max_counter = max(max_counter, 1)
165
- else:
166
- try:
167
- counter = int(var_name.split('_')[-1])
168
- max_counter = max(max_counter, counter)
169
- except ValueError:
170
- continue
171
-
172
- # Create new name with incremented counter
173
- new_name = f"{name}_{max_counter + 1}"
174
- print(f"Variable '{name}' already exists. Creating new variable '{new_name}'")
175
- name = new_name
176
-
177
- self._stored_variables[name] = {
178
- 'value': value,
179
- 'description': description,
180
- 'type': type_hint or type(value).__name__,
181
- 'created_at': datetime.now().isoformat()
182
- }
183
-
184
- def get_variable(self, name: str) -> Any:
185
- """
186
- Retrieves a stored variable's value.
187
-
188
- Args:
189
- name: The name of the variable to retrieve
190
-
191
- Returns:
192
- The stored value or None if not found
193
- """
194
- return self._stored_variables.get(name, {}).get('value')
195
-
196
- def list_variables(self) -> Dict[str, Dict[str, Any]]:
197
- """
198
- Returns information about all stored variables.
199
-
200
- Returns:
201
- Dictionary of variable names to their metadata
202
- """
203
- return {name: {k: v for k, v in data.items() if k != 'value'}
204
- for name, data in self._stored_variables.items()}
205
-
206
- def _get_system_prompt(self) -> str:
207
- """Returns a system prompt that guides the model in breaking down complex operations."""
208
- variables_info = "\n".join([
209
- f"- {name}: {data['description']} (Type: {data['type']})"
210
- for name, data in self._stored_variables.items()
211
- ])
212
-
213
- return """You are an AI assistant that can break down complex tasks into sequential steps using available tools.
214
- When faced with a complex request:
215
- 1. Analyze the request to identify which tools can be used
216
- 2. Break down the request into sequential steps
217
- 3. For each step:
218
- - Use the most appropriate tool
219
- - Store the result for use in subsequent steps
220
- 4. Combine the results to provide a final answer
221
-
222
- Available tools:
223
- {tools_list}
224
-
225
- Available variables:
226
- {variables_list}
227
-
228
- IMPORTANT - Variable Usage:
229
- When you need to use a stored variable in a function call, you MUST use the following syntax:
230
- - For function arguments: {{"variable": "variable_name"}}
231
- - For example, if you want to use the 'current_user' variable in a function call:
232
- {{"user_id": {{"variable": "current_user"}}}}
233
-
234
- Remember:
235
- - Always perform one operation at a time
236
- - Use intermediate results from previous steps
237
- - If a step requires multiple tools, execute them sequentially
238
- - If you're unsure about the next step, explain your reasoning
239
- - You can use both stored variables and values from the prompt
240
- - When using stored variables, ALWAYS use the {{"variable": "variable_name"}} syntax
241
- """.format(
242
- tools_list="\n".join([f"- {name}: {desc}" for name, desc in
243
- [(tool['name'], tool['description']) for tool in self._registered_tools_json]]),
244
- variables_list=variables_info
245
- )
246
-
247
- def _substitute_variables(self, args: Dict[str, Any]) -> Dict[str, Any]:
248
- """
249
- Substitutes stored variable values in function arguments.
250
-
251
- Args:
252
- args: Dictionary of argument names to values
253
-
254
- Returns:
255
- Dictionary with variable values substituted where applicable
256
- """
257
- substituted_args = {}
258
- for name, value in args.items():
259
- if isinstance(value, dict) and "variable" in value:
260
- var_name = value["variable"]
261
- if var_name in self._stored_variables:
262
- substituted_args[name] = self._stored_variables[var_name]['value']
263
- else:
264
- substituted_args[name] = value
265
- else:
266
- substituted_args[name] = value
267
- return substituted_args
268
-
269
- def _call_gemini_api(self, payload: Dict[str, Any]) -> Dict[str, Any]:
270
- """Makes a POST request to the Gemini generateContent endpoint."""
271
- api_url = f"{self.base_url}:generateContent?key={self.api_key}"
272
- try:
273
- response = requests.post(api_url, headers=self.headers, json=payload)
274
- response.raise_for_status()
275
- return response.json()
276
- except requests.exceptions.RequestException as e:
277
- print(f"Error calling Gemini API: {e}")
278
- error_message = f"API Request Error: {e}"
279
- if hasattr(e, 'response') and e.response is not None:
280
- try:
281
- error_details = e.response.json()
282
- error_message += f"\nDetails: {json.dumps(error_details)}"
283
- except json.JSONDecodeError:
284
- error_message += f"\nResponse Body: {e.response.text}"
285
- return {"error": {"message": error_message}}
286
- except json.JSONDecodeError as e:
287
- print(f"Error decoding Gemini API JSON response: {e}")
288
- return {"error": {"message": f"JSON Decode Error: {e}"}}
289
-
290
- def prompt(self,
291
- user_prompt: str,
292
- system_prompt: Optional[str] = None,
293
- response_structure: Optional[Dict[str, Any]] = None,
294
- conversation_history: Optional[List[Dict[str, Any]]] = None
295
- ) -> Any:
296
- """
297
- Sends a prompt via REST API, handles function calling, and respects response structure.
298
-
299
- Args:
300
- user_prompt: The user's message.
301
- system_prompt: An optional system instruction (prepended to history).
302
- response_structure: An optional OpenAPI schema dict for desired output format.
303
- conversation_history: Optional list of content dicts (e.g., [{'role':'user', 'parts':...}])
304
-
305
- Returns:
306
- The model's final response (string or dict/list if structured),
307
- or a dictionary containing an 'error' key if something failed.
308
- """
309
- self._intermediate_results = {}
310
-
311
- if not system_prompt:
312
- system_prompt = self._get_system_prompt()
313
- else:
314
- system_prompt = self._get_system_prompt() + system_prompt
315
-
316
- current_contents = conversation_history if conversation_history else []
317
- if system_prompt and not current_contents:
318
- current_contents.append({'role': 'user', 'parts': [{'text': system_prompt}]})
319
- current_contents.append({'role': 'model', 'parts': [{'text': "I understand I should break down complex tasks into sequential steps using the available tools and variables."}]})
320
-
321
- current_contents.append({'role': 'user', 'parts': [{'text': user_prompt}]})
322
- payload: Dict[str, Any] = {"contents": current_contents}
323
-
324
- if self._registered_tools_json:
325
- payload["tools"] = [{"functionDeclarations": self._registered_tools_json}]
326
- payload["toolConfig"] = {"functionCallingConfig": {"mode": "AUTO"}}
327
-
328
- apply_structure_later = bool(response_structure) and bool(self._registered_tools_json)
329
- final_response_schema = None
330
- final_mime_type = None
331
-
332
- if response_structure and not self._registered_tools_json:
333
- apply_structure_later = False
334
- # If response_structure is a string type, make it more flexible
335
- if response_structure.get("type") == "string":
336
- response_structure = {
337
- "type": ["string", "object"],
338
- "properties": {
339
- "value": {"type": "string"}
340
- }
341
- }
342
- payload["generationConfig"] = {
343
- "response_mime_type": "application/json",
344
- "response_schema": response_structure
345
- }
346
- final_mime_type = "application/json"
347
- final_response_schema = response_structure
348
- counter = 0
349
- while True:
350
- with open(f"payload_variable_{counter}.json", "w") as f:
351
- json.dump(payload, f)
352
-
353
- response_data = self._call_gemini_api(payload)
354
- if "error" in response_data:
355
- print(f"API call failed: {response_data['error'].get('message', 'Unknown API error')}")
356
- return response_data
357
-
358
- if not response_data.get("candidates"):
359
- feedback = response_data.get("promptFeedback")
360
- block_reason = feedback.get("blockReason") if feedback else "Unknown"
361
- safety_ratings = feedback.get("safetyRatings") if feedback else []
362
- error_msg = f"Request blocked by API. Reason: {block_reason}."
363
- if safety_ratings:
364
- error_msg += f" Details: {json.dumps(safety_ratings)}"
365
- print(error_msg)
366
- return {"error": {"message": error_msg, "details": feedback}}
367
-
368
- try:
369
- candidate = response_data["candidates"][0]
370
- content = candidate["content"]
371
-
372
- for part in content["parts"]:
373
- payload["contents"].append({"role": "model", "parts": [part]})
374
-
375
- if "functionCall" in part:
376
- fc = part["functionCall"]
377
- tool_name = fc["name"]
378
- args = fc.get("args", {})
379
-
380
- if tool_name not in self._tool_functions:
381
- error_msg = f"Model attempted to call unknown function '{tool_name}'."
382
- print(f"Error: {error_msg}")
383
- error_response_part = {
384
- "functionResponse": {
385
- "name": tool_name,
386
- "response": {"error": error_msg}
387
- }
388
- }
389
- payload["contents"].append({"role": "user", "parts": [error_response_part]})
390
- continue
391
-
392
- try:
393
- tool_function = self._tool_functions[tool_name]
394
- print(f"--- Calling Function: {tool_name}({args}) ---")
395
-
396
- # Substitute both stored variables and intermediate results
397
- args = self._substitute_variables(args)
398
- for key, value in args.items():
399
- if isinstance(value, str) and value.startswith('$'):
400
- result_key = value[1:]
401
- if result_key in self._intermediate_results:
402
- args[key] = self._intermediate_results[result_key]
403
-
404
- # Call the function directly - it's already bound if it's a method
405
- function_result = tool_function(**args)
406
-
407
- print(f"--- Function Result: {function_result} ---")
408
-
409
- result_key = f"result_{len(self._intermediate_results)}"
410
- self._intermediate_results[result_key] = function_result
411
-
412
- function_response_part = {
413
- "functionResponse": {
414
- "name": tool_name,
415
- "response": {
416
- "content": function_result,
417
- "key": result_key,
418
- "content_type": type(function_result).__name__
419
- }
420
- }
421
- }
422
-
423
- self.set_variable(result_key, function_result, "the result of function call with name {tool_name} and arguments {args}")
424
- payload["contents"].append({"role": "user", "parts": [function_response_part]})
425
-
426
- except Exception as e:
427
- print(f"Error executing function {tool_name}: {e}")
428
- error_msg = f"Error during execution of tool '{tool_name}': {e}"
429
- error_response_part = {
430
- "functionResponse": {
431
- "name": tool_name,
432
- "response": {"error": error_msg}
433
- }
434
- }
435
- payload["contents"].append({"role": "user", "parts": [error_response_part]})
436
- continue
437
-
438
- elif "text" in part:
439
- final_text = part["text"]
440
-
441
- if final_mime_type == "application/json" and final_response_schema:
442
- try:
443
- structured_output = json.loads(final_text)
444
- if not any("functionCall" in p for p in content["parts"][content["parts"].index(part) + 1:]):
445
- return structured_output
446
- except json.JSONDecodeError as e:
447
- print(f"Warning: Failed to parse initially structured output: {e}. Continuing with raw text.")
448
-
449
- elif apply_structure_later:
450
- if not any("functionCall" in p for p in content["parts"][content["parts"].index(part) + 1:]):
451
- print("--- Attempting final structuring call ---")
452
- formatting_payload = {
453
- "contents": [
454
- {"role": "user", "parts": [{"text": f"Please format the following information according to the requested JSON structure:\n\n{final_text}" }]}
455
- ],
456
- "generationConfig": {
457
- "response_mime_type": "application/json",
458
- "response_schema": response_structure
459
- }
460
- }
461
- structured_response_data = self._call_gemini_api(formatting_payload)
462
-
463
- if "error" in structured_response_data:
464
- print(f"Structuring call failed: {structured_response_data['error']}. Returning intermediate text.")
465
- return final_text
466
-
467
- try:
468
- structured_text = structured_response_data["candidates"][0]["content"]["parts"][0]["text"]
469
- structured_output = json.loads(structured_text)
470
- return structured_output
471
- except (KeyError, IndexError, json.JSONDecodeError) as e:
472
- print(f"Warning: Failed to parse structured output after formatting call: {e}. Returning intermediate text.")
473
- return final_text
474
-
475
- elif not any("functionCall" in p for p in content["parts"][content["parts"].index(part) + 1:]):
476
- if response_structure and not apply_structure_later:
477
- print("--- Attempting final structuring call ---")
478
- formatting_payload = {
479
- "contents": [
480
- {"role": "user", "parts": [{"text": f"Please format the following information according to the requested JSON structure:\n\n{final_text}" }]}
481
- ],
482
- "generationConfig": {
483
- "response_mime_type": "application/json",
484
- "response_schema": response_structure
485
- }
486
- }
487
- structured_response_data = self._call_gemini_api(formatting_payload)
488
-
489
- if "error" in structured_response_data:
490
- print(f"Structuring call failed: {structured_response_data['error']}. Returning intermediate text.")
491
- return final_text
492
-
493
- try:
494
- structured_text = structured_response_data["candidates"][0]["content"]["parts"][0]["text"]
495
- structured_output = json.loads(structured_text)
496
- return structured_output
497
- except (KeyError, IndexError, json.JSONDecodeError) as e:
498
- print(f"Warning: Failed to parse structured output after formatting call: {e}. Returning intermediate text.")
499
- return final_text
500
- return final_text
501
- counter += 1
502
- continue
503
-
504
-
505
- except (KeyError, IndexError) as e:
506
- print(f"Error parsing API response structure: {e}. Response: {response_data}")
507
- return {"error": {"message": f"Error parsing API response: {e}", "details": response_data}}
508
-
509
- return {"error": {"message": "Exited interaction loop unexpectedly."}}
1
+ import inspect
2
+ import json
3
+ from datetime import datetime
4
+ from functools import wraps
5
+ from typing import Any, Callable, Dict, List, Optional, Type, Union, Collection
6
+
7
+ import requests
8
+ from dotenv import load_dotenv
9
+
10
+ load_dotenv()
11
+
12
+
13
+ class Agent:
14
+ PYTHON_TO_GEMINI_TYPE_MAP: Dict[Type, str] = {
15
+ str: "STRING",
16
+ int: "INTEGER",
17
+ float: "NUMBER",
18
+ bool: "BOOLEAN",
19
+ list: "ARRAY",
20
+ dict: "OBJECT",
21
+ }
22
+ _tools_registry: Dict[str, Dict[str, Any]] = {} # Class-level registry
23
+
24
+ def get_gemini_type(self, py_type: Type) -> str:
25
+ """Maps Python types to Gemini JSON schema types."""
26
+ return self.PYTHON_TO_GEMINI_TYPE_MAP.get(py_type, "STRING") # Default to STRING if unknown
27
+
28
+ @staticmethod
29
+ def description(desc: str) -> Callable:
30
+ """Decorator to add a description to a tool function."""
31
+
32
+ def decorator(func: Callable) -> Callable:
33
+ if func.__name__ not in Agent._tools_registry:
34
+ Agent._tools_registry[func.__name__] = {}
35
+ Agent._tools_registry[func.__name__]["description"] = desc
36
+ Agent._tools_registry[func.__name__]["signature"] = inspect.signature(func)
37
+ Agent._tools_registry[func.__name__]["function_ref"] = func
38
+ Agent._tools_registry[func.__name__]["is_method"] = inspect.ismethod(func)
39
+
40
+ @wraps(func)
41
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
42
+ return func(*args, **kwargs)
43
+
44
+ return wrapper
45
+
46
+ return decorator
47
+
48
+ @staticmethod
49
+ def parameters(params: Dict[str, Dict[str, Any]]) -> Callable:
50
+ """Decorator to define parameters for a tool function."""
51
+
52
+ def decorator(func: Callable) -> Callable:
53
+ if func.__name__ not in Agent._tools_registry:
54
+ Agent._tools_registry[func.__name__] = {}
55
+ Agent._tools_registry[func.__name__]["parameters_def"] = params
56
+ if "signature" not in Agent._tools_registry[func.__name__]:
57
+ Agent._tools_registry[func.__name__]["signature"] = inspect.signature(func)
58
+ if "function_ref" not in Agent._tools_registry[func.__name__]:
59
+ Agent._tools_registry[func.__name__]["function_ref"] = func
60
+ Agent._tools_registry[func.__name__]["is_method"] = inspect.ismethod(func)
61
+
62
+ @wraps(func)
63
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
64
+ return func(*args, **kwargs)
65
+
66
+ return wrapper
67
+
68
+ return decorator
69
+
70
+ def __init__(
71
+ self, api_key: str, tools: Optional[List[Callable[..., Any]]] = None, model_name: str = "gemini-1.5-flash"
72
+ ) -> None:
73
+ """
74
+ Initializes the Agent using REST API calls.
75
+
76
+ Args:
77
+ api_key: Your Google Generative AI API key.
78
+ tools: A list of Python functions or class methods decorated as tools.
79
+ model_name: The name of the Gemini model to use.
80
+ """
81
+ if not api_key:
82
+ raise ValueError("API key is required.")
83
+ self.api_key = api_key
84
+ self.model_name = model_name
85
+ self.base_url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}"
86
+ self.headers = {"Content-Type": "application/json"}
87
+
88
+ self._registered_tools_json: List[Dict[str, Any]] = [] # Store JSON representation
89
+ self._tool_functions: Dict[str, Callable[..., Any]] = {} # Map name to actual function
90
+ self._tool_instances: Dict[str, Any] = {} # Store instances for class methods
91
+ self._intermediate_results: Dict[str, Any] = {} # Store intermediate results
92
+ self._stored_variables: Dict[str, Dict[str, Any]] = {} # Store variables with metadata
93
+
94
+ if tools:
95
+ self._process_tools(tools)
96
+
97
+ def _process_tools(self, tools: List[Callable[..., Any]]) -> None:
98
+ """Converts decorated Python functions into the JSON format for the REST API."""
99
+ for func in tools:
100
+ tool_name = func.__name__
101
+ if tool_name not in Agent._tools_registry:
102
+ print(
103
+ f"Warning: Function '{tool_name}' was passed "
104
+ "but has no @Agent decorators. Skipping."
105
+ )
106
+ continue
107
+
108
+ metadata = Agent._tools_registry[tool_name]
109
+ if "description" not in metadata:
110
+ print(f"Warning: Function '{tool_name}' is missing @Agent.description. Skipping.")
111
+ continue
112
+
113
+ # Store the bound method directly if it's a class method
114
+ if inspect.ismethod(func):
115
+ self._tool_functions[tool_name] = func
116
+ else:
117
+ self._tool_functions[tool_name] = metadata["function_ref"]
118
+
119
+ # Build the parameters schema JSON
120
+ gemini_params_schema = {"type": "OBJECT", "properties": {}, "required": []}
121
+ params_def = metadata.get("parameters_def", {})
122
+ signature = metadata.get("signature") # inspect.signature object
123
+
124
+ if not params_def and signature:
125
+ params_def = {}
126
+ for name, param in signature.parameters.items():
127
+ # Skip 'self' parameter for class methods
128
+ if name == "self" and inspect.ismethod(func):
129
+ continue
130
+ py_type = (
131
+ param.annotation if param.annotation != inspect.Parameter.empty else str
132
+ )
133
+ params_def[name] = {"type": py_type, "description": f"Parameter {name}"}
134
+
135
+ for name, definition in params_def.items():
136
+ py_type = definition.get("type", str)
137
+ gemini_type = self.get_gemini_type(py_type)
138
+ gemini_params_schema["properties"][name] = {
139
+ "type": gemini_type,
140
+ "description": definition.get("description", ""),
141
+ }
142
+ if signature and signature.parameters[name].default == inspect.Parameter.empty:
143
+ gemini_params_schema["required"].append(name)
144
+
145
+ # Create the Function Declaration JSON dictionary
146
+ declaration_json = {
147
+ "name": tool_name,
148
+ "description": metadata["description"],
149
+ "parameters": gemini_params_schema if gemini_params_schema["properties"] else None,
150
+ }
151
+ if declaration_json["parameters"] is None:
152
+ del declaration_json["parameters"]
153
+
154
+ self._registered_tools_json.append(declaration_json)
155
+
156
+ def set_variable(
157
+ self, name: str, value: Any, description: str = "", type_hint: Optional[Type] = None
158
+ ) -> str:
159
+ """
160
+ Stores a variable in the agent's memory with metadata.
161
+ If a variable with the same name exists, creates a new variable with a counter suffix.
162
+
163
+ Args:
164
+ name: The name of the variable
165
+ value: The actual value to store
166
+ description: A description of what the variable represents
167
+ type_hint: Optional type hint for the variable
168
+
169
+ Returns:
170
+ The name of the stored variable
171
+ """
172
+ # Check if the base name exists
173
+ if name in self._stored_variables:
174
+ # Find all variables that start with the base name
175
+ existing_vars = [
176
+ var_name
177
+ for var_name in self._stored_variables.keys()
178
+ if var_name.startswith(name + "_") or var_name == name
179
+ ]
180
+
181
+ # Find the highest counter used
182
+ max_counter = 0
183
+ for var_name in existing_vars:
184
+ if var_name == name:
185
+ max_counter = max(max_counter, 1)
186
+ else:
187
+ try:
188
+ counter = int(var_name.split("_")[-1])
189
+ max_counter = max(max_counter, counter)
190
+ except ValueError:
191
+ continue
192
+
193
+ # Create new name with incremented counter
194
+ new_name = f"{name}_{max_counter + 1}"
195
+ print(f"Variable '{name}' already exists. Creating new variable '{new_name}'")
196
+ name = new_name
197
+
198
+ self._stored_variables[name] = {
199
+ "value": value,
200
+ "description": description,
201
+ "type": type_hint or type(value).__name__,
202
+ "created_at": datetime.now().isoformat(),
203
+ }
204
+
205
+ return name
206
+
207
+ def get_variable(self, name: str) -> Any:
208
+ """
209
+ Retrieves a stored variable's value.
210
+
211
+ Args:
212
+ name: The name of the variable to retrieve
213
+
214
+ Returns:
215
+ The stored value or None if not found
216
+ """
217
+ return self._stored_variables.get(name, {}).get("value")
218
+
219
+ def list_variables(self) -> Dict[str, Dict[str, Any]]:
220
+ """
221
+ Returns information about all stored variables.
222
+
223
+ Returns:
224
+ Dictionary of variable names to their metadata
225
+ """
226
+ return {
227
+ name: {k: v for k, v in data.items() if k != "value"}
228
+ for name, data in self._stored_variables.items()
229
+ }
230
+
231
+ def _get_system_prompt(self) -> str:
232
+ """Returns a system prompt that guides the model in breaking down complex operations."""
233
+ variables_info = "\n".join(
234
+ [
235
+ f"- {name}: {data['description']} (Type: {data['type']})"
236
+ for name, data in self._stored_variables.items()
237
+ ]
238
+ )
239
+
240
+ return """You are an AI assistant that can break down complex tasks into sequential steps using available tools.
241
+ When faced with a complex request:
242
+ 1. Analyze the request to identify which tools can be used
243
+ 2. Break down the request into sequential steps
244
+ 3. For each step:
245
+ - Use the most appropriate tool
246
+ - Store the result for use in subsequent steps
247
+ 4. Combine the results to provide a final answer
248
+
249
+ Available tools:
250
+ {tools_list}
251
+
252
+ Available variables:
253
+ {variables_list}
254
+
255
+ IMPORTANT - Variable Usage:
256
+ When you need to use a stored variable in a function call, you MUST use the following syntax:
257
+ - For function arguments: {{"variable": "variable_name"}}
258
+ - For example, if you want to use the 'current_user' variable in a function call:
259
+ {{"user_id": {{"variable": "current_user"}}}}
260
+
261
+ Remember:
262
+ - Always perform one operation at a time
263
+ - Use intermediate results from previous steps
264
+ - If a step requires multiple tools, execute them sequentially
265
+ - If you're unsure about the next step, explain your reasoning
266
+ - You can use both stored variables and values from the prompt
267
+ - When using stored variables, ALWAYS use the {{"variable": "variable_name"}} syntax
268
+ """.format(
269
+ tools_list="\n".join(
270
+ [
271
+ f"- {name}: {desc}"
272
+ for name, desc in [
273
+ (tool["name"], tool["description"]) for tool in self._registered_tools_json
274
+ ]
275
+ ]
276
+ ),
277
+ variables_list=variables_info,
278
+ )
279
+
280
+ def _substitute_variables(self, args: Dict[str, Any]) -> Dict[str, Any]:
281
+ """Substitutes variable references in arguments with their actual values."""
282
+ result = {}
283
+ for key, value in args.items():
284
+ if isinstance(value, str) and value.startswith("$"):
285
+ var_name = value[1:]
286
+ if var_name in self._stored_variables:
287
+ result[key] = self._stored_variables[var_name]["value"]
288
+ else:
289
+ result[key] = value
290
+ else:
291
+ result[key] = value
292
+ return result
293
+
294
+ def _call_gemini_api(self, payload: Dict[str, Any]) -> Dict[str, Any]:
295
+ """Makes a call to the Gemini API."""
296
+ response = requests.post(
297
+ f"{self.base_url}:generateContent?key={self.api_key}",
298
+ headers=self.headers,
299
+ json=payload,
300
+ )
301
+ response.raise_for_status()
302
+ return response.json()
303
+
304
+ def prompt(
305
+ self,
306
+ user_prompt: str,
307
+ system_prompt: Optional[str] = None,
308
+ response_structure: Optional[Dict[str, Any]] = None,
309
+ conversation_history: Optional[List[Dict[str, Any]]] = None,
310
+ ) -> Any:
311
+ """
312
+ Sends a prompt to the Gemini model and processes the response.
313
+
314
+ Args:
315
+ user_prompt: The user's input prompt
316
+ system_prompt: Optional system prompt to override the default
317
+ response_structure: Optional structure to enforce on the response
318
+ conversation_history: Optional list of previous conversation turns
319
+
320
+ Returns:
321
+ The model's response, processed according to the response structure if provided
322
+ """
323
+ self._intermediate_results = {}
324
+
325
+ if not system_prompt:
326
+ system_prompt = self._get_system_prompt()
327
+ else:
328
+ system_prompt = self._get_system_prompt() + system_prompt
329
+
330
+ current_contents = conversation_history if conversation_history else []
331
+ if system_prompt and not current_contents:
332
+ current_contents.append({"role": "user", "parts": [{"text": system_prompt}]})
333
+ current_contents.append(
334
+ {
335
+ "role": "model",
336
+ "parts": [
337
+ {
338
+ "text": "I understand I should break down complex tasks into sequential steps using the available tools and variables."
339
+ }
340
+ ],
341
+ }
342
+ )
343
+
344
+ current_contents.append({"role": "user", "parts": [{"text": user_prompt}]})
345
+ payload: Dict[str, Any] = {"contents": current_contents}
346
+
347
+ if self._registered_tools_json:
348
+ payload["tools"] = [{"functionDeclarations": self._registered_tools_json}]
349
+ payload["toolConfig"] = {"functionCallingConfig": {"mode": "AUTO"}}
350
+
351
+ apply_structure_later = bool(response_structure) and bool(self._registered_tools_json)
352
+ final_response_schema = None
353
+ final_mime_type = None
354
+
355
+ if response_structure and not self._registered_tools_json:
356
+ apply_structure_later = False
357
+ # If response_structure is a string type, make it more flexible
358
+ if response_structure.get("type") == "string":
359
+ response_structure = {
360
+ "type": ["string", "object"],
361
+ "properties": {"value": {"type": "string"}},
362
+ }
363
+ payload["generationConfig"] = {
364
+ "response_mime_type": "application/json",
365
+ "response_schema": response_structure,
366
+ }
367
+ final_mime_type = "application/json"
368
+ final_response_schema = response_structure
369
+
370
+ while True:
371
+
372
+ response_data = self._call_gemini_api(payload)
373
+ if "error" in response_data:
374
+ print(
375
+ f"API call failed: {response_data['error'].get('message', 'Unknown API error')}"
376
+ )
377
+ return response_data
378
+
379
+ if not response_data.get("candidates"):
380
+ feedback = response_data.get("promptFeedback")
381
+ block_reason = feedback.get("blockReason") if feedback else "Unknown"
382
+ safety_ratings = feedback.get("safetyRatings") if feedback else []
383
+ error_msg = f"Request blocked by API. Reason: {block_reason}."
384
+ if safety_ratings:
385
+ error_msg += f" Details: {json.dumps(safety_ratings)}"
386
+ print(error_msg)
387
+ return {"error": {"message": error_msg, "details": feedback}}
388
+
389
+ try:
390
+ candidate = response_data["candidates"][0]
391
+ content = candidate["content"]
392
+
393
+ for part in content["parts"]:
394
+ payload["contents"].append({"role": "model", "parts": [part]})
395
+
396
+ if "functionCall" in part:
397
+ fc = part["functionCall"]
398
+ tool_name = fc["name"]
399
+ args = fc.get("args", {})
400
+
401
+ if tool_name not in self._tool_functions:
402
+ error_msg = f"Model attempted to call unknown function '{tool_name}'."
403
+ print(f"Error: {error_msg}")
404
+ error_response_part = {
405
+ "functionResponse": {
406
+ "name": tool_name,
407
+ "response": {"error": error_msg},
408
+ }
409
+ }
410
+ payload["contents"].append(
411
+ {"role": "user", "parts": [error_response_part]}
412
+ )
413
+ continue
414
+
415
+ try:
416
+ tool_function = self._tool_functions[tool_name]
417
+ print(f"--- Calling Function: {tool_name}({args}) ---")
418
+
419
+ # Substitute both stored variables and intermediate results
420
+ args = self._substitute_variables(args)
421
+ for key, value in args.items():
422
+ if isinstance(value, str) and value.startswith("$"):
423
+ result_key = value[1:]
424
+ if result_key in self._intermediate_results:
425
+ args[key] = self._intermediate_results[result_key]
426
+
427
+ # Call the function directly - it's already bound if it's a method
428
+ function_result = tool_function(**args)
429
+
430
+ print(f"--- Function Result: {function_result} ---")
431
+
432
+ result_key = f"result_{len(self._intermediate_results)}"
433
+ self._intermediate_results[result_key] = function_result
434
+
435
+ varaible_name = self.set_variable(
436
+ result_key,
437
+ function_result,
438
+ "the result of function call with name {tool_name} and arguments {args}",
439
+ )
440
+ function_response_part = {
441
+ "functionResponse": {
442
+ "name": tool_name,
443
+ "response": {
444
+ "content": function_result,
445
+ "key": varaible_name,
446
+ "content_type": type(function_result).__name__,
447
+ },
448
+ }
449
+ }
450
+
451
+ payload["contents"].append(
452
+ {
453
+ "role": "user",
454
+ "parts": [
455
+ {
456
+ "text": f"the return value of the function stored in the variable {varaible_name}"
457
+ }
458
+ ],
459
+ }
460
+ )
461
+
462
+ payload["contents"].append(
463
+ {"role": "user", "parts": [function_response_part]}
464
+ )
465
+
466
+ except Exception as e:
467
+ print(f"Error executing function {tool_name}: {e}")
468
+ error_msg = f"Error during execution of tool '{tool_name}': {e}"
469
+ error_response_part = {
470
+ "functionResponse": {
471
+ "name": tool_name,
472
+ "response": {"error": error_msg},
473
+ }
474
+ }
475
+ payload["contents"].append(
476
+ {"role": "user", "parts": [error_response_part]}
477
+ )
478
+ continue
479
+
480
+ elif "text" in part:
481
+ final_text = part["text"]
482
+
483
+ if final_mime_type == "application/json" and final_response_schema:
484
+ try:
485
+ structured_output = json.loads(final_text)
486
+ if not any(
487
+ "functionCall" in p
488
+ for p in content["parts"][content["parts"].index(part) + 1 :]
489
+ ):
490
+ return structured_output
491
+ except json.JSONDecodeError as e:
492
+ print(
493
+ f"Warning: Failed to parse initially structured output: {e}. Continuing with raw text."
494
+ )
495
+
496
+ elif apply_structure_later:
497
+ if not any(
498
+ "functionCall" in p
499
+ for p in content["parts"][content["parts"].index(part) + 1 :]
500
+ ):
501
+ print("--- Attempting final structuring call ---")
502
+ formatting_payload = {
503
+ "contents": [
504
+ {
505
+ "role": "user",
506
+ "parts": [
507
+ {
508
+ "text": f"Please format the following information according to the requested JSON structure:\n\n{final_text}"
509
+ }
510
+ ],
511
+ }
512
+ ],
513
+ "generationConfig": {
514
+ "response_mime_type": "application/json",
515
+ "response_schema": response_structure,
516
+ },
517
+ }
518
+ structured_response_data = self._call_gemini_api(formatting_payload)
519
+
520
+ if "error" in structured_response_data:
521
+ print(
522
+ f"Structuring call failed: {structured_response_data['error']}. Returning intermediate text."
523
+ )
524
+ return final_text
525
+
526
+ try:
527
+ structured_text = structured_response_data["candidates"][0][
528
+ "content"
529
+ ]["parts"][0]["text"]
530
+ structured_output = json.loads(structured_text)
531
+ return structured_output
532
+ except (KeyError, IndexError, json.JSONDecodeError) as e:
533
+ print(
534
+ f"Warning: Failed to parse structured output after formatting call: {e}. Returning intermediate text."
535
+ )
536
+ return final_text
537
+
538
+ elif not any(
539
+ "functionCall" in p
540
+ for p in content["parts"][content["parts"].index(part) + 1 :]
541
+ ):
542
+ if response_structure and not apply_structure_later:
543
+ print("--- Attempting final structuring call ---")
544
+ formatting_payload = {
545
+ "contents": [
546
+ {
547
+ "role": "user",
548
+ "parts": [
549
+ {
550
+ "text": f"Please format the following information according to the requested JSON structure:\n\n{final_text}"
551
+ }
552
+ ],
553
+ }
554
+ ],
555
+ "generationConfig": {
556
+ "response_mime_type": "application/json",
557
+ "response_schema": response_structure,
558
+ },
559
+ }
560
+ structured_response_data = self._call_gemini_api(formatting_payload)
561
+
562
+ if "error" in structured_response_data:
563
+ print(
564
+ f"Structuring call failed: {structured_response_data['error']}. Returning intermediate text."
565
+ )
566
+ return final_text
567
+
568
+ try:
569
+ structured_text = structured_response_data["candidates"][0][
570
+ "content"
571
+ ]["parts"][0]["text"]
572
+ structured_output = json.loads(structured_text)
573
+ return structured_output
574
+ except (KeyError, IndexError, json.JSONDecodeError) as e:
575
+ print(
576
+ f"Warning: Failed to parse structured output after formatting call: {e}. Returning intermediate text."
577
+ )
578
+ return final_text
579
+ return final_text
580
+
581
+ continue
582
+
583
+ except (KeyError, IndexError) as e:
584
+ print(f"Error parsing API response structure: {e}. Response: {response_data}")
585
+ return {
586
+ "error": {
587
+ "message": f"Error parsing API response: {e}",
588
+ "details": response_data,
589
+ }
590
+ }
591
+
592
+ return {"error": {"message": "Exited interaction loop unexpectedly."}}
593
+ def test_agent(self) -> None:
594
+ print("Testing agent...")