Update main.py

main.py

@@ -1,5 +1,3 @@
-# main.py
-
 import os
 import re
 import random
@@ -9,22 +7,18 @@ 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, Union, Tuple, AsyncGenerator
+from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 
 from datetime import datetime
 
 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, validator
-from io import BytesIO
-import base64
-
-from dotenv import load_dotenv
-
-# Load environment variables from .env file
-load_dotenv()
+from pydantic import BaseModel
+from PIL import Image
 
 # Configure logging
 logging.basicConfig(
@@ -105,50 +99,37 @@ class ModelNotWorkingException(Exception):
         self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
         super().__init__(self.message)
 
-# Image Handling Functions
-ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'svg'}
+# Mock implementations for ImageResponse and to_data_uri
+class ImageResponse:
+    def __init__(self, url: str, alt: str):
+        self.url = url
+        self.alt = alt
 
-def is_allowed_extension(filename: str) -> bool:
-    """
-    Checks if the given filename has an allowed extension.
-    """
-    return '.' in filename and \
-           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+def to_data_uri(image: Any) -> str:
+    return "data:image/png;base64,..."  # Replace with actual base64 data
 
-def is_data_uri_an_image(data_uri: str) -> bool:
-    """
-    Checks if the given data URI represents an image.
-    """
-    match = re.match(r'data:image/(\w+);base64,', data_uri)
-    if not match:
-        raise ValueError("Invalid data URI image.")
-    image_format = match.group(1).lower()
-    if image_format not in ALLOWED_EXTENSIONS and image_format != "svg+xml":
-        raise ValueError("Invalid image format (from MIME type).")
-    return True
-
-def extract_data_uri(data_uri: str) -> bytes:
-    """
-    Extracts the binary data from the given data URI.
-    """
-    return base64.b64decode(data_uri.split(",")[1])
+# 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 to_data_uri(image: str) -> str:
+def analyze_image(image: Image.Image) -> str:
     """
-    Validates and returns the data URI for an image.
+    Placeholder for image analysis.
+    Replace this with actual image analysis logic.
     """
-    is_data_uri_an_image(image)
-    return image
-
-class ImageResponseCustom:
-    def __init__(self, url: str, alt: str):
-        self.url = url
-        self.alt = alt
+    # Example: Return image size as analysis
+    width, height = image.size
+    return f"Image analyzed successfully. Width: {width}px, Height: {height}px."
 
-# Placeholder for Blackbox AI Integration
 class Blackbox:
     url = "https://www.blackbox.ai"
-    api_endpoint = "https://www.blackbox.ai/api/chat"  # Placeholder endpoint
+    api_endpoint = "https://www.blackbox.ai/api/chat"
     working = True
     supports_stream = True
     supports_system_message = True
@@ -159,7 +140,6 @@ class Blackbox:
     models = [
         default_model,
         'blackboxai-pro',
-        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -180,13 +160,18 @@ class Blackbox:
         'ReactAgent',
         'XcodeAgent',
         'AngularJSAgent',
+        *image_models,
+        'Niansuh',
     ]
 
+    # Filter models based on AVAILABLE_MODELS
+    if AVAILABLE_MODELS:
+        models = [model for model in models if model in AVAILABLE_MODELS]
+
     agentMode = {
         'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
         'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
     }
-
     trendingAgentMode = {
         "blackboxai": {},
         "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
@@ -266,13 +251,13 @@ class Blackbox:
     async def create_async_generator(
         cls,
         model: str,
-        messages: List[Dict[str, Any]],
+        messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Optional[str] = None,
+        image: Any = None,
         image_name: Optional[str] = None,
         webSearchMode: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponseCustom], None]:
+    ) -> AsyncGenerator[Any, None]:
         model = cls.get_model(model)
         if model is None:
             logger.error(f"Model {model} is not available.")
@@ -283,7 +268,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 +292,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'
@@ -318,12 +303,12 @@ class Blackbox:
         if image is not None:
             messages[-1]['data'] = {
                 'fileText': '',
-                'imageBase64': image,
+                '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.")
-
+        
         data = {
             "messages": messages,
             "id": random_id,
@@ -365,14 +350,13 @@ class Blackbox:
                     async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                         response.raise_for_status()
                         logger.info(f"Received response with status {response.status}")
-                        if model in cls.image_models:
+                        if model == 'ImageGeneration':
                             response_text = await response.text()
-                            # Extract image URL from the response
                             url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
                             if url_match:
                                 image_url = url_match.group(0)
-                                logger.info(f"Image URL found: {image_url}")
-                                yield ImageResponseCustom(url=image_url, alt=messages[-1]['content'])
+                                logger.info(f"Image URL found.")
+                                yield ImageResponse(image_url, alt=messages[-1]['content'])
                             else:
                                 logger.error("Image URL not found in the response.")
                                 raise Exception("Image URL not found in the response")
@@ -421,7 +405,7 @@ class Blackbox:
                 if attempt == retry_attempts - 1:
                     raise HTTPException(status_code=500, detail=str(e))
 
-# Initialize FastAPI app
+# FastAPI app setup
 app = FastAPI()
 
 # Add the cleanup task when the app starts
@@ -453,39 +437,10 @@ async def security_middleware(request: Request, call_next):
     response = await call_next(request)
     return response
 
-# Pydantic 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]
-
+# Request Models
 class Message(BaseModel):
     role: str
-    content: Union[str, List[ContentItem]]
-
-    @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
+    content: Union[str, List[Any]]  # Adjusted to accept list if needed
 
 class ChatRequest(BaseModel):
     model: str
@@ -501,12 +456,11 @@ 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
 
-# Utility Functions
-
 def calculate_estimated_cost(prompt_tokens: int, completion_tokens: int) -> float:
     """
     Calculate the estimated cost based on the number of tokens.
@@ -516,18 +470,6 @@ def calculate_estimated_cost(prompt_tokens: int, completion_tokens: int) -> floa
     cost_per_token = 0.00000268
     return round((prompt_tokens + completion_tokens) * cost_per_token, 8)
 
-def count_tokens(text: str) -> int:
-    """
-    Counts the number of tokens in a given text using tiktoken.
-    """
-    try:
-        import tiktoken
-        encoding = tiktoken.get_encoding("cl100k_base")
-        return len(encoding.encode(text))
-    except ImportError:
-        # Fallback if tiktoken is not installed
-        return len(text.split())
-
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -547,60 +489,6 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
-def extract_all_images_from_content(content: Union[str, List[ContentItem]]) -> List[Tuple[str, str]]:
-    """
-    Extracts all images from the content.
-    Returns a list of tuples containing (alt_text, image_data_uri).
-    """
-    images = []
-    if isinstance(content, list):
-        for item in content:
-            if isinstance(item, ImageContent):
-                alt_text = item.image_url.get('alt', '')  # Optional alt text
-                image_data_uri = item.image_url.get('url', '')
-                if image_data_uri:
-                    images.append((alt_text, image_data_uri))
-    return images
-
-# Image Analysis Function (Placeholder)
-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.
-    """
-    try:
-        # Extract base64 data
-        image_data = image_data_uri.split(",")[1]
-        # Decode the image
-        image_bytes = base64.b64decode(image_data)
-        
-        # Here, integrate with an image analysis API or implement your own logic
-        # For demonstration, we'll simulate analysis with a dummy response.
-        await asyncio.sleep(1)  # Simulate processing delay
-        return "Image analysis result: The image depicts a beautiful sunset over the mountains."
-    except Exception as e:
-        logger.error(f"Failed to analyze image: {e}")
-        raise HTTPException(status_code=400, detail="Failed to process the provided image.")
-
-# Helper Function for Token Counting
-def count_prompt_tokens(request: ChatRequest) -> int:
-    """
-    Counts the number of tokens in the prompt (input messages).
-    Handles both string and list types for the 'content' field.
-    """
-    total = 0
-    for msg in request.messages:
-        if isinstance(msg.content, str):
-            total += count_tokens(msg.content)
-        elif isinstance(msg.content, list):
-            for item in msg.content:
-                if isinstance(item, TextContent):
-                    total += count_tokens(item.text)
-                elif isinstance(item, ImageContent):
-                    total += count_tokens(item.image_url['url'])
-    return total
-
-# Endpoint: POST /v1/chat/completions
 @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
@@ -609,61 +497,69 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
     logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
 
+    analysis_result = None
+    if request.image:
+        try:
+            image = decode_base64_image(request.image)
+            analysis_result = analyze_image(image)
+            logger.info("Image analysis completed successfully.")
+        except HTTPException as he:
+            logger.error(f"Image analysis failed: {he.detail}")
+            raise he
+        except Exception as e:
+            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.")
 
-        # 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(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 using the helper function
-        prompt_tokens = count_prompt_tokens(request)
-        completion_tokens = count_tokens(assistant_content)
-        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}")
+        # 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
+        )
 
         if request.stream:
             async def generate():
                 try:
-                    for msg in request.messages:
-                        if msg.role == "user":
-                            images = extract_all_images_from_content(msg.content)
-                            for alt_text, image_data_uri in images:
-                                analysis_result = await analyze_image(image_data_uri)
-                                response_chunk = {
-                                    "id": f"chatcmpl-{uuid.uuid4()}",
-                                    "object": "chat.completion.chunk",
-                                    "created": int(datetime.now().timestamp()),
-                                    "model": request.model,
-                                    "choices": [
-                                        {
-                                            "index": 0,
-                                            "delta": {"content": analysis_result + "\n", "role": "assistant"},
-                                            "finish_reason": None,
-                                        }
-                                    ],
-                                    "usage": None,
-                                }
-                                yield f"data: {json.dumps(response_chunk)}\n\n"
-
-                    # Final message
+                    assistant_content = ""
+                    async for chunk in async_generator:
+                        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)
+                        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",
@@ -673,7 +569,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                             {
                                 "message": {
                                     "role": "assistant",
-                                    "content": assistant_content.strip()
+                                    "content": assistant_content
                                 },
                                 "finish_reason": "stop",
                                 "index": 0
@@ -686,6 +582,9 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                             "estimated_cost": estimated_cost
                         },
                     }
+                    if analysis_result:
+                        final_response["choices"][0]["message"]["content"] += f"\n\n**Image Analysis:** {analysis_result}"
+                    
                     yield f"data: {json.dumps(final_response)}\n\n"
                     yield "data: [DONE]\n\n"
                 except HTTPException as he:
@@ -698,7 +597,21 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
             return StreamingResponse(generate(), media_type="text/event-stream")
         else:
-            return {
+            response_content = ""
+            async for chunk in async_generator:
+                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)
+            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}")
+
+            response = {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
                 "created": int(datetime.now().timestamp()),
@@ -707,7 +620,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     {
                         "message": {
                             "role": "assistant",
-                            "content": assistant_content.strip()
+                            "content": response_content
                         },
                         "finish_reason": "stop",
                         "index": 0
@@ -720,6 +633,11 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     "estimated_cost": estimated_cost
                 },
             }
+
+            if analysis_result:
+                response["choices"][0]["message"]["content"] += f"\n\n**Image Analysis:** {analysis_result}"
+
+            return response
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -732,23 +650,23 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
 # Endpoint: POST /v1/tokenizer
 @app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
-async def tokenizer(request: TokenizerRequest, req: Request, api_key: str = Depends(get_api_key)):
+async def tokenizer(request: TokenizerRequest, req: Request):
     client_ip = req.client.host
     text = request.text
-    token_count = count_tokens(text)
+    token_count = len(text.split())
     logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
     return {"text": text, "tokens": token_count}
 
 # Endpoint: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
-async def get_models(req: Request, api_key: str = Depends(get_api_key)):
+async def get_models(req: Request):
     client_ip = req.client.host
     logger.info(f"Fetching available models from IP: {client_ip}")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
 # Endpoint: GET /v1/models/{model}/status
 @app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
-async def model_status(model: str, req: Request, api_key: str = Depends(get_api_key)):
+async def model_status(model: str, req: Request):
     client_ip = req.client.host
     logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
     if model in Blackbox.models:
@@ -762,14 +680,14 @@ async def model_status(model: str, req: Request, api_key: str = Depends(get_api_
 
 # Endpoint: GET /v1/health
 @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
-async def health_check(req: Request, api_key: str = Depends(get_api_key)):
+async def health_check(req: Request):
     client_ip = req.client.host
     logger.info(f"Health check requested from IP: {client_ip}")
     return {"status": "ok"}
 
 # Endpoint: GET /v1/chat/completions (GET method)
 @app.get("/v1/chat/completions")
-async def chat_completions_get(req: Request, api_key: str = Depends(get_api_key)):
+async def chat_completions_get(req: Request):
     client_ip = req.client.host
     logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
     return RedirectResponse(url='about:blank')
@@ -794,4 +712,4 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 # Run the application
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run(app, host="0.0.0.0", port=8000)