82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
|
|
from enum import Enum, auto
|
|
from openai import OpenAI
|
|
import anthropic
|
|
|
|
def formatPrompt(role, content):
|
|
return {"role": role, "content": content}
|
|
|
|
class AI(Enum):
|
|
OPEN_AI = "OPEN_AI"
|
|
CLAUDE = "CLAUDE"
|
|
GEMINI = "GEMINI"
|
|
OLLAMA = "OLLAMA"
|
|
|
|
class AISystem:
|
|
def __init__(self, processor, system_string="", model="", type=AI.OPEN_AI):
|
|
"""
|
|
Initialize the ChatSystem with a system string and empty messages list.
|
|
|
|
:param system_string: Optional initial system string description
|
|
"""
|
|
self.processor = processor
|
|
self.system = system_string
|
|
self.model = model
|
|
self.messages = []
|
|
self.type = type
|
|
|
|
def call(self, message):
|
|
self.messages.append(message)
|
|
toSend = self.messages
|
|
|
|
if self.type == AI.CLAUDE:
|
|
message = self.processor.messages.create(
|
|
model=self.model,
|
|
system=self.system,
|
|
messages=self.messages,
|
|
max_tokens=500
|
|
)
|
|
return message.content[0].text
|
|
else:
|
|
toSend.insert(0,self.system)
|
|
completion = self.processor.chat.completions.create(
|
|
model=self.model,
|
|
messages= toSend
|
|
)
|
|
return completion.choices[0].message.content
|
|
|
|
def stream(self, message, usingGradio=False):
|
|
self.messages.append(message)
|
|
|
|
if self.type == AI.CLAUDE:
|
|
result = self.processor.messages.stream(
|
|
model=self.model,
|
|
system=self.system,
|
|
messages=self.messages,
|
|
temperature=0.7,
|
|
max_tokens=500
|
|
)
|
|
response_chunks = ""
|
|
with result as stream:
|
|
for text in stream.text_stream:
|
|
if usingGradio:
|
|
response_chunks += text or ""
|
|
yield response_chunks
|
|
else:
|
|
yield text
|
|
else:
|
|
toSend = self.messages
|
|
toSend.insert(0,self.system)
|
|
stream = self.processor.chat.completions.create(
|
|
model=self.model,
|
|
messages= toSend,
|
|
stream=True
|
|
)
|
|
response_chunks = ""
|
|
for chunk in stream:
|
|
if usingGradio:
|
|
response_chunks += chunk.choices[0].delta.content or "" # need to yield the total cumulative results to gradio and not chunk by chunk
|
|
yield response_chunks
|
|
else:
|
|
yield chunk.choices[0].delta.content
|