add-chat-completio

src/server/main.py

@@ -1281,6 +1281,174 @@ async def indic_custom_prompt_kannada_pdf(
     finally:
         # Close the temporary file to ensure it's fully written
         temp_file.close()
+from typing import List, Optional, Dict, Any
+
+class ChatCompletionRequest(BaseModel):
+    model: str = Field(default="gemma-3-12b-it", description="Model identifier (e.g., gemma-3-12b-it)")
+    messages: List[Dict[str, str]] = Field(..., description="List of messages in the conversation")
+    max_tokens: Optional[int] = Field(None, description="Maximum number of tokens to generate")
+    temperature: Optional[float] = Field(1.0, description="Sampling temperature")
+    top_p: Optional[float] = Field(1.0, description="Nucleus sampling parameter")
+    stream: Optional[bool] = Field(False, description="Whether to stream the response")
+
+# OpenAI-compatible response model
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: Dict[str, str]
+    finish_reason: Optional[str]
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Optional[Dict[str, int]] = None
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "id": "chatcmpl-123",
+                "object": "chat.completion",
+                "created": 1698765432,
+                "model": "gemma-3-12b-it",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "Hello! How can I assist you today?"
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 10,
+                    "total_tokens": 20
+                }
+            }
+        }
+
+# Helper function to convert OpenAI messages to a prompt for llama-server
+def messages_to_prompt(messages: List[Dict[str, str]]) -> str:
+    prompt = ""
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        if role == "system":
+            prompt += f"System: {content}\n"
+        elif role == "user":
+            prompt += f"User: {content}\n"
+        elif role == "assistant":
+            prompt += f"Assistant: {content}\n"
+    prompt += "Assistant: "
+    return prompt
+
+@app.post("/v1/chat/completions",
+          response_model=ChatCompletionResponse,
+          summary="OpenAI-Compatible Chat Completions",
+          description="Proxy endpoint to generate chat completions using llama-server with gemma-3-12b-it model, compatible with OpenAI's API.",
+          tags=["Chat"],
+          responses={
+              200: {"description": "Chat completion response", "model": ChatCompletionResponse},
+              400: {"description": "Invalid request parameters"},
+              500: {"description": "External llama-server error"},
+              504: {"description": "External llama-server timeout"}
+          })
+async def chat_completions(
+    request: Request,
+    body: ChatCompletionRequest
+):
+    logger.info("Processing chat completion request", extra={
+        "endpoint": "/v1/chat/completions",
+        "model": body.model,
+        "messages_count": len(body.messages),
+        "client_ip": request.client.host
+    })
+
+    # Validate messages
+    if not body.messages:
+        raise HTTPException(status_code=400, detail="Messages cannot be empty")
+
+    # Prepare payload for llama-server
+    # Adjust this based on the actual llama-server API requirements
+    llama_payload = {
+        "prompt": messages_to_prompt(body.messages),
+        "max_tokens": body.max_tokens if body.max_tokens is not None else 512,
+        "temperature": body.temperature,
+        "top_p": body.top_p,
+        "stream": body.stream
+    }
+
+    external_url = f"{os.getenv('DWANI_AI_LLM_URL')}/v1/chat/completions"
+
+    # llama-server endpoint (adjust if different)
+    start_time = time()
+
+    try:
+        response = requests.post(
+            external_url,
+            json=llama_payload,
+            headers={
+                "accept": "application/json",
+                "Content-Type": "application/json"
+            },
+            timeout=30
+        )
+        response.raise_for_status()
+
+        # Parse llama-server response
+        response_data = response.json()
+
+        # Transform llama-server response to OpenAI-compatible format
+        # Adjust based on actual response structure
+        completion_text = response_data.get("choices", [{}])[0].get("text", "")
+        finish_reason = response_data.get("choices", [{}])[0].get("finish_reason", "stop")
+
+        # Generate a unique ID for the response
+        completion_id = f"chatcmpl-{int(time.time())}"
+
+        # Build OpenAI-compatible response
+        openai_response = ChatCompletionResponse(
+            id=completion_id,
+            created=int(time.time()),
+            model=body.model,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message={
+                        "role": "assistant",
+                        "content": completion_text.strip()
+                    },
+                    finish_reason=finish_reason
+                )
+            ],
+            usage={
+                "prompt_tokens": len(llama_payload["prompt"].split()),  # Rough estimate
+                "completion_tokens": len(completion_text.split()),      # Rough estimate
+                "total_tokens": len(llama_payload["prompt"].split()) + len(completion_text.split())
+            }
+        )
+
+        logger.info(f"Chat completion successful in {time() - start_time:.2f} seconds", extra={
+            "response_length": len(completion_text)
+        })
+        return openai_response
+
+    except requests.Timeout:
+        logger.error("llama-server request timed out")
+        raise HTTPException(status_code=504, detail="llama-server timeout")
+    except requests.RequestException as e:
+        logger.error(f"llama-server request failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"llama-server error: {str(e)}")
+    except ValueError as e:
+        logger.error(f"Invalid JSON response from llama-server: {str(e)}")
+        raise HTTPException(status_code=500, detail="Invalid response format from llama-server")
+    except Exception as e:
+        logger.error(f"Unexpected error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
 
 
 if __name__ == "__main__":