iQuentin commited on
Commit
5c40ea2
·
verified ·
1 Parent(s): e5ab152

QAgent good first initialization for test

Browse files
Files changed (1) hide show
  1. agent.py +101 -22
agent.py CHANGED
@@ -1,23 +1,102 @@
1
  import os
2
- from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, Tool, tool, InferenceClientModel
3
-
4
-
5
- class QAgent(Tool):
6
- def __init__(self):
7
- print("BasicAgent initialized.")
8
- def __call__(self, question: str) -> str:
9
- print(f"Agent received question (first 50 chars): {question[:50]}...")
10
- # provider = "nebius"
11
- model = "Qwen/Qwen2.5-Coder-32B-Instruct"
12
- api_key = os.environ.get('SP_HF_TOK')
13
-
14
- agent = CodeAgent(tools=[DuckDuckGoSearchTool, VisitWebpageTool],
15
- model=InferenceClientModel(
16
- model=model,
17
- # provider=provider,
18
- api_key=api_key),
19
- add_base_tools=True)
20
-
21
- fixed_answer = agent.run(question)
22
- print(f"Agent returning fixed answer: {fixed_answer}")
23
- return fixed_answer
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
2
+ from smolagents import (
3
+ CodeAgent,
4
+ DuckDuckGoSearchTool,
5
+ VisitWebpageTool,
6
+ InferenceClientModel,
7
+ HfApiModel,
8
+ tool
9
+ )
10
+
11
+
12
+
13
+ class QAgent:
14
+ def __init__(
15
+ self,
16
+ model_type: str = "InferenceClientModel",
17
+ model_id: Optional[str] = None,
18
+ api_key: Optional[str] = None,
19
+ provider: Optional[str] = None, # for InferenceClientModel
20
+ timeout: Optional[int] = None, # for InferenceClientModel
21
+ system_prompt: Optional[str] = None,
22
+ verbose: bool = false # Verbose logging or not
23
+ ):
24
+ """
25
+ QAgent description
26
+ """
27
+
28
+ self.verbose = verbose
29
+ self.system_prompt = system_prompt
30
+
31
+
32
+ if model_type == "HfApiModel":
33
+ if api_key is None:
34
+ api_key = os.getenv("SP_HF_TOK")
35
+ if not api_key:
36
+ raise ValueError("No API Key found for HuggingFace. Please set SP_HF_TOK or pass api_key.")
37
+
38
+ if self.verbose:
39
+ print(f"Using Hugging Face token: {api_key[:5]}... (HfApiModel mode)")
40
+
41
+ self.model = HfApiModel(
42
+ model_id=model_id or "Qwen/Qwen2.5-Coder-32B-Instruct", # précédemment : or "meta-llama/Llama-3-70B-Instruct",
43
+ token=api_key
44
+ # temperature=temperature
45
+ )
46
+ elif model_type == "InferenceClientModel":
47
+ if api_key is None:
48
+ api_key = os.getenv("SP_HF_TOK")
49
+ if not api_key:
50
+ raise ValueError("No API Key found for HuggingFace. Please set SP_HF_TOK or pass api_key.")
51
+
52
+ if self.verbose:
53
+ print(f"Using Hugging Face token: {api_key[:5]}... (InferenceClientModel mode)")
54
+
55
+ self.model = InferenceClientModel(
56
+ model_id=model_id or "Qwen/Qwen2.5-Coder-32B-Instruct", # précédemment : or "meta-llama/Llama-3-70B-Instruct",
57
+ provider=provider or "hf-inference",
58
+ token=api_key,
59
+ timeout=timeout or 120
60
+ # temperature=temperature
61
+ )
62
+ else:
63
+ raise ValueError(f"Unknown model type: {model_type}")
64
+
65
+ if self.verbose:
66
+ print(f"Model initialized: {model_type} - {self.model.model_id} - prov: {self.model.provider}")
67
+
68
+ # Initialize tools
69
+ self.tools = [
70
+ DuckDuckGoSearchTool(),
71
+ PythonInterpreterTool(),
72
+ # save_and_read_file,
73
+ # download_file_from_url,
74
+ # analyze_csv_file,
75
+ # analyze_excel_file
76
+ ]
77
+
78
+ # Setup imports
79
+ self.imports = ["pandas", "numpy", "datetime", "json", "re", "math", "os", "requests", "csv", "urllib"]
80
+
81
+ # Create CodeAgent
82
+ self.agent = CodeAgent(
83
+ tools=self.tools,
84
+ model=self.model,
85
+ # additional_authorized_imports=self.imports,
86
+ # executor_type=executor_type,
87
+ # executor_kwargs=executor_kwargs,
88
+ verbosity_level=2 if self.verbose else 0
89
+ )
90
+
91
+ if self.verbose:
92
+ print("CodeAgent initialized")
93
+
94
+ def invoke(self, prompt: str) -> str:
95
+ print(f"Agent invoked with prompt: {prompt[:80]}...")
96
+
97
+ result = self.agent.run(prompt)
98
+ print(result)
99
+ )
100
+ return result
101
+
102
+