h2oaimichalmarszalek's picture
redesign
d3b823a
import os
from pathlib import Path
from typing import Optional
from smolagents import CodeAgent, PythonInterpreterTool, WikipediaSearchTool, VisitWebpageTool, FinalAnswerTool
from tools.utils import reverse_string, process_excel_file, is_text_file, execute_python_file, get_ingredients
from tools.youtube import load_youtube
from tools.audio import transcribe_audio
from tools.web import optimized_web_search
# from langchain.agents import load_tools
# model = LiteLLMModel(
# model_id="ollama/qwen2.5:7b",
# api_base="http://localhost:11434"
# )
class BasicAgent:
def __init__(self, model):
self._model = model
self._agent = CodeAgent(
tools=[PythonInterpreterTool(), WikipediaSearchTool(), VisitWebpageTool(), FinalAnswerTool(),optimized_web_search, reverse_string, process_excel_file, is_text_file, load_youtube, execute_python_file, transcribe_audio, get_ingredients],
additional_authorized_imports=['*', 'subprocess','markdownify', 'chess', 'random', 'time', 'itertools', 'pandas', 'webbrowser', 'requests', 'beautifulsoup4', 'csv', 'openpyxl', 'json', 'yaml'],
model=model,
add_base_tools=True,
max_steps=10
)
print("BasicAgent initialized.")
def __call__(self, question: str, file_path: Optional[str] = None) -> str:
# print(f"Agent received question (first 50 chars): {question[:50]}...")
if file_path and os.path.exists(file_path):
try:
file = Path(file_path)
if is_text_file(file_path):
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
file_content = f.read()
if file.suffix == '.py':
question += f"\nAttached Python code:\n{file_path}"
else:
question += f"\nAttached File Content:\n{file_content}"
else:
if file.suffix == '.mp3':
question += f"\nUse tool transcribe_audio to extract text from mp3 file.\nAttached mp3 file to transcribes:\n{file_path}"
else:
question += f"\nAttached File Path:\n{file_path}"
except Exception as e:
print(f"Failed to read file: {e}")
print(question)
answer = self._agent.run(question)
print(f"Agent returning answer: {answer}")
return answer