Spaces:
Running
Running
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
|
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 =
|
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 |
-
|
261 |
-
|
262 |
-
|
263 |
-
|
264 |
-
|
265 |
-
|
266 |
-
|
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 |
-
|
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):
|