refactor: restructure the chat and conversations.

Fix regression issues along the way

Author: ksarang90

chatbot-api/__init__.py

@@ -1 +1 @@
-"""Package initialization file""" 
+"""ChatBot API package""" 

chatbot-api/src/app.py

@@ -1,11 +1,21 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Union, Dict
 from datetime import datetime
 
+@dataclass
+class ImageContent:
+    type: str = "image_url"
+    image_url: Dict[str, str] = None
+
+@dataclass
+class TextContent:
+    type: str = "text"
+    text: str = ""
+
 @dataclass
 class Message:
     role: str
-    content: str
+    content: Union[str, List[Union[TextContent, ImageContent]]]
     model: Optional[str] = None
     timestamp: datetime = datetime.utcnow()
 

chatbot-api/src/chat/models.py

@@ -1,33 +1,72 @@
 from typing import Protocol, AsyncIterator, List
 from openai import OpenAI
-from .models import Message, ChatResponse
+from .models import Message, ChatResponse, TextContent, ImageContent
+import logging
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
 
 class AIProvider(Protocol):
     """Protocol for AI providers"""
-    async def generate_response(self, messages: List[Message]) -> ChatResponse: ...
-    async def generate_stream(self, messages: List[Message]) -> AsyncIterator[str]: ...
+    def generate_response(self, messages: List[Message]) -> ChatResponse: ...
+    def generate_stream(self, messages: List[Message]) -> AsyncIterator[str]: ...
 
 class OpenAIProvider:
     """OpenAI implementation of AIProvider"""
     def __init__(self, api_key: str, api_base: str | None = None):
         self.client = OpenAI(api_key=api_key, base_url=api_base)
 
-    async def generate_response(self, messages: List[Message]) -> ChatResponse:
-        response = await self.client.chat.completions.create(
-            model="gpt-4o",
-            messages=[{"role": m.role, "content": m.content} for m in messages]
+    def _format_message(self, message: Message) -> dict:
+        """Format message for OpenAI API"""
+        if isinstance(message.content, str):
+            return {"role": message.role, "content": message.content}
+        
+        # For messages with text and images
+        formatted_content = []
+        for item in message.content:
+            if isinstance(item, TextContent):
+                formatted_content.append({"type": "text", "text": item.text})
+            elif isinstance(item, ImageContent):
+                formatted_content.append({
+                    "type": "image_url",
+                    "image_url": item.image_url
+                })
+        return {"role": message.role, "content": formatted_content}
+
+    def generate_response(self, messages: List[Message]) -> ChatResponse:
+        """Generate a response for messages"""
+        # Use the model from the last message, or default to claude-3-5-sonnet
+        model = next((m.model for m in reversed(messages) if m.model), "claude-3-5-sonnet")
+        
+        formatted_messages = [self._format_message(m) for m in messages]
+        response = self.client.chat.completions.create(
+            model=model,
+            messages=formatted_messages
         )
+        
         return ChatResponse(
             content=response.choices[0].message.content,
-            model=response.model
+            model=response.model,
+            timestamp=datetime.utcnow()
         )
 
     async def generate_stream(self, messages: List[Message]) -> AsyncIterator[str]:
-        stream = await self.client.chat.completions.create(
-            model="gpt-4o",
-            messages=[{"role": m.role, "content": m.content} for m in messages],
-            stream=True
-        )
-        async for chunk in stream:
-            if chunk.choices[0].delta.content:
-                yield chunk.choices[0].delta.content 
+        """Stream response for messages"""
+        try:
+            # Use the model from the last message, or default to claude-3-5-sonnet
+            model = next((m.model for m in reversed(messages) if m.model), "claude-3-5-sonnet")
+            
+            formatted_messages = [self._format_message(m) for m in messages]
+            stream = self.client.chat.completions.create(
+                model=model,
+                messages=formatted_messages,
+                stream=True
+            )
+            
+            for chunk in stream:
+                if chunk.choices[0].delta.content:
+                    content = chunk.choices[0].delta.content
+                    yield content
+        except Exception as e:
+            logger.error(f"Error in stream: {str(e)}")
+            raise