Update app.py

app.py

@@ -12,22 +12,41 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
+# Define valid models (replace with actual models supported by your AI server)
+VALID_MODELS = ["default-model", "another-model"]  # Update this list
+
 class GenerateRequest(BaseModel):
     prompt: str
+    publisher: Optional[str] = None  # Allow publisher in the body if needed
 
-async def generate_ai_response(prompt: str, model: str):
-    logger.debug(f"Received prompt: {prompt}, model: {model}")
+async def generate_ai_response(prompt: str, model: str, publisher: Optional[str]):
+    logger.debug(f"Received prompt: {prompt}, model: {model}, publisher: {publisher}")
+    
+    # Configuration for AI endpoint
     token = os.getenv("GITHUB_TOKEN")
     endpoint = os.getenv("AI_SERVER_URL", "https://models.github.ai/inference")
+    default_publisher = os.getenv("DEFAULT_PUBLISHER", "abdullahalioo")  # Fallback publisher
     
     if not token:
         logger.error("GitHub token not configured")
         raise HTTPException(status_code=500, detail="GitHub token not configured")
 
-    logger.debug(f"Using endpoint: {endpoint}")
+    # Use provided publisher or fallback to environment variable
+    final_publisher = publisher or default_publisher
+    if not final_publisher:
+        logger.error("Publisher is required")
+        raise HTTPException(status_code=400, detail="Publisher is required")
+
+    # Validate model
+    if model not in VALID_MODELS:
+        logger.error(f"Invalid model: {model}. Valid models: {VALID_MODELS}")
+        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)
 
     try:
+        # Include publisher in the request payload (modify as needed based on AI server requirements)
         stream = await client.chat.completions.create(
             messages=[
                 {"role": "system", "content": "You are a helpful assistant named Orion, created by Abdullah Ali"},
@@ -36,7 +55,8 @@ async def generate_ai_response(prompt: str, model: str):
             model=model,
             temperature=1.0,
             top_p=1.0,
-            stream=True
+            stream=True,
+            extra_body={"publisher": final_publisher}  # Add publisher to extra_body
         )
 
         async for chunk in stream:
@@ -52,10 +72,22 @@ async def generate_ai_response(prompt: str, model: str):
 async def generate_response(
     model: str = Query("default-model", 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
 ):
-    logger.debug(f"Request received - model: {model}, prompt: {prompt}, body: {request}")
+    """
+    Generate a streaming AI response based on the provided prompt, model, and publisher.
+    
+    - **model**: The AI model to use (e.g., default-model)
+    - **prompt**: The input text prompt for the AI (query param or body)
+    - **publisher**: The publisher identifier (optional, defaults to DEFAULT_PUBLISHER env var)
+    """
+    logger.debug(f"Request received - model: {model}, prompt: {prompt}, publisher: {publisher}, body: {request}")
+    
+    # Determine prompt source: query parameter or request body
     final_prompt = prompt if prompt is not None else (request.prompt if request is not None else None)
+    # Determine publisher source: query parameter or request body
+    final_publisher = publisher if publisher is not None else (request.publisher if request is not None else None)
     
     if not final_prompt or not final_prompt.strip():
         logger.error("Prompt cannot be empty")
@@ -66,7 +98,7 @@ async def generate_response(
         raise HTTPException(status_code=400, detail="Model cannot be empty")
     
     return StreamingResponse(
-        generate_ai_response(final_prompt, model),
+        generate_ai_response(final_prompt, model, final_publisher),
         media_type="text/event-stream"
     )