Spaces:
Running
Running
File size: 4,689 Bytes
3703b6a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import os
from openai import OpenAI
from dotenv import load_dotenv
from typing import List, Optional
from utils import log
class GPTClient:
def __init__(self, model: str, api_key: Optional[str] = None):
# Initialize
self.model = model
self._init_environment(api_key)
self.histories = list()
self.token_usages = dict()
self.__first_turn = True
def _init_environment(self, api_key: Optional[str] = None) -> None:
"""
Initialize OpenAI client.
Args:
api_key (Optional[str]): API key for OpenAI. If not provided, it will
be loaded from environment variables.
"""
if not api_key:
load_dotenv(override=True)
api_key = os.environ.get("OPENAI_API_KEY", None)
self.client = OpenAI(api_key=api_key)
def reset_history(self, verbose: bool = True) -> None:
"""
Reset the conversation history.
Args:
verbose (bool): Whether to print verbose output. Defaults to True.
"""
self.__first_turn = True
self.histories = list()
self.token_usages = dict()
if verbose:
log('Conversation history has been reset.', color=True)
def __make_payload(self, user_prompt: str) -> List[dict]:
"""
Create a payload for API calls to the GPT model.
Args:
user_prompt (str): User prompt.
Returns:
List[dict]: Payload including prompts and image data.
"""
payloads = list()
user_contents = {"role": "user", "content": []}
# User prompts
user_contents["content"].append(
{"type": "text", "text": user_prompt}
)
payloads.append(user_contents)
return payloads
def __call__(self,
user_prompt: str,
system_prompt: Optional[str] = None,
using_multi_turn: bool = True,
greeting: Optional[str] = None,
verbose: bool = True,
**kwargs) -> str:
"""
Sends a chat completion request to the model with optional image input and system prompt.
Args:
user_prompt (str): The main user prompt or query to send to the model.
system_prompt (Optional[str], optional): An optional system-level prompt to set context or behavior. Defaults to None.
using_multi_turn (bool): Whether to structure it as multi-turn. Defaults to True.
greeting (Optional[str]): An optional greeting message to include in the conversation. Defaults to None.
verbose (bool): Whether to print verbose output. Defaults to True.
Raises:
e: Any exception raised during the API call is re-raised.
Returns:
str: The model's response message.
"""
try:
# To ensure empty history
if not using_multi_turn:
self.reset_history(verbose)
if self.__first_turn:
# System prompt
if system_prompt:
self.histories.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
# Greeting
if greeting and self.__first_turn:
self.histories.append({"role": "assistant", "content": [{"type": "text", "text": greeting}]})
self.__first_turn = False
# User prompt
self.histories += self.__make_payload(user_prompt)
# Model response
response = self.client.chat.completions.create(
model=self.model,
messages=self.histories,
**kwargs
)
assistant_msg = response.choices[0].message
self.histories.append({"role": assistant_msg.role, "content": [{"type": "text", "text": assistant_msg.content}]})
# Logging token usage
if response.usage:
self.token_usages.setdefault("prompt_tokens", []).append(response.usage.prompt_tokens)
self.token_usages.setdefault("completion_tokens", []).append(response.usage.completion_tokens)
self.token_usages.setdefault("total_tokens", []).append(response.usage.total_tokens)
self.token_usages.setdefault("reasoning_tokens", []).append(response.usage.completion_tokens_details.reasoning_tokens)
return assistant_msg.content
except Exception as e:
raise e
|