Final_Assignment_Template / geminiAgent.py
ninooo96's picture
bug fixing
02121c6
raw
history blame
2.18 kB
import os
import time
import operator
from smolagents import CodeAgent,DuckDuckGoSearchTool, tool, WikipediaSearchTool, LiteLLMModel
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool
import pandas as pd
import pytesseract
import openpyxl
from PIL import Image
def create_agent():
return CodeAgent(
model=LiteLLMModel(model_id="gemini/gemini-2.0-flash-lite-001", api_key=os.getenv("GEMINI_KEY")),
tools=[DuckDuckGoSearchTool(), WikipediaSearchTool(), VisitWebpageTool(), ocr_tool, read_csv, read_excel],
add_base_tools=True,
additional_authorized_imports=['pandas','numpy','csv','subprocess', 'exec']
)
@tool
def ocr_tool(image: str) -> str:
"""
A tool that performs OCR processing on an image.
Args:
image: path for the image.
Returns:
Text extracted from the image.
"""
image = Image.open(image)
return pytesseract.image_to_string(image)
@tool
def read_csv(csv: str) -> str:
"""
Reads a CSV file and returns its contents as a human-readable string.
Args:
csv: path for the csv file.
Returns:
A string representation of the CSV file contents.
"""
csv_file = pd.read_csv(csv)
if csv_file.empty:
return "The CSV file is empty."
# Format the DataFrame as a string with clear headers and a separator
output = ""
output += "CSV Data:\n"
output += "--------------------------------\n"
output += csv_file.to_string(index=False)
output += "\n--------------------------------\n"
return output
@tool
def read_excel(excel_path: str) -> str:
"""
Reads an Excel file and returns its contents as a human-readable string.
Args:
excel_path: The path to the Excel file.
Returns:
A string representation of the Excel file contents.
"""
try:
df = pd.read_excel(excel_path)
return df.to_string()
except Exception as e:
return f"Error reading Excel file: {e}"