helloparthshah commited on
Commit
a2ebc4b
·
1 Parent(s): bc17645

include cost per 1000 tokes

Browse files
Files changed (1) hide show
  1. src/manager/agent_manager.py +7 -3
src/manager/agent_manager.py CHANGED
@@ -361,17 +361,21 @@ class AgentManager():
361
  agent: Agent = self.get_agent(agent_name)
362
  print(agent.get_type())
363
  print(agent_name)
364
- print(self.is_local_invocation_enabled, self.is_cloud_invocation_enabled)
 
365
  if not self.is_local_invocation_enabled and agent.get_type() == "local":
366
  raise ValueError("Local invocation mode is disabled.")
367
 
368
  if not self.is_cloud_invocation_enabled and agent.get_type() == "cloud":
369
  raise ValueError("Cloud invocation mode is disabled.")
370
 
 
 
371
  self.validate_budget(agent.invoke_resource_cost,
372
- agent.invoke_expense_cost)
373
 
374
- self.budget_manager.add_to_expense_budget(agent.invoke_expense_cost)
 
375
 
376
  response = agent.ask_agent(prompt)
377
  return (response,
 
361
  agent: Agent = self.get_agent(agent_name)
362
  print(agent.get_type())
363
  print(agent_name)
364
+ print(self.is_local_invocation_enabled,
365
+ self.is_cloud_invocation_enabled)
366
  if not self.is_local_invocation_enabled and agent.get_type() == "local":
367
  raise ValueError("Local invocation mode is disabled.")
368
 
369
  if not self.is_cloud_invocation_enabled and agent.get_type() == "cloud":
370
  raise ValueError("Cloud invocation mode is disabled.")
371
 
372
+ n_tokens = len(prompt.split())/1000
373
+
374
  self.validate_budget(agent.invoke_resource_cost,
375
+ agent.invoke_expense_cost*n_tokens)
376
 
377
+ self.budget_manager.add_to_expense_budget(
378
+ agent.invoke_expense_cost*n_tokens)
379
 
380
  response = agent.ask_agent(prompt)
381
  return (response,