Kunal Pai commited on
Commit
1b3d55c
·
1 Parent(s): 4a2efcf

Refactor GroqAgent initialization and enhance error handling; update AgentCostManager with Groq model costs

Browse files
src/manager/agent_manager.py CHANGED
@@ -147,43 +147,72 @@ class GroqAgent(Agent):
147
  def __init__(
148
  self,
149
  agent_name: str,
150
- base_model: str = "llama-3.3-70b-versatile",
151
- system_prompt: str = "system.prompt",
 
 
 
 
152
  ):
153
- self.agent_name = agent_name
154
- self.base_model = base_model
155
- # load API key from environment
 
 
 
156
  api_key = os.getenv("GROQ_API_KEY")
 
 
157
  self.client = Groq(api_key=api_key)
158
- # read system prompt content
159
- with open(system_prompt, 'r') as f:
160
- self.system_instruction = f.read()
 
 
 
 
 
 
161
 
162
  def create_model(self) -> None:
163
- # Groq models are available by name; no creation step
 
 
 
 
164
  pass
165
 
166
  def ask_agent(self, prompt: str) -> str:
 
 
 
 
 
 
167
  messages = [
168
- {"role": "system", "content": self.system_instruction},
169
  {"role": "user", "content": prompt},
170
  ]
171
- response = self.client.chat.completions.create(
172
- messages=messages,
173
- model=self.base_model,
174
- )
175
- result = response.choices[0].message.content
176
- print(result)
177
- return result
 
 
 
 
178
 
179
  def delete_agent(self) -> None:
180
- # No delete support for Groq
181
  pass
182
 
183
- def get_type(self):
 
184
  return self.type
185
 
186
-
187
  @singleton
188
  class AgentManager():
189
  budget_manager: BudgetManager = BudgetManager()
 
147
  def __init__(
148
  self,
149
  agent_name: str,
150
+ base_model: str,
151
+ system_prompt: str,
152
+ create_resource_cost: int,
153
+ invoke_resource_cost: int,
154
+ create_expense_cost: int = 0,
155
+ invoke_expense_cost: int = 0,
156
  ):
157
+ # Call the parent class constructor first
158
+ super().__init__(agent_name, base_model, system_prompt,
159
+ create_resource_cost, invoke_resource_cost,
160
+ create_expense_cost, invoke_expense_cost)
161
+
162
+ # Groq-specific API client setup
163
  api_key = os.getenv("GROQ_API_KEY")
164
+ if not api_key:
165
+ raise ValueError("GROQ_API_KEY environment variable not set. Please set it in your .env file or environment.")
166
  self.client = Groq(api_key=api_key)
167
+
168
+ if self.base_model and "groq-" in self.base_model:
169
+ self.groq_api_model_name = self.base_model.split("groq-", 1)[1]
170
+ else:
171
+ # Fallback or error if the naming convention isn't followed.
172
+ # This ensures that if a non-prefixed model name is somehow passed,
173
+ # it might still work, or you can raise an error.
174
+ self.groq_api_model_name = self.base_model
175
+ print(f"Warning: GroqAgent base_model '{self.base_model}' does not follow 'groq-' prefix convention.")
176
 
177
  def create_model(self) -> None:
178
+ """
179
+ Create and Initialize agent.
180
+ For Groq, models are pre-existing on their cloud.
181
+ This method is called by Agent's __init__.
182
+ """
183
  pass
184
 
185
  def ask_agent(self, prompt: str) -> str:
186
+ """Ask agent a question"""
187
+ if not self.client:
188
+ raise ConnectionError("Groq client not initialized. Check API key and constructor.")
189
+ if not self.groq_api_model_name:
190
+ raise ValueError("Groq API model name not set. Check base_model configuration.")
191
+
192
  messages = [
193
+ {"role": "system", "content": self.system_prompt},
194
  {"role": "user", "content": prompt},
195
  ]
196
+ try:
197
+ response = self.client.chat.completions.create(
198
+ messages=messages,
199
+ model=self.groq_api_model_name, # Use the derived model name for Groq API
200
+ )
201
+ result = response.choices[0].message.content
202
+ return result
203
+ except Exception as e:
204
+ # Handle API errors or other exceptions during the call
205
+ print(f"Error calling Groq API: {e}")
206
+ raise # Re-raise the exception or handle it as appropriate
207
 
208
  def delete_agent(self) -> None:
209
+ """Delete agent"""
210
  pass
211
 
212
+ def get_type(self) -> str: # Ensure return type hint matches Agent ABC
213
+ """Get agent type"""
214
  return self.type
215
 
 
216
  @singleton
217
  class AgentManager():
218
  budget_manager: BudgetManager = BudgetManager()
src/tools/default_tools/agent_cost_manager.py CHANGED
@@ -60,6 +60,11 @@ class AgentCostManager():
60
  "create_expense_cost": 0,
61
  "invoke_expense_cost": 0.0375,
62
  },
 
 
 
 
 
63
  }
64
 
65
  def get_costs(self):
 
60
  "create_expense_cost": 0,
61
  "invoke_expense_cost": 0.0375,
62
  },
63
+ "groq-qwen-qwq-32b": {
64
+ "description": "Avg Accuracy: 60.0%, 70.0% on multi-task understanding, 80.0% on Math",
65
+ "create_expense_cost": 0,
66
+ "invoke_expense_cost": 0.05,
67
+ },
68
  }
69
 
70
  def get_costs(self):