Update main.py

main.py

@@ -1,7 +1,5 @@
 import os
 import re
-import random
-import string
 import uuid
 import json
 import logging
@@ -9,10 +7,11 @@ import asyncio
 import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator
+
 from datetime import datetime
 
 from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header, UploadFile, File
+from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
@@ -26,35 +25,26 @@ logger = logging.getLogger(__name__)
 
 # Load environment variables
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
-RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
+RATE_LIMIT_PER_MINUTE = int(os.getenv('RATE_LIMIT_PER_MINUTE', '60'))  # Requests per minute per IP
 
 if not API_KEYS or API_KEYS == ['']:
-    logger.error("No API keys found. Please set the API_KEYS environment variable. | NiansuhAI")
-    raise Exception("API_KEYS environment variable not set. | NiansuhAI")
-
-# Simple in-memory rate limiter
-rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+    logger.error("No API keys found. Please set the API_KEYS environment variable.")
+    raise Exception("API_KEYS environment variable not set.")
 
-async def get_api_key(authorization: str = Header(...)) -> str:
-    if not authorization.startswith('Bearer '):
-        logger.warning("Invalid authorization header format.")
-        raise HTTPException(status_code=401, detail='Invalid authorization header format | NiansuhAI')
-    api_key = authorization[7:]
-    if api_key not in API_KEYS:
-        logger.warning(f"Invalid API key attempted: {api_key}")
-        raise HTTPException(status_code=401, detail='Invalid API key | NiansuhAI')
-    return api_key
+# Simple in-memory rate limiter per IP
+rate_limit_store_ip = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
-async def rate_limiter(api_key: str = Depends(get_api_key)):
+async def rate_limiter(request: Request):
+    client_host = request.client.host
     current_time = time.time()
-    window_start = rate_limit_store[api_key]["timestamp"]
+    window_start = rate_limit_store_ip[client_host]["timestamp"]
     if current_time - window_start > 60:
-        rate_limit_store[api_key] = {"count": 1, "timestamp": current_time}
+        rate_limit_store_ip[client_host] = {"count": 1, "timestamp": current_time}
     else:
-        if rate_limit_store[api_key]["count"] >= RATE_LIMIT:
-            logger.warning(f"Rate limit exceeded for API key: {api_key}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded | NiansuhAI')
-        rate_limit_store[api_key]["count"] += 1
+        if rate_limit_store_ip[client_host]["count"] >= RATE_LIMIT_PER_MINUTE:
+            logger.warning(f"Rate limit exceeded for IP: {client_host}")
+            raise HTTPException(status_code=429, detail='Rate limit exceeded.')
+        rate_limit_store_ip[client_host]["count"] += 1
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -132,7 +122,7 @@ class Blackbox:
         'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
         'ReactAgent': {'mode': True, 'id': "React Agent"},
         'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
-        'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
+        'AngularJSAgent': {'mode': True, 'id": "AngularJS Agent"},
     }
 
     userSelectedModel = {
@@ -188,225 +178,361 @@ class Blackbox:
         else:
             return cls.default_model
 
-    # (Rest of the Blackbox class remains unchanged)
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        image: Any = None,
+        image_name: Optional[str] = None,
+        webSearchMode: bool = False,
+        **kwargs
+    ) -> AsyncGenerator[Any, None]:
+        model = cls.get_model(model)
+        logger.info(f"Selected model: {model}")
+
+        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",
+            "cache-control": "no-cache",
+            "content-type": "application/json",
+            "origin": cls.url,
+            "pragma": "no-cache",
+            "priority": "u=1, i",
+            "referer": cls.model_referers.get(model, cls.url),
+            "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
+            "sec-ch-ua-mobile": "?0",
+            "sec-ch-ua-platform": '"Linux"',
+            "sec-fetch-dest": "empty",
+            "sec-fetch-mode": "cors",
+            "sec-fetch-site": "same-origin",
+            "user-agent": "Mozilla/5.0 (X11; Linux x86_64)",
+        }
+
+        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.")
+                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'
+
+        # Don't log the full message content for privacy
+        logger.debug(f"Generated message ID: {random_id} for model: {model}")
+
+        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.")
+
+        data = {
+            "messages": messages,
+            "id": random_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": {},
+            "trendingAgentMode": {},
+            "isMicMode": False,
+            "userSystemPrompt": None,
+            "maxTokens": 99999999,
+            "playgroundTopP": 0.9,
+            "playgroundTemperature": 0.5,
+            "isChromeExt": False,
+            "githubToken": None,
+            "clickedAnswer2": False,
+            "clickedAnswer3": False,
+            "clickedForceWebSearch": False,
+            "visitFromDelta": False,
+            "mobileClient": False,
+            "userSelectedModel": None,
+            "webSearchMode": webSearchMode,
+        }
+
+        if model in cls.agentMode:
+            data["agentMode"] = cls.agentMode[model]
+        elif model in cls.trendingAgentMode:
+            data["trendingAgentMode"] = cls.trendingAgentMode[model]
+        elif model in cls.userSelectedModel:
+            data["userSelectedModel"] = cls.userSelectedModel[model]
+        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
+
+        timeout = ClientTimeout(total=30)  # Reduced timeout for faster response
+        retry_attempts = 3  # Reduced retry attempts for faster failure handling
+
+        for attempt in range(retry_attempts):
+            try:
+                async with ClientSession(headers=headers, timeout=timeout) as session:
+                    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 == 'ImageGeneration':
+                            response_text = await response.text()
+                            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.")
+                                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")
+                        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
+                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:
+                    raise HTTPException(status_code=502, detail="Error communicating with the external API.")
+            except asyncio.TimeoutError:
+                logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
+                if attempt == retry_attempts - 1:
+                    raise HTTPException(status_code=504, detail="External API request timed out.")
+            except Exception as e:
+                logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
+                if attempt == retry_attempts - 1:
+                    raise HTTPException(status_code=500, detail=str(e))
 
 # FastAPI app setup
 app = FastAPI()
 
+# Implement per-IP rate limiting middleware
+@app.middleware("http")
+async def rate_limit_middleware(request: Request, call_next):
+    await rate_limiter(request)
+    response = await call_next(request)
+    return response
+
+# Pydantic models for OpenAI API
 class Message(BaseModel):
     role: str
     content: str
 
-class ChatRequest(BaseModel):
+class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[Message]
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
     stream: Optional[bool] = False
-    webSearchMode: Optional[bool] = False
-
-def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
+    stop: Optional[Any] = None  # Can be a string or list of strings
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+
+def create_chat_completion_response(content: str, model: str, usage: Dict[str, int]) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
+        "object": "chat.completion",
         "created": int(datetime.now().timestamp()),
         "model": model,
         "choices": [
             {
                 "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": finish_reason,
+                "message": {
+                    "role": "assistant",
+                    "content": content
+                },
+                "finish_reason": "stop"
             }
         ],
-        "usage": None,
+        "usage": usage
     }
 
-@app.post("/niansuhai/v1/chat/completions", dependencies=[Depends(rate_limiter)])
-async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
-    # Redact user messages only for logging purposes
-    redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
+def create_stream_response_chunk(content: str, role: Optional[str] = None, finish_reason: Optional[str] = None):
+    delta = {}
+    if role:
+        delta['role'] = role
+    if content:
+        delta['content'] = content
+    return {
+        "object": "chat.completion.chunk",
+        "created": int(datetime.now().timestamp()),
+        "model": "",  # Model name can be added if necessary
+        "choices": [
+            {
+                "delta": delta,
+                "index": 0,
+                "finish_reason": finish_reason
+            }
+        ]
+    }
 
-    logger.info(f"Received chat completions request from API key: {api_key} | Model: {request.model} | Messages: {redacted_messages}")
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest, authorization: str = Header(None)):
+    # Verify API key
+    if not authorization or not authorization.startswith('Bearer '):
+        logger.warning("Invalid authorization header format.")
+        raise HTTPException(status_code=401, detail='Invalid authorization header format.')
+    api_key = authorization[7:]
+    if api_key not in API_KEYS:
+        logger.warning(f"Invalid API key attempted: {api_key}")
+        raise HTTPException(status_code=401, detail='Invalid API key.')
+
+    logger.info(f"Received chat completion request for model: {request.model}")
+
+    # Validate model
+    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}")
+        raise HTTPException(status_code=400, detail="The model is not available.")
 
+    # Process the request
     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}")
-            raise HTTPException(status_code=400, detail="Requested model is not available. | NiansuhAI")
-
-        # 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
-            image=None,
-            image_name=None,
-            webSearchMode=request.webSearchMode
+        # Convert messages to dicts
+        messages = [msg.dict() for msg in request.messages]
+
+        # Check if the user is requesting image generation
+        image_generation_requested = any(
+            re.search(r'\b(generate|create|draw)\b.*\b(image|picture|art)\b', msg['content'], re.IGNORECASE)
+            for msg in messages if msg['role'] == 'user'
         )
 
+        if image_generation_requested:
+            model = 'ImageGeneration'
+            # For image generation, use the last message as prompt
+            prompt = messages[-1]['content']
+            # Build messages for the Blackbox.create_async_generator
+            messages = [{"role": "user", "content": prompt}]
+            async_generator = Blackbox.create_async_generator(
+                model=model,
+                messages=messages,
+                image=None,
+                image_name=None,
+                webSearchMode=False
+            )
+
+            # Collect images
+            images = []
+            count = 0
+            async for response in async_generator:
+                if isinstance(response, ImageResponse):
+                    images.append(response.url)
+                    count += 1
+                    if count >= request.n:
+                        break
+
+            # Build response content with image URLs
+            response_content = "\n".join(f"![Generated Image]({url})" for url in images)
+            completion_tokens = len(response_content.split())
+        else:
+            # Use the requested model
+            async_generator = Blackbox.create_async_generator(
+                model=request.model,
+                messages=messages,
+                image=None,
+                image_name=None,
+                webSearchMode=False
+            )
+            # Usage tracking
+            completion_tokens = 0  # Will be updated as we process the response
+
+        prompt_tokens = sum(len(msg['content'].split()) for msg in messages)
+
         if request.stream:
             async def generate():
+                nonlocal completion_tokens
                 try:
+                    # Initial delta with role
+                    initial_chunk = create_stream_response_chunk(content=None, role="assistant")
+                    yield f"data: {json.dumps(initial_chunk)}\n\n"
+
                     async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponse):
-                            image_markdown = f"![image]({chunk.url})"
-                            response_chunk = create_response(image_markdown, request.model)
+                        if isinstance(chunk, str):
+                            completion_tokens += len(chunk.split())
+                            response_chunk = create_stream_response_chunk(content=chunk)
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                        elif isinstance(chunk, ImageResponse):
+                            content = f"![Generated Image]({chunk.url})"
+                            completion_tokens += len(content.split())
+                            response_chunk = create_stream_response_chunk(content=content)
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
                         else:
-                            response_chunk = create_response(chunk, request.model)
-                        
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
-                    
+                            pass  # Handle other types if necessary
+                    # Finish reason
+                    final_chunk = create_stream_response_chunk(content=None, finish_reason="stop")
+                    yield f"data: {json.dumps(final_chunk)}\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("Error during streaming response generation. | NiansuhAI")
-                    error_response = {"error": str(e)}
-                    yield f"data: {json.dumps(error_response)}\n\n"
-
+                    logger.exception("Error during streaming response generation.")
+                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
             return StreamingResponse(generate(), media_type="text/event-stream")
         else:
             response_content = ""
             async for chunk in async_generator:
-                if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.url})\n"
-                else:
+                if isinstance(chunk, str):
                     response_content += chunk
-
-            logger.info(f"Completed non-streaming response generation for API key: {api_key}")
-            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": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(response_content.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
-                },
+                elif isinstance(chunk, ImageResponse):
+                    response_content += f"![Generated Image]({chunk.url})\n"
+            completion_tokens = len(response_content.split())
+            usage = {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens
             }
+            return create_chat_completion_response(response_content, request.model, usage)
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e}")
         raise HTTPException(status_code=503, detail=str(e))
-    except HTTPException as he:
-        logger.warning(f"HTTPException: {he.detail}")
-        raise he
     except Exception as e:
-        logger.exception("An unexpected error occurred while processing the chat completions request. | NiansuhAI")
+        logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/niansuhai/v1/models", dependencies=[Depends(rate_limiter)])
-async def get_models(api_key: str = Depends(get_api_key)):
-    logger.info(f"Fetching available models for API key: {api_key}")
-    return {"data": [{"id": model} for model in Blackbox.models]}
-
-# Additional endpoints for better functionality
-
-@app.get("/niansuhai/v1/health", dependencies=[Depends(rate_limiter)])
-async def health_check(api_key: str = Depends(get_api_key)):
-    logger.info(f"Health check requested by API key: {api_key}")
-    return {"status": "ok"}
-
-@app.get("/niansuhai/v1/models/{model}/status", dependencies=[Depends(rate_limiter)])
-async def model_status(model: str, api_key: str = Depends(get_api_key)):
-    logger.info(f"Model status requested for '{model}' by API key: {api_key}")
-    if model in Blackbox.models:
-        return {"model": model, "status": "available | NiansuhAI"}
-    elif model in Blackbox.model_aliases:
-        actual_model = Blackbox.model_aliases[model]
-        return {"model": actual_model, "status": "available via alias | NiansuhAI"}
-    else:
-        logger.warning(f"Model not found: {model}")
-        raise HTTPException(status_code=404, detail="Model not found | NiansuhAI")
-
-# New endpoint to get model details
-@app.get("/niansuhai/v1/models/{model}/details", dependencies=[Depends(rate_limiter)])
-async def get_model_details(model: str, api_key: str = Depends(get_api_key)):
-    logger.info(f"Model details requested for '{model}' by API key: {api_key}")
-    actual_model = Blackbox.get_model(model)
-    if actual_model not in Blackbox.models:
-        logger.warning(f"Model not found: {model}")
-        raise HTTPException(status_code=404, detail="Model not found | NiansuhAI")
-    # For demonstration, we'll return mock details
-    model_details = {
-        "id": actual_model,
-        "description": f"Details about model {actual_model}",
-        "capabilities": ["chat", "completion", "image generation"] if actual_model in Blackbox.image_models else ["chat", "completion"],
-        "status": "available",
-    }
-    return {"data": model_details}
-
-# Session history endpoints
-session_histories = defaultdict(list)  # In-memory storage for session histories
-
-@app.post("/niansuhai/v1/sessions/{session_id}/messages", dependencies=[Depends(rate_limiter)])
-async def add_message_to_session(session_id: str, message: Message, api_key: str = Depends(get_api_key)):
-    logger.info(f"Adding message to session '{session_id}' by API key: {api_key}")
-    session_histories[session_id].append({"role": message.role, "content": message.content})
-    return {"status": "message added"}
-
-@app.get("/niansuhai/v1/sessions/{session_id}/messages", dependencies=[Depends(rate_limiter)])
-async def get_session_messages(session_id: str, api_key: str = Depends(get_api_key)):
-    logger.info(f"Fetching messages for session '{session_id}' by API key: {api_key}")
-    messages = session_histories.get(session_id)
-    if messages is None:
-        raise HTTPException(status_code=404, detail="Session not found | NiansuhAI")
-    return {"data": messages}
-
-# User preferences endpoints
-user_preferences = defaultdict(dict)  # In-memory storage for user preferences
-
-class UserPreferences(BaseModel):
-    theme: Optional[str] = "light"
-    notifications_enabled: Optional[bool] = True
-
-@app.post("/niansuhai/v1/users/{user_id}/preferences", dependencies=[Depends(rate_limiter)])
-async def update_user_preferences(user_id: str, preferences: UserPreferences, api_key: str = Depends(get_api_key)):
-    logger.info(f"Updating preferences for user '{user_id}' by API key: {api_key}")
-    user_preferences[user_id] = preferences.dict()
-    return {"status": "preferences updated"}
-
-@app.get("/niansuhai/v1/users/{user_id}/preferences", dependencies=[Depends(rate_limiter)])
-async def get_user_preferences(user_id: str, api_key: str = Depends(get_api_key)):
-    logger.info(f"Fetching preferences for user '{user_id}' by API key: {api_key}")
-    preferences = user_preferences.get(user_id)
-    if preferences is None:
-        raise HTTPException(status_code=404, detail="User not found | NiansuhAI")
-    return {"data": preferences}
-
-# Image upload endpoint
-@app.post("/niansuhai/v1/images/upload", dependencies=[Depends(rate_limiter)])
-async def upload_image(image: UploadFile = File(...), api_key: str = Depends(get_api_key)):
-    logger.info(f"Image upload requested by API key: {api_key}")
-    if not image.content_type.startswith('image/'):
-        logger.warning("Uploaded file is not an image.")
-        raise HTTPException(status_code=400, detail="Uploaded file is not an image | NiansuhAI")
-    # For demonstration, we'll just return the filename
-    return {"filename": image.filename, "status": "image uploaded"}
-
-# Component health check endpoint
-@app.get("/niansuhai/v1/health/{component}", dependencies=[Depends(rate_limiter)])
-async def component_health_check(component: str, api_key: str = Depends(get_api_key)):
-    logger.info(f"Health check for component '{component}' requested by API key: {api_key}")
-    # Mock health status for components
-    components_status = {
-        "database": "healthy",
-        "message_queue": "healthy",
-        "cache": "healthy",
-    }
-    status = components_status.get(component)
-    if status is None:
-        logger.warning(f"Component not found: {component}")
-        raise HTTPException(status_code=404, detail="Component not found | NiansuhAI")
-    return {"component": component, "status": status}
+@app.get("/v1/models")
+async def get_models(authorization: str = Header(None)):
+    # Verify API key
+    if not authorization or not authorization.startswith('Bearer '):
+        logger.warning("Invalid authorization header format.")
+        raise HTTPException(status_code=401, detail='Invalid authorization header format.')
+    api_key = authorization[7:]
+    if api_key not in API_KEYS:
+        logger.warning(f"Invalid API key attempted: {api_key}")
+        raise HTTPException(status_code=401, detail='Invalid API key.')
+
+    logger.info("Fetching available models.")
+    # Return models in OpenAI format
+    models_data = [{"id": model, "object": "model", "owned_by": "organization-owner", "permission": []} for model in Blackbox.models]
+    return {"data": models_data, "object": "list"}
 
 if __name__ == "__main__":
     import uvicorn