Co-authored-by: k-boikov <64261260+k-boikov@users.noreply.github.com>
This commit is contained in:
Konrad
2023-05-20 19:45:27 -04:00
committed by GitHub
parent c30f5b7d5e
commit 57ea7b5216
7 changed files with 122 additions and 3 deletions

View File

@@ -1,6 +1,9 @@
from __future__ import annotations
from typing import List, Optional
import openai
from openai import Model
from autogpt.config import Config
from autogpt.llm.modelsinfo import COSTS
@@ -14,12 +17,14 @@ class ApiManager(metaclass=Singleton):
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 create_chat_completion(
self,
@@ -127,3 +132,17 @@ class ApiManager(metaclass=Singleton):
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