Merge pull request #3650 from verifywise-ai/fix/ai-gateway-prompt-test-streaming

fix(ai-gateway): implement prompt test streaming and fix Compare panel

Author: gorkem-bwl

AIGateway/src/routers/prompts.py

@@ -1,11 +1,15 @@
+import json as _json
 import re
 from typing import Any, Optional
 
 from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, field_validator
 
 from crud import prompts as crud
 from middlewares.auth import verify_internal_key
+from services.llm_service import stream_chat_completion
+from services.proxy_service import resolve_endpoint_for_key
 from utils.auth import get_org_id, get_user_id, require_admin
 from utils.notifications import notify_config_change
 
@@ -377,27 +381,42 @@ async def delete_test_dataset(
     return {"deleted": True}
 
 
-@router.post("/test", status_code=501)
+@router.post("/test")
 async def test_prompt(request: Request, body: TestPromptRequest):
     """
-    Test a prompt by resolving variables and proxying to the LLM completion
-    endpoint. The LLM proxy integration lives in the Express backend for now;
-    this stub returns 501 until the proxy is co-located in the AIGateway
-    service.
+    Test a prompt by resolving variables and streaming the LLM response.
+    Returns an SSE stream compatible with the frontend's streamPromptTest().
     """
     verify_internal_key(request)
+    org_id = get_org_id(request)
 
-    # Resolve variables so callers can at least validate substitution locally
+    # Resolve variables in the prompt content
     resolved_content = crud.resolve_variables(
         body.content,
         body.variables or {},
     )
 
-    raise HTTPException(
-        status_code=501,
-        detail=(
-            "test-prompt requires proxy integration — "
-            "LLM proxy currently runs in the Express backend. "
-            "Resolved content is available but cannot be forwarded yet."
-        ),
-    )
+    # Resolve the endpoint to get provider, model, and API key
+    try:
+        endpoint = await resolve_endpoint_for_key(
+            organization_id=org_id,
+            endpoint_slug=body.endpoint_slug,
+            allowed_endpoint_ids=[],
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+    # Stream the LLM response
+    async def _stream():
+        try:
+            async for chunk_str in stream_chat_completion(
+                model=endpoint["model"],
+                messages=resolved_content,
+                api_key=endpoint["decrypted_key"],
+            ):
+                yield chunk_str
+        except Exception as e:
+            yield f"data: {_json.dumps({'error': str(e)})}\n\n"
+            yield "data: [DONE]\n\n"
+
+    return StreamingResponse(_stream(), media_type="text/event-stream")

Clients/src/presentation/pages/AIGateway/Prompts/ComparePanel.tsx

@@ -116,6 +116,7 @@ export default function ComparePanel({
             items={versionItems}
             placeholder="Select version"
             sx={{ flex: 1 }}
+            getOptionValue={(item) => item._id}
           />
           <Select
             id="compare-select-right"
@@ -125,6 +126,7 @@ export default function ComparePanel({
             items={versionItems}
             placeholder="Select version"
             sx={{ flex: 1 }}
+            getOptionValue={(item) => item._id}
           />
         </Box>
         <Typography fontSize={13} fontWeight={500} color="text.secondary" mt="8px">Endpoint</Typography>
@@ -135,6 +137,7 @@ export default function ComparePanel({
           items={endpoints.map((e) => ({ _id: e.slug, name: `${e.display_name} (${e.slug})` }))}
           placeholder="Select endpoint"
           sx={{ width: "100%" }}
+          getOptionValue={(item) => item._id}
         />
         <Typography fontSize={11} color="text.disabled" mt="4px">
           {detectedVars.length > 0

Clients/src/presentation/pages/AIGateway/shared.ts

@@ -245,7 +245,7 @@ export async function streamPromptTest(
 ): Promise<StreamPromptTestResult> {
   const startTime = Date.now();
 
-  const response = await fetch(`${GATEWAY_API_URL}/ai-gateway/prompts/test`, {
+  const response = await fetch(`/api/ai-gateway/prompts/test`, {
     method: "POST",
     headers: {
       "Content-Type": "application/json",
@@ -274,6 +274,7 @@ export async function streamPromptTest(
   let content = "";
   let tokens = 0;
   let cost = 0;
+  let streamError = "";
 
   try {
     while (true) {
@@ -285,6 +286,9 @@ export async function streamPromptTest(
         if (line.startsWith("data: ") && line !== "data: [DONE]") {
           try {
             const chunk = JSON.parse(line.slice(6));
+            if (chunk.error) {
+              streamError = chunk.error;
+            }
             const delta = chunk.choices?.[0]?.delta?.content;
             if (delta) {
               content += delta;
@@ -300,5 +304,13 @@ export async function streamPromptTest(
     reader.releaseLock();
   }
 
+  if (!content && streamError) {
+    content = `Error: ${streamError}`;
+    opts.onDelta(content);
+  } else if (!content) {
+    content = "No response received from the model. Check that the endpoint has a valid API key configured.";
+    opts.onDelta(content);
+  }
+
   return { content, tokens, cost, latency: Date.now() - startTime };
 }