week2/community-contributions/AISystem.py

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