agent.py

import os
from tools import all_tools
from dotenv import load_dotenv
from smolagents import CodeAgent, PromptTemplates, LiteLLMModel
import logfire
import litellm

load_dotenv()

logfire.configure(token=os.getenv("LOGFIRE_TOKEN"))
litellm.callbacks = ["logfire"]
litellm.success_callback = ["logfire"]

model = LiteLLMModel(
    model_id="openrouter/openai/gpt-4.1-mini",
    api_key=os.getenv("OPENROUTER_API_KEY"),
    max_tokens=5045
)

# Loads system prompt from the file
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    prompts = PromptTemplates(system_prompt = f.read())

agent = CodeAgent(
    model=model,
    tools=all_tools,
    #prompt_templates=prompts,
    additional_authorized_imports=['numpy',
                                   'pandas',
                                   'subprocess',
                                   'scipy',
                                   'csv'],
    add_base_tools=True
)

class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        self.agent = agent 
    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        fixed_answer = self.agent.run(question)
        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer