File size: 3,377 Bytes
932fded c531eac 932fded c531eac 932fded c531eac 932fded a225ae4 c531eac 9cf935f a225ae4 c531eac a225ae4 c531eac 932fded c531eac |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
from typing import Any, List, Optional
from smolagents import CodeAgent
from utils.logger import get_logger
logger = get_logger(__name__)
class Agent:
"""
Agent class that wraps a CodeAgent and provides a callable interface for answering questions.
Args:
model (Any): The language model to use.
tools (Optional[List[Any]]): List of tools to provide to the agent.
prompt (Optional[str]): Custom prompt template for the agent.
"""
def __init__(
self,
model: Any,
tools: Optional[List[Any]] = None,
prompt: Optional[str] = None,
):
logger.info("Initializing Agent")
self.model = model
self.tools = tools
self.imports = [
"pandas",
"numpy",
"os",
"requests",
"tempfile",
"datetime",
"json",
"time",
"re",
"openpyxl",
"pathlib",
"sys",
]
self.agent = CodeAgent(
model=self.model,
tools=self.tools,
add_base_tools=True,
additional_authorized_imports=self.imports,
)
self.prompt = prompt or (
"""
You are an advanced AI assistant specialized in solving complex, real-world tasks that require multi-step reasoning, factual accuracy, and use of external tools.
Follow these principles:
- Be precise and concise. The final answer must strictly match the required format with no extra commentary.
- Use tools intelligently. If a question involves external information, structured data, images, or audio, call the appropriate tool to retrieve or process it.
- Reason step-by-step. Think through the solution logically and plan your actions carefully before answering.
- Validate information. Always verify facts when possible instead of guessing.
- Use code if needed. For calculations, parsing, or transformations, generate Python code and execute it. But be careful, some questions contains time-consuming tasks, so you should be careful with the code you run. Better analyze the question and think about the best way to solve it.
- Don't forget to use `final_answer` to give the final answer.
- Use name of file ONLY FROM "FILE:" section. THIS IF ALWAYS A FILE.
IMPORTANT: When giving the final answer, output only the direct required result without any extra text like "Final Answer:" or explanations. YOU MUST RESPOND IN THE EXACT FORMAT AS THE QUESTION.
QUESTION: {question}
FILE: {context}
ANSWER:
"""
)
logger.info("Agent initialized")
def __call__(self, question: str, file_path: Optional[str] = None) -> str:
"""
Run the agent to answer a question, optionally using a file as context.
Args:
question (str): The question to answer.
file_path (Optional[str]): Path to a file to use as context (if any).
Returns:
str: The agent's answer as a string.
"""
answer = self.agent.run(
self.prompt.format(question=question, context=file_path)
)
answer = str(answer).strip("'").strip('"').strip()
return answer
|