Spaces:
Sleeping
Sleeping
File size: 2,732 Bytes
5fdb69e |
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 |
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
|