QueryHelper / gptManager.py
anumaurya114exp's picture
Overwrite queryHelper with queryHelper2
37bd4dd
raw
history blame
1.78 kB
import re
class ChatgptManager:
def __init__(self, openAIClient, model="gpt-3.5-turbo-1106", tokenLimit=8000):
self.client = openAIClient
self.tokenLimit = tokenLimit
self.model = model
def getResponseForUserInput(self, userInput, systemPrompt):
self.messages = []
newMessage = {"role":"system", "content":systemPrompt}
if not self.isTokeLimitExceeding(newMessage):
self.messages.append(newMessage)
else:
raise ValueError("System Prompt Too long.")
userMessage = {"role":"user", "content":userInput}
if not self.isTokeLimitExceeding(userMessage):
self.messages.append(userMessage)
else:
raise ValueError("Token Limit exceeding. With user input")
# completion = self.client.chat.completions.create(
# model="gpt-3.5-turbo-1106",
# messages=self.messages,
# temperature=0,
# )
completion = self.client.chat.completions.create(
model=self.model,
messages=self.messages,
temperature=0,
)
gptResponse = completion.choices[0].message.content
self.messages.append({"role": "assistant", "content": gptResponse})
return gptResponse
def isTokeLimitExceeding(self, newMessage=None, truncate=True, throwError=True):
if self.getTokenCount(newMessage=newMessage) > self.tokenLimit:
return True
return False
def getTokenCount(self, newMessage=None):
"""Token count including new Message"""
def getWordsCount(text):
return len(re.findall(r'\b\w+\b', text))
messages = self.messages[:]
if newMessage!=None:
messages.append(newMessage)
combinedContent = " ".join(msg["content"] for msg in messages)
currentTokensInMessages = getWordsCount(combinedContent)
return currentTokensInMessages