CDP token for CAII inference

Author: Khauneesh-AI

.project-metadata.yaml

@@ -30,6 +30,10 @@ environment_variables:
     default: "your huggingface username"
     description: >-
       hf_username
+  CDP_TOKEN:
+    default: "API key for Cloudera AI Inference"
+    description: >-
+      CDP_TOKEN
   
 
 
@@ -69,7 +73,7 @@ tasks:
     script: build/build_client.py
     arguments: None
     cpu: 2
-    memory: 2
+    memory: 4
     short_summary: Create job to build client application
     environment:
       TASK_TYPE: CREATE/RUN_JOB

app/core/config.py

@@ -5,6 +5,10 @@
 import requests
 import json
 from fastapi.responses import JSONResponse
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+load_dotenv() 
 
 class UseCase(str, Enum):
     CODE_GENERATION = "code_generation"
@@ -281,18 +285,55 @@ def get_examples_for_topic(use_case: UseCase, topic: str) -> List[Dict[str, str]
     }
 }
 
+JWT_PATH = Path("/tmp/jwt")
+
+def _get_caii_token() -> str:
+    if (tok := os.getenv("CDP_TOKEN")):
+        return tok
+    try:
+        payload = json.loads(open(JWT_PATH).read())
+    except FileNotFoundError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="No CDP_TOKEN env‑var and no /tmp/jwt file")
+    except json.JSONDecodeError:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Malformed /tmp/jwt")
+
+    if not (tok := payload.get("access_token")):
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="access_token missing in /tmp/jwt")
+    return tok
+
+def caii_check(endpoint: str, timeout: int = 3) -> requests.Response:
+    """
+    Return the GET /models response if everything is healthy.
+    Raise HTTPException on *any* problem.
+    """
+    if not endpoint:
+        raise HTTPException(400, "CAII endpoint not provided")
+
+    token = _get_caii_token()
+    url = endpoint.removesuffix("/chat/completions") + "/models"
+
+    try:
+        r = requests.get(url,
+                         headers={"Authorization": f"Bearer {token}"},
+                         timeout=timeout)
+    except requests.exceptions.RequestException as exc:
+        raise HTTPException(503, f"CAII endpoint unreachable: {exc}")
+
+    if r.status_code in (401, 403):
+        raise HTTPException(403, "Token is valid but has no access to this environment")
+    if r.status_code == 404:
+        raise HTTPException(404, "CAII endpoint or resource not found")
+    if 500 <= r.status_code < 600:
+        raise HTTPException(503, "CAII endpoint is downscaled; retry in ~15 min")
+    if r.status_code != 200:
+        raise HTTPException(r.status_code, r.text)
+
+    return r
 
-def caii_check(caii_endpoint):
-    API_KEY = json.load(open("/tmp/jwt"))["access_token"]
-    headers = {
-        "Authorization": f"Bearer {API_KEY}"
-    }
-    
-    
-    if caii_endpoint:
-        caii_endpoint = caii_endpoint.removesuffix('/chat/completions') 
-        caii_endpoint = caii_endpoint + "/models"
-        response = requests.get(caii_endpoint, headers=headers, timeout=3)  # Will raise RequestException if fails
-        
-    return response
 

app/core/model_handlers.py

@@ -11,6 +11,7 @@
 from openai import OpenAI
 from app.core.exceptions import APIError, InvalidModelError, ModelHandlerError, JSONParsingError
 from app.core.telemetry_integration import track_llm_operation
+from app.core.config import  _get_caii_token
 
 
 
@@ -280,7 +281,8 @@ def _handle_bedrock_request(self, prompt: str, retry_with_reduced_tokens: bool):
     def _handle_caii_request(self, prompt: str):
         """Original CAII implementation"""
         try:
-            API_KEY = json.load(open("/tmp/jwt"))["access_token"]
+            #API_KEY = json.load(open("/tmp/jwt"))["access_token"]
+            API_KEY = _get_caii_token()
             MODEL_ID = self.model_id
             caii_endpoint = self.caii_endpoint
 

app/main.py

@@ -408,15 +408,10 @@ async def generate_examples(request: SynthesisRequest):
     # Generate a request ID
     request_id = str(uuid.uuid4())
 
-    if request.inference_type== "CAII":
+    if request.inference_type == "CAII":
         caii_endpoint = request.caii_endpoint
-        response = caii_check(caii_endpoint)
-        message = "The CAII endpoint you are tring to reach is downscaled, please try after >15 minutes while it autoscales, meanwhile please try another model"
-        if response.status_code != 200:
-                    return JSONResponse(
-                        status_code=503,  # Service Unavailable
-                        content={"status": "failed", "error": message}
-                    )
+       
+        caii_check(caii_endpoint)
 
 
     is_demo = request.is_demo
@@ -464,13 +459,7 @@ async def generate_freeform_data(request: SynthesisRequest):
 
     if request.inference_type == "CAII":
         caii_endpoint = request.caii_endpoint
-        response = caii_check(caii_endpoint)
-        message = "The CAII endpoint you are trying to reach is downscaled, please try after >15 minutes while it autoscales, meanwhile please try another model"
-        if response.status_code != 200:
-            return JSONResponse(
-                status_code=503,  # Service Unavailable
-                content={"status": "failed", "error": message}
-            )
+        caii_check(caii_endpoint)
 
     is_demo = request.is_demo
     mem = 4
@@ -514,15 +503,9 @@ async def evaluate_examples(request: EvaluationRequest):
     """Evaluate generated QA pairs"""
     request_id = str(uuid.uuid4())
 
-    if request.inference_type== "CAII":
+    if request.inference_type == "CAII":
         caii_endpoint = request.caii_endpoint
-        response = caii_check(caii_endpoint)
-        message = "The CAII endpoint you are tring to reach is downscaled, please try after >15 minutes while it autoscales, meanwhile please try another model"
-        if response.status_code != 200:
-                    return JSONResponse(
-                        status_code=503,  # Service Unavailable
-                        content={"status": "failed", "error": message}
-                    )
+        caii_check(caii_endpoint)
 
     is_demo = request.is_demo
     if is_demo:
@@ -541,13 +524,8 @@ async def evaluate_freeform(request: EvaluationRequest):
 
     if request.inference_type == "CAII":
         caii_endpoint = request.caii_endpoint
-        response = caii_check(caii_endpoint)
-        message = "The CAII endpoint you are trying to reach is downscaled, please try after >15 minutes while it autoscales, meanwhile please try another model"
-        if response.status_code != 200:
-            return JSONResponse(
-                status_code=503,  # Service Unavailable
-                content={"status": "failed", "error": message}
-            )
+        caii_check(caii_endpoint)
+        
 
     is_demo = getattr(request, 'is_demo', True)
     if is_demo: