Update tools/ai_tools.py
Browse files- tools/ai_tools.py +20 -20
tools/ai_tools.py
CHANGED
@@ -223,27 +223,23 @@ class AITools():
|
|
223 |
except Exception as e:
|
224 |
raise RuntimeError(f"Processing failed: {str(e)}")
|
225 |
|
226 |
-
@tool("Code
|
227 |
-
def
|
228 |
-
"""Given a question and
|
229 |
-
|
230 |
Args:
|
231 |
question (str): Question to answer
|
232 |
-
|
233 |
-
|
234 |
Returns:
|
235 |
str: Answer to the question
|
236 |
-
|
237 |
Raises:
|
238 |
RuntimeError: If processing fails"""
|
239 |
try:
|
240 |
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
|
241 |
-
|
242 |
-
file = client.files.upload(file=file_path)
|
243 |
-
|
244 |
response = client.models.generate_content(
|
245 |
-
model=
|
246 |
-
contents=[
|
247 |
config=types.GenerateContentConfig(
|
248 |
tools=[types.Tool(code_execution=types.ToolCodeExecution)]
|
249 |
),
|
@@ -254,24 +250,28 @@ class AITools():
|
|
254 |
return part.code_execution_result.output
|
255 |
except Exception as e:
|
256 |
raise RuntimeError(f"Processing failed: {str(e)}")
|
|
|
|
|
|
|
|
|
257 |
|
258 |
-
@tool("Code Generation Tool")
|
259 |
-
def code_generation_tool(question: str, json_data: str) -> str:
|
260 |
-
"""Given a question and JSON data, generate and execute code to answer the question.
|
261 |
Args:
|
262 |
question (str): Question to answer
|
263 |
-
|
|
|
264 |
Returns:
|
265 |
str: Answer to the question
|
266 |
-
|
267 |
Raises:
|
268 |
RuntimeError: If processing fails"""
|
269 |
try:
|
270 |
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
|
271 |
-
|
|
|
|
|
272 |
response = client.models.generate_content(
|
273 |
-
model=
|
274 |
-
contents=[
|
275 |
config=types.GenerateContentConfig(
|
276 |
tools=[types.Tool(code_execution=types.ToolCodeExecution)]
|
277 |
),
|
|
|
223 |
except Exception as e:
|
224 |
raise RuntimeError(f"Processing failed: {str(e)}")
|
225 |
|
226 |
+
@tool("Code Generation Tool")
|
227 |
+
def code_generation_tool(question: str, json_data: str) -> str:
|
228 |
+
"""Given a question and JSON data, generate and execute code to answer the question.
|
|
|
229 |
Args:
|
230 |
question (str): Question to answer
|
231 |
+
file_path (str): The JSON data
|
|
|
232 |
Returns:
|
233 |
str: Answer to the question
|
234 |
+
|
235 |
Raises:
|
236 |
RuntimeError: If processing fails"""
|
237 |
try:
|
238 |
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
|
239 |
+
|
|
|
|
|
240 |
response = client.models.generate_content(
|
241 |
+
model=CODE_GENERATION_MODEL,
|
242 |
+
contents=[f"{question}\n{json_data}"],
|
243 |
config=types.GenerateContentConfig(
|
244 |
tools=[types.Tool(code_execution=types.ToolCodeExecution)]
|
245 |
),
|
|
|
250 |
return part.code_execution_result.output
|
251 |
except Exception as e:
|
252 |
raise RuntimeError(f"Processing failed: {str(e)}")
|
253 |
+
|
254 |
+
@tool("Code Execution Tool")
|
255 |
+
def code_execution_tool(question: str, file_path: str) -> str:
|
256 |
+
"""Given a question and Python file, execute the file to answer the question.
|
257 |
|
|
|
|
|
|
|
258 |
Args:
|
259 |
question (str): Question to answer
|
260 |
+
file_path (str): The Python file path
|
261 |
+
|
262 |
Returns:
|
263 |
str: Answer to the question
|
264 |
+
|
265 |
Raises:
|
266 |
RuntimeError: If processing fails"""
|
267 |
try:
|
268 |
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
|
269 |
+
|
270 |
+
file = client.files.upload(file=file_path)
|
271 |
+
|
272 |
response = client.models.generate_content(
|
273 |
+
model=CODE_EXECUTION_MODEL,
|
274 |
+
contents=[file, question],
|
275 |
config=types.GenerateContentConfig(
|
276 |
tools=[types.Tool(code_execution=types.ToolCodeExecution)]
|
277 |
),
|