Update routers/llm_chat.py

routers/llm_chat.py

@@ -1,17 +1,15 @@
 # routers/llm_chat.py
 from fastapi import APIRouter, HTTPException, Header
-import requests
-import os
 from pydantic import BaseModel
+from helpers.ai_client import AIClient
 
 router = APIRouter(
     prefix="/api/v1",  # Prefix for all routes in this router
     tags=["LLM Chat"],  # Tag for OpenAPI documentation
 )
 
-# Load environment variables
-LLM_API_URL = os.getenv("LLM_API_URL")
-API_KEY = os.getenv("X_API_KEY")
+# Initialize the AI client
+ai_client = AIClient()
 
 # Pydantic model for request validation
 class LLMChatRequest(BaseModel):
@@ -26,26 +24,17 @@ async def llm_chat(
     request: LLMChatRequest,
     x_api_key: str = Header(None, description="API Key for authentication")
 ):
-    if x_api_key != API_KEY:
-        raise HTTPException(status_code=403, detail="Invalid API Key")
-
-    payload = {
-        "prompt": request.prompt,
-        "system_message": request.system_message,
-        "model_id": request.model_id,
-        "conversation_id": request.conversation_id,
-        "user_id": request.user_id
-    }
-
-    headers = {
-        "accept": "application/json",
-        "X-API-Key": x_api_key,
-        "Content-Type": "application/json"
-    }
-
-    # Use requests to call the external API
-    response = requests.post(LLM_API_URL, json=payload, headers=headers)
-    if response.status_code != 200:
-        raise HTTPException(status_code=response.status_code, detail="Error from LLM API")
-
-    return response.json()
+    try:
+        # Use the AI client to send the prompt
+        response = ai_client.send_prompt(
+            prompt=request.prompt,
+            system_message=request.system_message,
+            model_id=request.model_id,
+            conversation_id=request.conversation_id,
+            user_id=request.user_id,
+            api_key=x_api_key
+        )
+        return response
+    except Exception as e:
+        logger.error(f"Error in llm_chat: {e}")
+        raise HTTPException(status_code=500, detail=str(e))