Exercise week2 day 2: offer multi llms dropdown and reusable AISystem with stream option

week2/community-contributions/AISystem.py

@@ -0,0 +1,81 @@
+
+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