Update main.py

main.py

@@ -7,8 +7,6 @@ import json
 import logging
 import asyncio
 import time
-import base64
-from io import BytesIO
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 
@@ -18,7 +16,10 @@ from aiohttp import ClientSession, ClientTimeout, ClientError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
 from pydantic import BaseModel
+
 from PIL import Image
+import base64
+from io import BytesIO
 
 # Configure logging
 logging.basicConfig(
@@ -108,25 +109,6 @@ class ImageResponse:
 def to_data_uri(image: Any) -> str:
     return "data:image/png;base64,..."  # Replace with actual base64 data
 
-# Utility functions for image processing
-def decode_base64_image(base64_str: str) -> Image.Image:
-    try:
-        image_data = base64.b64decode(base64_str)
-        image = Image.open(BytesIO(image_data))
-        return image
-    except Exception as e:
-        logger.error("Failed to decode base64 image.")
-        raise HTTPException(status_code=400, detail="Invalid base64 image data.") from e
-
-def analyze_image(image: Image.Image) -> str:
-    """
-    Placeholder for image analysis.
-    Replace this with actual image analysis logic.
-    """
-    # Example: Return image size as analysis
-    width, height = image.size
-    return f"Image analyzed successfully. Width: {width}px, Height: {height}px."
-
 class Blackbox:
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
@@ -440,7 +422,7 @@ async def security_middleware(request: Request, call_next):
 # Request Models
 class Message(BaseModel):
     role: str
-    content: Union[str, List[Any]]  # Adjusted to accept list if needed
+    content: Union[str, List[Any]]  # content can be a string or a list (for images)
 
 class ChatRequest(BaseModel):
     model: str
@@ -510,31 +492,59 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             logger.exception("Unexpected error during image analysis.")
             raise HTTPException(status_code=500, detail="Image analysis failed.") from e
 
-    try:
-        # Validate that the requested model is available
-        if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
-            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 request with actual message content and image data
-        async_generator = Blackbox.create_async_generator(
-            model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            image=request.image,
-            image_name="uploaded_image",  # You can modify this as needed
-            webSearchMode=request.webSearchMode
-        )
+    # Prepare messages to send to the external API, excluding image data
+    processed_messages = []
+    for msg in request.messages:
+        if isinstance(msg.content, list) and len(msg.content) == 2:
+            # Assume the second item is image data, skip it
+            processed_messages.append({
+                "role": msg.role,
+                "content": msg.content[0]["text"]  # Only include the text part
+            })
+        else:
+            processed_messages.append({
+                "role": msg.role,
+                "content": msg.content
+            })
+
+    # Create a modified ChatRequest without the image
+    modified_request = ChatRequest(
+        model=request.model,
+        messages=[msg for msg in processed_messages],
+        stream=request.stream,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        max_tokens=request.max_tokens,
+        presence_penalty=request.presence_penalty,
+        frequency_penalty=request.frequency_penalty,
+        logit_bias=request.logit_bias,
+        user=request.user,
+        webSearchMode=request.webSearchMode,
+        image=None  # Exclude image from external API
+    )
 
+    try:
         if request.stream:
-            async def generate():
+            logger.info("Streaming response")
+            streaming_response = await Blackbox.create_async_generator(
+                model=modified_request.model,
+                messages=[{"role": msg["role"], "content": msg["content"]} for msg in modified_request.messages],
+                proxy=None,
+                image=None,
+                image_name=None,
+                webSearchMode=modified_request.webSearchMode
+            )
+
+            # Wrap the streaming generator to include image analysis at the end
+            async def generate_with_analysis():
+                assistant_content = ""
                 try:
-                    assistant_content = ""
-                    async for chunk in async_generator:
+                    async for chunk in streaming_response:
                         if isinstance(chunk, ImageResponse):
                             # 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)
+                            response_chunk = create_response(image_markdown, modified_request.model, finish_reason=None)
                         else:
                             assistant_content += chunk
                             # Yield the chunk as a partial choice
@@ -542,7 +552,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                                 "id": f"chatcmpl-{uuid.uuid4()}",
                                 "object": "chat.completion.chunk",
                                 "created": int(datetime.now().timestamp()),
-                                "model": request.model,
+                                "model": modified_request.model,
                                 "choices": [
                                     {
                                         "index": 0,
@@ -555,7 +565,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         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)
+                    prompt_tokens = sum(len(msg["content"].split()) for msg in modified_request.messages)
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
@@ -564,7 +574,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         "id": f"chatcmpl-{uuid.uuid4()}",
                         "object": "chat.completion",
                         "created": int(datetime.now().timestamp()),
-                        "model": request.model,
+                        "model": modified_request.model,
                         "choices": [
                             {
                                 "message": {
@@ -596,16 +606,26 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     error_response = {"error": str(e)}
                     yield f"data: {json.dumps(error_response)}\n\n"
 
-            return StreamingResponse(generate(), media_type="text/event-stream")
+            return StreamingResponse(generate_with_analysis(), media_type="text/event-stream")
         else:
+            logger.info("Non-streaming response")
+            streaming_response = await Blackbox.create_async_generator(
+                model=modified_request.model,
+                messages=[{"role": msg["role"], "content": msg["content"]} for msg in modified_request.messages],
+                proxy=None,
+                image=None,
+                image_name=None,
+                webSearchMode=modified_request.webSearchMode
+            )
+
             response_content = ""
-            async for chunk in async_generator:
+            async for chunk in streaming_response:
                 if isinstance(chunk, ImageResponse):
                     response_content += f"![image]({chunk.url})\n"
                 else:
                     response_content += chunk
 
-            prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
+            prompt_tokens = sum(len(msg["content"].split()) for msg in modified_request.messages)
             completion_tokens = len(response_content.split())
             total_tokens = prompt_tokens + completion_tokens
             estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
@@ -616,7 +636,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
                 "created": int(datetime.now().timestamp()),
-                "model": request.model,
+                "model": modified_request.model,
                 "choices": [
                     {
                         "message": {
@@ -710,6 +730,25 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         },
     )
 
+# Image Processing Utilities
+def decode_base64_image(base64_str: str) -> Image.Image:
+    try:
+        image_data = base64.b64decode(base64_str)
+        image = Image.open(BytesIO(image_data))
+        return image
+    except Exception as e:
+        logger.error("Failed to decode base64 image.")
+        raise HTTPException(status_code=400, detail="Invalid base64 image data.") from e
+
+def analyze_image(image: Image.Image) -> str:
+    """
+    Placeholder for image analysis.
+    Replace this with actual image analysis logic.
+    """
+    # Example: Return image size as analysis
+    width, height = image.size
+    return f"Image analyzed successfully. Width: {width}px, Height: {height}px."
+
 # Run the application
 if __name__ == "__main__":
     import uvicorn