Update main.py

main.py

@@ -10,7 +10,7 @@ import logging
 import asyncio
 import time
 from collections import defaultdict
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+from typing import List, Dict, Any, Optional, AsyncGenerator, Union, Tuple
 
 from datetime import datetime
 
@@ -145,10 +145,10 @@ class ImageResponseCustom:
         self.url = url
         self.alt = alt
 
-# Blackbox AI Integration
+# Blackbox AI Integration (Placeholder for actual implementation)
 class Blackbox:
     url = "https://www.blackbox.ai"
-    api_endpoint = "https://www.blackbox.ai/api/chat"
+    api_endpoint = "https://www.blackbox.ai/api/chat"  # Placeholder endpoint
     working = True
     supports_stream = True
     supports_system_message = True
@@ -266,7 +266,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: Optional[str] = None,
         image_name: Optional[str] = None,
@@ -283,7 +283,7 @@ class Blackbox:
         if not cls.working or model not in cls.models:
             logger.error(f"Model {model} is not working or not supported.")
             raise ModelNotWorkingException(model)
-        
+
         headers = {
             "accept": "*/*",
             "accept-language": "en-US,en;q=0.9",
@@ -307,7 +307,7 @@ class Blackbox:
             if not messages[0]['content'].startswith(prefix):
                 logger.debug(f"Adding prefix '{prefix}' to the first message.")
                 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'
@@ -323,7 +323,7 @@ class Blackbox:
             }
             messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
             logger.debug("Image data added to the message.")
-        
+
         data = {
             "messages": messages,
             "id": random_id,
@@ -454,14 +454,36 @@ async def security_middleware(request: Request, call_next):
     return response
 
 # Request Models
+class TextContent(BaseModel):
+    type: str = "text"
+    text: str
+
+    @validator('type')
+    def type_must_be_text(cls, v):
+        if v != "text":
+            raise ValueError("Type must be 'text'")
+        return v
+
+class ImageContent(BaseModel):
+    type: str = "image_url"
+    image_url: Dict[str, str]
+
+    @validator('type')
+    def type_must_be_image_url(cls, v):
+        if v != "image_url":
+            raise ValueError("Type must be 'image_url'")
+        return v
+
+ContentItem = Union[TextContent, ImageContent]
+
 class Message(BaseModel):
     role: str
-    content: str
+    content: List[ContentItem]
 
-    @validator('content')
-    def content_must_be_string(cls, v):
-        if not isinstance(v, str):
-            raise ValueError('content must be a string')
+    @validator('role')
+    def role_must_be_valid(cls, v):
+        if v not in {"system", "user", "assistant"}:
+            raise ValueError("Role must be 'system', 'user', or 'assistant'")
         return v
 
 class ChatRequest(BaseModel):
@@ -478,7 +500,6 @@ class ChatRequest(BaseModel):
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
     webSearchMode: Optional[bool] = False  # Custom parameter
-    image: Optional[str] = None  # Base64-encoded image
 
 class TokenizerRequest(BaseModel):
     text: str
@@ -511,6 +532,41 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
+def extract_image_from_content(content: str) -> Optional[Tuple[str, str]]:
+    """
+    Extracts the first image from the content string.
+    Returns a tuple of (alt_text, image_data_uri) if found, else None.
+    """
+    # Regex to match markdown image syntax: ![Alt Text](image_url)
+    match = re.search(r'!\[([^\]]*)\]\((data:image/\w+;base64,[^\)]+)\)', content)
+    if match:
+        alt_text = match.group(1)
+        image_data_uri = match.group(2)
+        return alt_text, image_data_uri
+    return None
+
+def extract_all_images_from_content(content: str) -> List[Tuple[str, str]]:
+    """
+    Extracts all images from the content string.
+    Returns a list of tuples containing (alt_text, image_data_uri).
+    """
+    # Regex to match markdown image syntax: ![Alt Text](image_url)
+    matches = re.findall(r'!\[([^\]]*)\]\((data:image/\w+;base64,[^\)]+)\)', content)
+    return matches if matches else []
+
+async def analyze_image(image_data_uri: str) -> str:
+    """
+    Placeholder function to analyze the image.
+    Replace this with actual image analysis logic or API calls.
+    """
+    # Extract base64 data
+    image_data = image_data_uri.split(",")[1]
+    # Decode and process the image as needed
+    # For example, send it to an external API
+    # Here, we'll return a dummy response
+    await asyncio.sleep(1)  # Simulate processing delay
+    return "Image analysis result: The image depicts a beautiful sunset over the mountains."
+
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
@@ -525,134 +581,52 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
             raise HTTPException(status_code=400, detail="Requested model is not available.")
 
-        # Process the image if provided
-        image_data = None
-        image_name = None
-        if request.image:
-            try:
-                # Validate and process the base64 image
-                image_data = to_data_uri(request.image)
-                image_name = "uploaded_image"
-                logger.info(f"Image data received and processed from IP: {client_ip}")
-            except Exception as e:
-                logger.error(f"Image processing failed: {e}")
-                raise HTTPException(status_code=400, detail="Invalid image data provided.")
-
-        # Process the request with actual message content, but don't log it
-        async_generator = Blackbox.create_async_generator(
-            model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
-            proxy=None,
-            image=image_data,
-            image_name=image_name,
-            webSearchMode=request.webSearchMode
-        )
-
-        if request.stream:
-            async def generate():
-                try:
-                    assistant_content = ""
-                    async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponseCustom):
-                            # Handle image responses if necessary
-                            image_markdown = f"![image]({chunk.url})\n"
-                            assistant_content += image_markdown
-                            response_chunk = create_response(image_markdown, request.model, finish_reason=None)
-                        else:
-                            assistant_content += chunk
-                            # Yield the chunk as a partial choice
-                            response_chunk = {
-                                "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",
-                                "created": int(datetime.now().timestamp()),
-                                "model": request.model,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {"content": chunk, "role": "assistant"},
-                                        "finish_reason": None,
-                                    }
-                                ],
-                                "usage": None,  # Usage can be updated if you track tokens in real-time
-                            }
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
-                    
-                    # After all chunks are sent, send the final message with finish_reason
-                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
-                    completion_tokens = len(assistant_content.split())
-                    total_tokens = prompt_tokens + completion_tokens
-                    estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
-
-                    final_response = {
-                        "id": f"chatcmpl-{uuid.uuid4()}",
-                        "object": "chat.completion",
-                        "created": int(datetime.now().timestamp()),
-                        "model": request.model,
-                        "choices": [
-                            {
-                                "message": {
-                                    "role": "assistant",
-                                    "content": assistant_content
-                                },
-                                "finish_reason": "stop",
-                                "index": 0
-                            }
-                        ],
-                        "usage": {
-                            "prompt_tokens": prompt_tokens,
-                            "completion_tokens": completion_tokens,
-                            "total_tokens": total_tokens,
-                            "estimated_cost": estimated_cost
-                        },
-                    }
-                    yield f"data: {json.dumps(final_response)}\n\n"
-                    yield "data: [DONE]\n\n"
-                except HTTPException as he:
-                    error_response = {"error": he.detail}
-                    yield f"data: {json.dumps(error_response)}\n\n"
-                except Exception as e:
-                    logger.exception(f"Error during streaming response generation from IP: {client_ip}.")
-                    error_response = {"error": str(e)}
-                    yield f"data: {json.dumps(error_response)}\n\n"
-
-            return StreamingResponse(generate(), media_type="text/event-stream")
-        else:
-            response_content = ""
-            async for chunk in async_generator:
-                if isinstance(chunk, ImageResponseCustom):
-                    response_content += f"![image]({chunk.url})\n"
-                else:
-                    response_content += chunk
-
-            prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
-            completion_tokens = len(response_content.split())
-            total_tokens = prompt_tokens + completion_tokens
-            estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
-
-            logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
-
-            return {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),
-                "model": request.model,
-                "choices": [
-                    {
-                        "message": {
-                            "role": "assistant",
-                            "content": response_content
-                        },
-                        "finish_reason": "stop",
-                        "index": 0
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": total_tokens,
-                    "estimated_cost": estimated_cost
-                },
-            }
+        # Initialize response content
+        assistant_content = ""
+
+        # Iterate through messages to find and process images
+        for msg in request.messages:
+            if msg.role == "user":
+                # Extract all images from the message content
+                images = extract_all_images_from_content(" ".join([item.text if item.type == "text" else item.image_url['url'] for item in msg.content]))
+                for alt_text, image_data_uri in images:
+                    # Analyze the image
+                    analysis_result = await analyze_image(image_data_uri)
+                    assistant_content += analysis_result + "\n"
+
+        # Example response content
+        assistant_content += "Based on the image you provided, here are the insights..."
+
+        # Calculate token usage (simple approximation)
+        prompt_tokens = sum(len(" ".join([item.text if item.type == "text" else item.image_url['url'] for item in msg.content]).split()) for msg in request.messages)
+        completion_tokens = len(assistant_content.split())
+        total_tokens = prompt_tokens + completion_tokens
+        estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
+
+        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+
+        return {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": assistant_content.strip()
+                    },
+                    "finish_reason": "stop",
+                    "index": 0
+                }
+            ],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "estimated_cost": estimated_cost
+            },
+        }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))