Update app.py

app.py

@@ -18,10 +18,10 @@ app = FastAPI(
 
 # Define valid models (replace with actual models supported by https://models.github.ai/inference)
 VALID_MODELS = [
-    "deepseek/DeepSeek-V3-0324",  # Added based on your request
-    "gpt-3.5-turbo",              # Common model (placeholder)
-    "llama-3",                    # Common model (placeholder)
-    "mistral-7b"                  # Common model (placeholder)
+    "DeepSeek-V3-0324",  # Try without deepseek/ prefix
+    "gpt-3.5-turbo",     # Placeholder
+    "llama-3",           # Placeholder
+    "mistral-7b"         # Placeholder
 ]
 
 class GenerateRequest(BaseModel):
@@ -52,10 +52,13 @@ async def generate_ai_response(prompt: str, model: str, publisher: Optional[str]
         raise HTTPException(status_code=400, detail=f"Invalid model. Valid models: {VALID_MODELS}")
 
     logger.debug(f"Using endpoint: {endpoint}, publisher: {final_publisher}")
-    client = AsyncOpenAI(base_url=endpoint, api_key=token)
+    client = AsyncOpenAI(
+        base_url=endpoint,
+        api_key=token,
+        default_headers={"X-Publisher": final_publisher}  # Pass publisher as header
+    )
 
     try:
-        # Include publisher in the request payload
         stream = await client.chat.completions.create(
             messages=[
                 {"role": "system", "content": "You are a helpful assistant named Orion, created by Abdullah Ali"},
@@ -64,8 +67,7 @@ async def generate_ai_response(prompt: str, model: str, publisher: Optional[str]
             model=model,
             temperature=1.0,
             top_p=1.0,
-            stream=True,
-            extra_body={"publisher": final_publisher}  # Add publisher to extra_body
+            stream=True
         )
 
         async for chunk in stream:
@@ -74,12 +76,16 @@ async def generate_ai_response(prompt: str, model: str, publisher: Optional[str]
 
     except Exception as err:
         logger.error(f"AI generation failed: {str(err)}")
-        yield f"Error: {str(err)}"
-        raise HTTPException(status_code=500, detail=f"AI generation failed: {str(err)}")
+        # Check for specific AI server errors
+        error_msg = str(err)
+        if "unknown_model" in error_msg.lower():
+            raise HTTPException(status_code=400, detail=f"AI server error: {error_msg}")
+        yield f"Error: {error_msg}"
+        raise HTTPException(status_code=500, detail=f"AI generation failed: {error_msg}")
 
 @app.post("/generate", summary="Generate AI response", response_description="Streaming AI response")
 async def generate_response(
-    model: str = Query("deepseek/DeepSeek-V3-0324", description="The AI model to use"),
+    model: str = Query("DeepSeek-V3-0324", description="The AI model to use"),
     prompt: Optional[str] = Query(None, description="The input text prompt for the AI"),
     publisher: Optional[str] = Query(None, description="Publisher identifier (optional, defaults to DEFAULT_PUBLISHER env var)"),
     request: Optional[GenerateRequest] = None
@@ -87,7 +93,7 @@ async def generate_response(
     """
     Generate a streaming AI response based on the provided prompt, model, and publisher.
     
-    - **model**: The AI model to use (e.g., deepseek/DeepSeek-V3-0324)
+    - **model**: The AI model to use (e.g., DeepSeek-V3-0324)
     - **prompt**: The input text prompt for the AI (query param or body)
     - **publisher**: The publisher identifier (optional, defaults to DEFAULT_PUBLISHER env var)
     """
@@ -120,3 +126,4 @@ async def list_models():
 
 def get_app():
     return app
+