sanjeed5's picture
Refactor agent.py to implement BasicAgent class with enhanced file handling capabilities and integrate smolagents.CodeAgent for improved task processing.
097f999
raw
history blame
3.83 kB
from smolagents import CodeAgent, InferenceClientModel, WebSearchTool
import os
from dotenv import load_dotenv
load_dotenv()
# Configure the underlying LLM model for the agent
smol_model = InferenceClientModel(token=os.getenv("HF_TOKEN")) # You can choose to not pass any model_id to InferenceClientModel to use a default model
# you can also specify a particular provider e.g. provider="together" or provider="sambanova"
# Instantiate the CodeAgent with authorized imports for file operations
# WebSearchTool is kept as an example, you might want to add/remove tools as needed.
smol_code_agent_instance = CodeAgent(
tools=[WebSearchTool()],
model=smol_model,
add_base_tools=True,
additional_authorized_imports=["os", "io"] # Authorize os and io for file access
)
# This is the agent class that app.py will instantiate and use.
# It wraps the smolagents.CodeAgent.
class BasicAgent:
def __init__(self):
self.smol_agent = smol_code_agent_instance
print("BasicAgent (using smolagents.CodeAgent) initialized.")
def __call__(self, question: str, file_path: str | None = None) -> str:
print(f"BasicAgent received question (first 100 chars): {question[:100]}...")
task_description = question
run_additional_args = {}
if file_path:
print(f"BasicAgent received file_path: {file_path}")
# Augment the task description to inform the agent about the file
# and how to access it via additional_args.
task_description += (
f"\n\nA file relevant to this task has been provided. "
f"You can access this file's content using standard Python file operations (e.g., open(), read()). "
f"The path to this file is available in the 'additional_args' dictionary, passed to the run method, under the key 'task_file'. "
f"The value is: '{file_path}'."
)
run_additional_args["additional_args"] = {"task_file": file_path}
print(f"Passing to smolagent: task='{task_description[:150]}...', additional_args={run_additional_args['additional_args']}")
else:
print("BasicAgent received no file_path.")
print(f"Passing to smolagent: task='{task_description[:150]}...'")
try:
# The CodeAgent's .run() method returns the final answer.
# Pass additional_args if they exist.
if run_additional_args:
answer = self.smol_agent.run(task_description, **run_additional_args)
else:
answer = self.smol_agent.run(task_description)
print(f"smolagents.CodeAgent returned: {str(answer)[:200]}...") # Log a snippet of the answer
return str(answer) # Ensure the output is a string
except Exception as e:
print(f"Error during smolagents.CodeAgent run: {e}")
# It's often good to return a more informative error message
return f"AGENT_ERROR: An error occurred while processing with CodeAgent: {str(e)}"
# Example test run (optional, can be commented out for deployment)
# if __name__ == '__main__':
# print("Testing CodeAgent with a simple task...")
# test_agent = BasicAgent()
# # Test without file
# # response = test_agent("What is the capital of France?")
# # print(f"Test response (no file): {response}")
# # To test with a file, you'd create a dummy file and pass its path:
# # with open("dummy_test_file.txt", "w") as f:
# # f.write("This is a test file for the agent.")
# # response_with_file = test_agent("Summarize the content of the provided file.", file_path="dummy_test_file.txt")
# # print(f"Test response (with file): {response_with_file}")
# # os.remove("dummy_test_file.txt")