File size: 1,307 Bytes
4ffeb3d
694192c
a5c5a54
5403288
 
4d3e43e
5403288
a5c5a54
5403288
2b5fdbb
297ebdf
 
5403288
39cd766
c45c190
bdeca1f
597f2ad
39cd766
88b462f
 
 
450dce9
88b462f
a5c5a54
 
694192c
88511f2
38669c9
 
 
 
 
0466935
a5c5a54
5c2b2a9
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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