sanjeed5 commited on
Commit
097f999
·
1 Parent(s): a2cbab6

Refactor agent.py to implement BasicAgent class with enhanced file handling capabilities and integrate smolagents.CodeAgent for improved task processing.

Browse files
Files changed (1) hide show
  1. agent.py +64 -10
agent.py CHANGED
@@ -4,19 +4,73 @@ from dotenv import load_dotenv
4
 
5
  load_dotenv()
6
 
7
- model = InferenceClientModel( token=os.getenv("HF_TOKEN")) # You can choose to not pass any model_id to InferenceClientModel to use a default model
 
8
  # you can also specify a particular provider e.g. provider="together" or provider="sambanova"
9
- agent = CodeAgent(tools=[WebSearchTool()], model=model, add_base_tools=True)
10
 
11
- agent.run(
12
- "Could you give me the 118th number in the Fibonacci sequence?",
 
 
 
 
 
13
  )
14
 
 
 
15
  class BasicAgent:
16
  def __init__(self):
17
- print("BasicAgent initialized.")
18
- def __call__(self, question: str) -> str:
19
- print(f"Agent received question (first 50 chars): {question[:50]}...")
20
- fixed_answer = "This is a default answer."
21
- print(f"Agent returning fixed answer: {fixed_answer}")
22
- return fixed_answer
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
 
5
  load_dotenv()
6
 
7
+ # Configure the underlying LLM model for the agent
8
+ smol_model = InferenceClientModel(token=os.getenv("HF_TOKEN")) # You can choose to not pass any model_id to InferenceClientModel to use a default model
9
  # you can also specify a particular provider e.g. provider="together" or provider="sambanova"
 
10
 
11
+ # Instantiate the CodeAgent with authorized imports for file operations
12
+ # WebSearchTool is kept as an example, you might want to add/remove tools as needed.
13
+ smol_code_agent_instance = CodeAgent(
14
+ tools=[WebSearchTool()],
15
+ model=smol_model,
16
+ add_base_tools=True,
17
+ additional_authorized_imports=["os", "io"] # Authorize os and io for file access
18
  )
19
 
20
+ # This is the agent class that app.py will instantiate and use.
21
+ # It wraps the smolagents.CodeAgent.
22
  class BasicAgent:
23
  def __init__(self):
24
+ self.smol_agent = smol_code_agent_instance
25
+ print("BasicAgent (using smolagents.CodeAgent) initialized.")
26
+
27
+ def __call__(self, question: str, file_path: str | None = None) -> str:
28
+ print(f"BasicAgent received question (first 100 chars): {question[:100]}...")
29
+
30
+ task_description = question
31
+ run_additional_args = {}
32
+
33
+ if file_path:
34
+ print(f"BasicAgent received file_path: {file_path}")
35
+ # Augment the task description to inform the agent about the file
36
+ # and how to access it via additional_args.
37
+ task_description += (
38
+ f"\n\nA file relevant to this task has been provided. "
39
+ f"You can access this file's content using standard Python file operations (e.g., open(), read()). "
40
+ f"The path to this file is available in the 'additional_args' dictionary, passed to the run method, under the key 'task_file'. "
41
+ f"The value is: '{file_path}'."
42
+ )
43
+ run_additional_args["additional_args"] = {"task_file": file_path}
44
+ print(f"Passing to smolagent: task='{task_description[:150]}...', additional_args={run_additional_args['additional_args']}")
45
+ else:
46
+ print("BasicAgent received no file_path.")
47
+ print(f"Passing to smolagent: task='{task_description[:150]}...'")
48
+
49
+ try:
50
+ # The CodeAgent's .run() method returns the final answer.
51
+ # Pass additional_args if they exist.
52
+ if run_additional_args:
53
+ answer = self.smol_agent.run(task_description, **run_additional_args)
54
+ else:
55
+ answer = self.smol_agent.run(task_description)
56
+
57
+ print(f"smolagents.CodeAgent returned: {str(answer)[:200]}...") # Log a snippet of the answer
58
+ return str(answer) # Ensure the output is a string
59
+ except Exception as e:
60
+ print(f"Error during smolagents.CodeAgent run: {e}")
61
+ # It's often good to return a more informative error message
62
+ return f"AGENT_ERROR: An error occurred while processing with CodeAgent: {str(e)}"
63
+
64
+ # Example test run (optional, can be commented out for deployment)
65
+ # if __name__ == '__main__':
66
+ # print("Testing CodeAgent with a simple task...")
67
+ # test_agent = BasicAgent()
68
+ # # Test without file
69
+ # # response = test_agent("What is the capital of France?")
70
+ # # print(f"Test response (no file): {response}")
71
+ # # To test with a file, you'd create a dummy file and pass its path:
72
+ # # with open("dummy_test_file.txt", "w") as f:
73
+ # # f.write("This is a test file for the agent.")
74
+ # # response_with_file = test_agent("Summarize the content of the provided file.", file_path="dummy_test_file.txt")
75
+ # # print(f"Test response (with file): {response_with_file}")
76
+ # # os.remove("dummy_test_file.txt")