Update main.py

main.py

@@ -27,11 +27,18 @@ 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
+AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '')  # Comma-separated available models
 
 if not API_KEYS or API_KEYS == ['']:
     logger.error("No API keys found. Please set the API_KEYS environment variable.")
     raise Exception("API_KEYS environment variable not set.")
 
+# Process available models
+if AVAILABLE_MODELS:
+    AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
+else:
+    AVAILABLE_MODELS = []  # If empty, all models are available
+
 # Simple in-memory rate limiter
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 ip_rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
@@ -122,6 +129,10 @@ class Blackbox:
         '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"},
@@ -194,12 +205,12 @@ class Blackbox:
     def get_model(cls, model: str) -> str:
         if model in cls.models:
             return model
-        elif model in cls.userSelectedModel:
+        elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
             return model
-        elif model in cls.model_aliases:
+        elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
             return cls.model_aliases[model]
         else:
-            return cls.default_model
+            return cls.default_model if cls.default_model in cls.models else None
 
     @classmethod
     async def create_async_generator(
@@ -213,6 +224,10 @@ class Blackbox:
         **kwargs
     ) -> AsyncGenerator[Any, None]:
         model = cls.get_model(model)
+        if model is None:
+            logger.error(f"Model {model} is not available.")
+            raise ModelNotWorkingException(model)
+
         logger.info(f"Selected model: {model}")
 
         if not cls.working or model not in cls.models:
@@ -477,23 +492,23 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/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}")
+@app.get("/v1/models")
+async def get_models():
+    logger.info("Fetching available models")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
 # Additional endpoints for better functionality
 @app.get("/v1/health", dependencies=[Depends(rate_limiter)])
-async def health_check(api_key: str = Depends(get_api_key)):
+async def health_check(req: Request, api_key: str = Depends(get_api_key)):
     logger.info(f"Health check requested by API key: {api_key}")
     return {"status": "ok"}
 
-@app.get("/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}")
+@app.get("/v1/models/{model}/status")
+async def model_status(model: str):
+    logger.info(f"Model status requested for '{model}'")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
-    elif model in Blackbox.model_aliases:
+    elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
         actual_model = Blackbox.model_aliases[model]
         return {"model": actual_model, "status": "available via alias"}
     else:
@@ -515,6 +530,73 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         },
     )
 
+# New endpoint: /v1/tokenizer to calculate token counts
+class TokenizerRequest(BaseModel):
+    text: str
+
+@app.post("/v1/tokenizer")
+async def tokenizer(request: TokenizerRequest):
+    text = request.text
+    token_count = len(text.split())
+    return {"text": text, "tokens": token_count}
+
+# New endpoint: /v1/completions to support text completions
+class CompletionRequest(BaseModel):
+    model: str
+    prompt: str
+    max_tokens: Optional[int] = 16
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    logprobs: Optional[int] = None
+    echo: Optional[bool] = False
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    best_of: Optional[int] = 1
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+
+@app.post("/v1/completions", dependencies=[Depends(rate_limiter)])
+async def completions(request: CompletionRequest, req: Request, api_key: str = Depends(get_api_key)):
+    logger.info(f"Received completion request from API key: {api_key} | Model: {request.model}")
+
+    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.")
+
+        # Simulate a simple completion by echoing the prompt
+        completion_text = f"{request.prompt} [Completed by {request.model}]"
+
+        return {
+            "id": f"cmpl-{uuid.uuid4()}",
+            "object": "text_completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "text": completion_text,
+                    "index": 0,
+                    "logprobs": None,
+                    "finish_reason": "length"
+                }
+            ],
+            "usage": {
+                "prompt_tokens": len(request.prompt.split()),
+                "completion_tokens": len(completion_text.split()),
+                "total_tokens": len(request.prompt.split()) + len(completion_text.split())
+            }
+        }
+    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 completions request.")
+        raise HTTPException(status_code=500, detail=str(e))
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)