mirror of
https://github.com/aljazceru/Auto-GPT.git
synced 2026-01-25 00:44:32 +01:00
6e6e7fcc9a
* Extract open ai api calls and retry at lowest level * Forgot a test * Gotta fix my local docker config so I can let pre-commit hooks run, ugh * fix: merge artiface * Fix linting * Update memory.vector.utils * feat: make sure resp exists * fix: raise error message if created * feat: rename file * fix: partial test fix * fix: update comments * fix: linting * fix: remove broken test * fix: require a model to exist * fix: BaseError issue * fix: runtime error * Fix mock response in test_make_agent * add 429 as errors to retry --------- Co-authored-by: k-boikov <64261260+k-boikov@users.noreply.github.com> Co-authored-by: Nicholas Tindle <nick@ntindle.com> Co-authored-by: Reinier van der Leer <github@pwuts.nl> Co-authored-by: Nicholas Tindle <nicktindle@outlook.com> Co-authored-by: Luke K (pr-0f3t) <2609441+lc0rp@users.noreply.github.com> Co-authored-by: Merwane Hamadi <merwanehamadi@gmail.com>
106 lines
2.9 KiB
Python
106 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import List, Optional
|
|
|
|
import openai
|
|
from openai import Model
|
|
|
|
from autogpt.llm.modelsinfo import COSTS
|
|
from autogpt.logs import logger
|
|
from autogpt.singleton import Singleton
|
|
|
|
|
|
class ApiManager(metaclass=Singleton):
|
|
def __init__(self):
|
|
self.total_prompt_tokens = 0
|
|
self.total_completion_tokens = 0
|
|
self.total_cost = 0
|
|
self.total_budget = 0
|
|
self.models: Optional[list[Model]] = None
|
|
|
|
def reset(self):
|
|
self.total_prompt_tokens = 0
|
|
self.total_completion_tokens = 0
|
|
self.total_cost = 0
|
|
self.total_budget = 0.0
|
|
self.models = None
|
|
|
|
def update_cost(self, prompt_tokens, completion_tokens, model):
|
|
"""
|
|
Update the total cost, prompt tokens, and completion tokens.
|
|
|
|
Args:
|
|
prompt_tokens (int): The number of tokens used in the prompt.
|
|
completion_tokens (int): The number of tokens used in the completion.
|
|
model (str): The model used for the API call.
|
|
"""
|
|
# the .model property in API responses can contain version suffixes like -v2
|
|
model = model[:-3] if model.endswith("-v2") else model
|
|
|
|
self.total_prompt_tokens += prompt_tokens
|
|
self.total_completion_tokens += completion_tokens
|
|
self.total_cost += (
|
|
prompt_tokens * COSTS[model]["prompt"]
|
|
+ completion_tokens * COSTS[model]["completion"]
|
|
) / 1000
|
|
logger.debug(f"Total running cost: ${self.total_cost:.3f}")
|
|
|
|
def set_total_budget(self, total_budget):
|
|
"""
|
|
Sets the total user-defined budget for API calls.
|
|
|
|
Args:
|
|
total_budget (float): The total budget for API calls.
|
|
"""
|
|
self.total_budget = total_budget
|
|
|
|
def get_total_prompt_tokens(self):
|
|
"""
|
|
Get the total number of prompt tokens.
|
|
|
|
Returns:
|
|
int: The total number of prompt tokens.
|
|
"""
|
|
return self.total_prompt_tokens
|
|
|
|
def get_total_completion_tokens(self):
|
|
"""
|
|
Get the total number of completion tokens.
|
|
|
|
Returns:
|
|
int: The total number of completion tokens.
|
|
"""
|
|
return self.total_completion_tokens
|
|
|
|
def get_total_cost(self):
|
|
"""
|
|
Get the total cost of API calls.
|
|
|
|
Returns:
|
|
float: The total cost of API calls.
|
|
"""
|
|
return self.total_cost
|
|
|
|
def get_total_budget(self):
|
|
"""
|
|
Get the total user-defined budget for API calls.
|
|
|
|
Returns:
|
|
float: The total budget for API calls.
|
|
"""
|
|
return self.total_budget
|
|
|
|
def get_models(self) -> List[Model]:
|
|
"""
|
|
Get list of available GPT models.
|
|
|
|
Returns:
|
|
list: List of available GPT models.
|
|
|
|
"""
|
|
if self.models is None:
|
|
all_models = openai.Model.list()["data"]
|
|
self.models = [model for model in all_models if "gpt" in model["id"]]
|
|
|
|
return self.models
|