Kunal Pai commited on
Commit
fcdfb63
·
1 Parent(s): 2187cd8

Refactor LambdaAgent to use OpenAI client and update cost manager with new Lambda model expenses

Browse files
src/manager/agent_manager.py CHANGED
@@ -3,7 +3,7 @@ from typing import Dict, Type, Any, Optional, Tuple
3
  import os
4
  import json
5
  import ollama
6
- import requests
7
  from src.manager.utils.singleton import singleton
8
  from src.manager.utils.streamlit_interface import output_assistant_response
9
  from google import genai
@@ -233,15 +233,19 @@ class LambdaAgent(Agent):
233
  create_expense_cost: int = 0,
234
  invoke_expense_cost: int = 0,
235
  output_expense_cost: int = 0,
236
- lambda_url: str = "",
237
  api_key: str = ""):
238
- if not lambda_url:
239
- raise ValueError("Lambda URL must be provided for LambdaAgent.")
240
 
241
- self.lambda_url = lambda_url
242
  self.api_key = api_key or os.getenv("LAMBDA_API_KEY")
 
 
243
  if not self.api_key:
244
  raise ValueError("Lambda API key must be provided or set in LAMBDA_API_KEY environment variable.")
 
 
 
 
 
245
 
246
  super().__init__(agent_name,
247
  base_model,
@@ -256,20 +260,18 @@ class LambdaAgent(Agent):
256
  pass # Lambda already deployed
257
 
258
  def ask_agent(self, prompt: str) -> str:
 
259
  try:
260
- headers = {
261
- "Content-Type": "application/json",
262
- "x-api-key": self.api_key # Required by API Gateway if enabled
263
- }
264
- payload = {
265
- "prompt": prompt,
266
- "system_prompt": self.system_prompt
267
- }
268
- response = requests.post(self.lambda_url, headers=headers, json=payload)
269
- response.raise_for_status()
270
- return response.json().get("response", "")
271
  except Exception as e:
272
- print(f"Error calling Lambda agent: {e}")
273
  raise
274
 
275
  def delete_agent(self) -> None:
 
3
  import os
4
  import json
5
  import ollama
6
+ from openai import OpenAI
7
  from src.manager.utils.singleton import singleton
8
  from src.manager.utils.streamlit_interface import output_assistant_response
9
  from google import genai
 
233
  create_expense_cost: int = 0,
234
  invoke_expense_cost: int = 0,
235
  output_expense_cost: int = 0,
 
236
  api_key: str = ""):
 
 
237
 
238
+ self.lambda_url = "https://api.lambda.ai/v1"
239
  self.api_key = api_key or os.getenv("LAMBDA_API_KEY")
240
+
241
+ self.lambda_model = base_model.split("lambda-")[1] if base_model.startswith("lambda-") else base_model
242
  if not self.api_key:
243
  raise ValueError("Lambda API key must be provided or set in LAMBDA_API_KEY environment variable.")
244
+
245
+ self.client = client = OpenAI(
246
+ api_key=self.api_key,
247
+ base_url=self.lambda_url,
248
+ )
249
 
250
  super().__init__(agent_name,
251
  base_model,
 
260
  pass # Lambda already deployed
261
 
262
  def ask_agent(self, prompt: str) -> str:
263
+ """Ask agent a question"""
264
  try:
265
+ response = self.client.chat.completions.create(
266
+ model=self.lambda_model,
267
+ messages=[
268
+ {"role": "system", "content": self.system_prompt},
269
+ {"role": "user", "content": prompt}
270
+ ],
271
+ )
272
+ return response.choices[0].message.content
 
 
 
273
  except Exception as e:
274
+ output_assistant_response(f"Error asking agent: {e}")
275
  raise
276
 
277
  def delete_agent(self) -> None:
src/tools/default_tools/agent_cost_manager.py CHANGED
@@ -72,6 +72,12 @@ class AgentCostManager():
72
  "invoke_expense_cost": 0.29,
73
  "output_expense_cost": 0.39,
74
  },
 
 
 
 
 
 
75
  }
76
 
77
  def get_costs(self):
 
72
  "invoke_expense_cost": 0.29,
73
  "output_expense_cost": 0.39,
74
  },
75
+ "lambda-hermes3-8b": {
76
+ "description": "High volume and lower intelligence tasks, 60.0% on MMLU, 58.0% on MATH",
77
+ "create_expense_cost": 0,
78
+ "invoke_expense_cost": 0.025,
79
+ "output_expense_cost": 0.04,
80
+ },
81
  }
82
 
83
  def get_costs(self):