Update main.py

main.py

@@ -10,8 +10,8 @@ from aiohttp import ClientSession, ClientTimeout, ClientError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header, status
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
-from typing import List, Dict, Any, Optional, AsyncGenerator
+from pydantic import BaseModel, Field, validator
+from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 from datetime import datetime
 from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
@@ -100,7 +100,7 @@ def to_data_uri(image: Any) -> str:
     return "data:image/png;base64,..."  # Replace with actual base64 data if needed
 
 # Token Counting using tiktoken
-def count_tokens(messages: List[Dict[str, str]], model: str) -> int:
+def count_tokens(messages: List[Dict[str, Any]], model: str) -> int:
     """
     Counts the number of tokens in the messages using tiktoken.
     Adjust the encoding based on the model.
@@ -111,7 +111,14 @@ def count_tokens(messages: List[Dict[str, str]], model: str) -> int:
         encoding = tiktoken.get_encoding("cl100k_base")  # Default encoding
     tokens = 0
     for message in messages:
-        tokens += len(encoding.encode(message['content']))
+        if isinstance(message['content'], list):
+            for content_part in message['content']:
+                if content_part.get('type') == 'text':
+                    tokens += len(encoding.encode(content_part['text']))
+                elif content_part.get('type') == 'image_url':
+                    tokens += len(encoding.encode(content_part['image_url']['url']))
+        else:
+            tokens += len(encoding.encode(message['content']))
     return tokens
 
 # Blackbox Class: Handles interaction with the external AI service
@@ -235,7 +242,7 @@ class Blackbox:
     async def create_async_generator(
         cls,
         model: str,
-        messages: List[Dict[str, str]],
+        messages: List[Dict[str, Any]],
         proxy: Optional[str] = None,
         image: Any = None,
         image_name: Optional[str] = None,
@@ -269,22 +276,33 @@ class Blackbox:
 
         if model in cls.model_prefixes:
             prefix = cls.model_prefixes[model]
-            if not messages[0]['content'].startswith(prefix):
-                logger.debug(f"Adding prefix '{prefix}' to the first message.")
+            if messages and isinstance(messages[0]['content'], list):
+                # Prepend prefix to the first text message
+                for content_part in messages[0]['content']:
+                    if content_part.get('type') == 'text' and not content_part['text'].startswith(prefix):
+                        logger.debug(f"Adding prefix '{prefix}' to the first text message.")
+                        content_part['text'] = f"{prefix} {content_part['text']}"
+                        break
+            elif messages and isinstance(messages[0]['content'], str) and not messages[0]['content'].startswith(prefix):
                 messages[0]['content'] = f"{prefix} {messages[0]['content']}"
-        
+
         random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
-        messages[-1]['id'] = random_id
-        messages[-1]['role'] = 'user'
-        if image is not None:
-            messages[-1]['data'] = {
-                'fileText': '',
-                'imageBase64': to_data_uri(image),
-                'title': image_name
-            }
-            messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
-            logger.debug("Image data added to the message.")
+        # Assuming the last message is from the user
+        if messages:
+            last_message = messages[-1]
+            if isinstance(last_message['content'], list):
+                for content_part in last_message['content']:
+                    if content_part.get('type') == 'text':
+                        content_part['role'] = 'user'
+            else:
+                last_message['id'] = random_id
+                last_message['role'] = 'user'
         
+        if image is not None:
+            # Process image if required
+            # This implementation assumes that image URLs are handled by the external service
+            pass  # Implement as needed
+
         data = {
             "messages": messages,
             "id": random_id,
@@ -337,36 +355,13 @@ class Blackbox:
                                 logger.error("Image URL not found in the response.")
                                 raise Exception("Image URL not found in the response")
                         else:
-                            full_response = ""
-                            search_results_json = ""
-                            try:
-                                async for chunk, _ in response.content.iter_chunks():
-                                    if chunk:
-                                        decoded_chunk = chunk.decode(errors='ignore')
-                                        decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
-                                        if decoded_chunk.strip():
-                                            if '$~~~$' in decoded_chunk:
-                                                search_results_json += decoded_chunk
-                                            else:
-                                                full_response += decoded_chunk
-                                                yield decoded_chunk
-                                logger.info("Finished streaming response chunks.")
-                            except Exception as e:
-                                logger.exception("Error while iterating over response chunks.")
-                                raise e
-                            if data["webSearchMode"] and search_results_json:
-                                match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
-                                if match:
-                                    try:
-                                        search_results = json.loads(match.group(1))
-                                        formatted_results = "\n\n**Sources:**\n"
-                                        for i, result in enumerate(search_results[:5], 1):
-                                            formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
-                                        logger.info("Formatted search results.")
-                                        yield formatted_results
-                                    except json.JSONDecodeError as je:
-                                        logger.error("Failed to parse search results JSON.")
-                                        raise je
+                            async for chunk in response.content.iter_chunks():
+                                if chunk:
+                                    decoded_chunk = chunk.decode(errors='ignore')
+                                    decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
+                                    if decoded_chunk.strip():
+                                        yield decoded_chunk
+                break  # Exit the retry loop if successful
             except ClientError as ce:
                 logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
                 if attempt == retry_attempts - 1:
@@ -381,9 +376,28 @@ class Blackbox:
                     raise HTTPException(status_code=500, detail=str(e))
 
 # Pydantic Models
+class TextContent(BaseModel):
+    type: str = Field(..., description="Type of content, e.g., 'text'.")
+    text: str = Field(..., description="The text content.")
+
+class ImageURLContent(BaseModel):
+    type: str = Field(..., description="Type of content, e.g., 'image_url'.")
+    image_url: Dict[str, str] = Field(..., description="Dictionary containing the image URL.")
+
+Content = Union[TextContent, ImageURLContent]
+
 class Message(BaseModel):
     role: str = Field(..., description="The role of the message author.")
-    content: str = Field(..., description="The content of the message.")
+    content: Union[str, List[Content]] = Field(..., description="The content of the message. Can be a string or a list of content parts.")
+
+    @validator('content', pre=True)
+    def validate_content(cls, v):
+        if isinstance(v, list):
+            return [Content(**item) for item in v]
+        elif isinstance(v, str):
+            return v
+        else:
+            raise ValueError("Content must be either a string or a list of content parts.")
 
 class ChatRequest(BaseModel):
     model: str = Field(..., description="ID of the model to use.")
@@ -431,12 +445,26 @@ async def chat_completions(
 ):
     logger.info(f"Received chat completions request: {chat_request}")
     try:
-        messages = [{"role": msg.role, "content": msg.content} for msg in chat_request.messages]
-        prompt_tokens = count_tokens(messages, chat_request.model)
+        # Process messages for token counting and sending to Blackbox
+        processed_messages = []
+        for msg in chat_request.messages:
+            if isinstance(msg.content, list):
+                # Convert list of content parts to a structured format
+                combined_content = []
+                for part in msg.content:
+                    if isinstance(part, TextContent):
+                        combined_content.append({"type": part.type, "text": part.text})
+                    elif isinstance(part, ImageURLContent):
+                        combined_content.append({"type": part.type, "image_url": part.image_url})
+                processed_messages.append({"role": msg.role, "content": combined_content})
+            else:
+                processed_messages.append({"role": msg.role, "content": msg.content})
+
+        prompt_tokens = count_tokens(processed_messages, chat_request.model)
 
         async_generator = Blackbox.create_async_generator(
             model=chat_request.model,
-            messages=messages,
+            messages=processed_messages,
             image=None,  # Adjust if image handling is required
             image_name=None,
             webSearchMode=chat_request.webSearchMode