Update app.py

app.py

@@ -1,43 +1,49 @@
+import os
+import uuid
+from datetime import datetime
+from typing import List
+
+from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException, Depends, Security
 from fastapi.security import APIKeyHeader
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-import requests
-import uuid
-from datetime import datetime
 import uvicorn
-import os
+from huggingface_hub import InferenceClient
+
+# Load environment variables
+load_dotenv()
 
 app = FastAPI()
 
 # Configuration
 API_KEY_NAME = "X-API-KEY"
-API_KEYS = os.getenv("API_KEYS", "").split(",")  # Comma-separated list of valid API keys
+API_KEYS = os.getenv("API_KEYS", "").split(",")
+HF_MODEL = os.getenv("HF_MODEL_NAME", "deepseek-ai/deepseek-llm-7b")
+HF_TOKEN = os.getenv("HF_API_TOKEN")
+
+# Initialize Hugging Face client
+client = InferenceClient(token=HF_TOKEN)
 
 # Security
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
-async def get_api_key(api_key_header: str = Security(api_key_header)):
-    if not api_key_header:
-        raise HTTPException(
-            status_code=401, 
-            detail="API key is missing"
-        )
-    if api_key_header not in API_KEYS:
-        raise HTTPException(
-            status_code=401, 
-            detail="Invalid API key"
-        )
-    return api_key_header
+async def get_api_key(api_key: str = Security(api_key_header)):
+    if not api_key:
+        raise HTTPException(status_code=401, detail="API key is missing")
+    if api_key not in API_KEYS:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    return api_key
 
 class Message(BaseModel):
     role: str
     content: str
 
 class ChatRequest(BaseModel):
-    messages: list[Message]
-    model: str = "deepseek-ai/DeepSeek-R1"
+    messages: List[Message]
+    model: str = HF_MODEL
     temperature: float = 0.7
+    max_tokens: int = 200
 
 @app.get("/")
 async def health_check():
@@ -49,23 +55,19 @@ async def chat_completion(
     api_key: str = Depends(get_api_key)
 ):
     try:
-        response = requests.post(
-            "https://www.multichatai.com/api/chat/deepinfra",
-            json={
-                "chatSettings": {
-                    "model": request.model,
-                    "prompt": "You are helpful.",
-                    "temperature": request.temperature,
-                    "contextLength": 4096
-                },
-                "messages": [
-                    {"role": "system", "content": "You are helpful."},
-                    *[{"role": m.role, "content": m.content} for m in request.messages]
-                ]
-            },
-            timeout=10
+        # Format messages for HF inference
+        prompt = "\n".join(
+            f"{msg.role}: {msg.content}" 
+            for msg in request.messages
+        )
+        
+        # Call Hugging Face model
+        response = client.text_generation(
+            prompt=prompt,
+            model=request.model,
+            temperature=request.temperature,
+            max_new_tokens=request.max_tokens
         )
-        response.raise_for_status()
         
         return JSONResponse({
             "id": str(uuid.uuid4()),
@@ -75,7 +77,7 @@ async def chat_completion(
             "choices": [{
                 "message": {
                     "role": "assistant",
-                    "content": response.text.strip()
+                    "content": response.strip()
                 }
             }]
         })
@@ -83,4 +85,7 @@ async def chat_completion(
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.getenv("PORT", 7860))