add-chat-completio

requirements.txt

@@ -8,4 +8,5 @@ pillow
 pyjwt
 sqlalchemy
 passlib[bcrypt]
-pycryptodome
+pycryptodome
+openai

src/server/main.py

@@ -1283,10 +1283,12 @@ async def indic_custom_prompt_kannada_pdf(
         temp_file.close()
 from typing import List, Optional, Dict, Any
 
+from openai import AsyncOpenAI, OpenAIError
+# OpenAI-compatible request model
 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")
+    model: str = Field(default="gemma-3-12b-it", description="Model identifier")
+    messages: List[Dict[str, str]] = Field(..., description="List of messages")
+    max_tokens: Optional[int] = Field(None, description="Maximum 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")
@@ -1305,151 +1307,86 @@ class ChatCompletionResponse(BaseModel):
     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
+# Initialize OpenAI client
+openai_client = AsyncOpenAI(
+    base_url=os.getenv("DWANI_AI_LLM_URL"),  # e.g., https://<ngrok-url>.ngrok.io or http://localhost:7860
+    api_key=os.getenv("DWANI_AI_LLM_API_KEY", ""),  # Optional API key
+    timeout=30.0
+)
 
 @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={
+          description="Proxies chat completions to llama-server using OpenAI API format.",
+          tags=["Chat"])
+async def chat_completions(request: Request, body: ChatCompletionRequest):
+    logger.info("Received chat completion request", extra={
         "endpoint": "/v1/chat/completions",
         "model": body.model,
-        "messages_count": len(body.messages),
+        "messages": body.messages,
         "client_ip": request.client.host
     })
 
     # Validate messages
     if not body.messages:
+        logger.error("Messages field is empty", extra={"client_ip": request.client.host})
         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
+        # Proxy request to llama-server using OpenAI client
+        response = await openai_client.chat.completions.create(
+            model=body.model,
+            messages=body.messages,
+            max_tokens=body.max_tokens,
+            temperature=body.temperature,
+            top_p=body.top_p,
+            stream=body.stream
         )
-        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())}"
+        # Streaming not supported in this simple version
+        if body.stream:
+            logger.error("Streaming requested but not supported")
+            raise HTTPException(status_code=400, detail="Streaming not supported")
 
-        # Build OpenAI-compatible response
+        # Map OpenAI response to Pydantic model
         openai_response = ChatCompletionResponse(
-            id=completion_id,
-            created=int(time.time()),
-            model=body.model,
+            id=response.id,
+            created=response.created,
+            model=response.model,
             choices=[
                 ChatCompletionChoice(
-                    index=0,
+                    index=choice.index,
                     message={
-                        "role": "assistant",
-                        "content": completion_text.strip()
+                        "role": choice.message.role,
+                        "content": choice.message.content
                     },
-                    finish_reason=finish_reason
-                )
+                    finish_reason=choice.finish_reason
+                ) for choice in response.choices
             ],
-            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())
-            }
+            usage=(
+                {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens
+                } if response.usage else None
+            )
         )
 
         logger.info(f"Chat completion successful in {time() - start_time:.2f} seconds", extra={
-            "response_length": len(completion_text)
+            "response_length": len(response.choices[0].message.content if response.choices else 0)
         })
         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 OpenAIError as e:
+        logger.error(f"llama-server error: {str(e)}", extra={"client_ip": request.client.host})
+        status_code = 504 if "timeout" in str(e).lower() else 500
+        raise HTTPException(status_code=status_code, detail=f"llama-server error: {str(e)}")
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
+        logger.error(f"Internal error: {str(e)}", extra={"client_ip": request.client.host})
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
-
-
+    
 
 if __name__ == "__main__":
     # Ensure EXTERNAL_API_BASE_URL is set