Add is_api_key at the query parameter for system prompt

Author: sm-stack

src/app/api/helper/get_system_prompt.py

@@ -6,12 +6,12 @@
 
 _system_prompts_cache = TTLCache(maxsize=100, ttl=3)  # 5 minutes TTL
 
-async def get_system_prompt(model: str, usage: str) -> str:
+async def get_system_prompt(model: str, usage: str, is_api_key: bool = False) -> str:
     """
     Get the system prompt for the model and usage.
     """
     try:
-        cached_prompts = _system_prompts_cache.get(f"prompt-{model}-{usage}")
+        cached_prompts = _system_prompts_cache.get(f"prompt-{model}-{usage}-{is_api_key}")
         if cached_prompts is not None:
             log.info("Retrieved system prompts from cache")
             return cached_prompts
@@ -20,7 +20,7 @@ async def get_system_prompt(model: str, usage: str) -> str:
         api_key = get_settings().PANDA_APP_SERVER_TOKEN
         client = httpx.AsyncClient()
         response = await client.get(
-            f"{base_url}/system-prompt?model={model}&usage={usage}",
+            f"{base_url}/system-prompt?model={model}&usage={usage}&is_api_key={is_api_key}",
             headers={"X-API-Key": f"{api_key}"}
         )
 
@@ -34,7 +34,7 @@ async def get_system_prompt(model: str, usage: str) -> str:
             else:
                 log.error(f"Failed to get system prompt for model {model} and usage {usage}", response.text)
                 raise HTTPException(status_code=500, detail="Failed to get system prompt")
-        _system_prompts_cache[f"prompt-{model}-{usage}"] = response.json()["system_prompt"]
+        _system_prompts_cache[f"prompt-{model}-{usage}-{is_api_key}"] = response.json()["system_prompt"]
         return response.json()["system_prompt"]
     except Exception as e:
         log.error(f"Error getting system prompt for model {model} and usage {usage}: {str(e)}")

src/app/api/helper/request_llm.py

@@ -14,7 +14,14 @@
 LLMSuccessResponse = Union[Dict[str, Any], List[Any]]
 THRESHOLD = 0.6
 
-async def arequest_llm(request_body: str, stream: bool = True, vllm_url: str = get_settings().VLLM_URL, user_id: str = None, use_vector_db: bool = False) -> Union[httpx.Response, JSONResponse, LLMSuccessResponse]:
+async def arequest_llm(
+    request_body: str,
+    stream: bool = True,
+    vllm_url: str = get_settings().VLLM_URL,
+    user_id: str | None = None,
+    is_api_key: bool = False,
+    use_vector_db: bool = False
+) -> Union[httpx.Response, JSONResponse, LLMSuccessResponse]:
     """
     Request LLM.
     - Returns httpx.Response if stream=True and status=200 (for caller to handle streaming).
@@ -25,7 +32,7 @@ async def arequest_llm(request_body: str, stream: bool = True, vllm_url: str = g
     response: Optional[httpx.Response] = None
 
     # Add system prompt to the request body
-    request_body = await _add_system_prompt(request_body)
+    request_body = await _add_system_prompt(request_body, is_api_key)
 
     # Apply vector DB if enabled
     if use_vector_db:
@@ -80,7 +87,14 @@ async def arequest_llm(request_body: str, stream: bool = True, vllm_url: str = g
             await client.aclose()
 
 
-def request_llm(request_body: str, stream: bool = True, vllm_url: str = get_settings().VLLM_URL) -> Union[httpx.Response, JSONResponse, LLMSuccessResponse]:
+def request_llm(
+    request_body: str,
+    stream: bool = True,
+    vllm_url: str = get_settings().VLLM_URL,
+    user_id: str | None = None,
+    is_api_key: bool = False,
+    use_vector_db: bool = False
+) -> Union[httpx.Response, JSONResponse, LLMSuccessResponse]:
     """
     Request LLM (Synchronous version).
     - Returns httpx.Response if stream=True and status=200 (for caller to handle streaming).
@@ -91,7 +105,7 @@ def request_llm(request_body: str, stream: bool = True, vllm_url: str = get_sett
     response: Optional[httpx.Response] = None
 
     # Add system prompt to the request body
-    request_body = asyncio.run(_add_system_prompt(request_body))
+    request_body = asyncio.run(_add_system_prompt(request_body, is_api_key))
 
     try:
         headers = { "Content-Type": "application/json" }
@@ -141,10 +155,10 @@ def request_llm(request_body: str, stream: bool = True, vllm_url: str = get_sett
         if client and not client.is_closed and not is_streaming_success:
             client.close()
 
-async def _add_system_prompt(request_body: str) -> str:
+async def _add_system_prompt(request_body: str, is_api_key: bool) -> str:
     """Add a system prompt to the messages."""
     request_body = json.loads(request_body)
-    default_prompt = await get_system_prompt(request_body["model"], "default")
+    default_prompt = await get_system_prompt(request_body["model"], "default", is_api_key)
     if default_prompt:
         request_body["messages"] = [{"role": "system", "content": default_prompt.format(current_date=datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}] + request_body["messages"]
     return json.dumps(request_body)