Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 33 additions & 14 deletions AIGateway/src/routers/prompts.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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")
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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>
Expand All @@ -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
Expand Down
14 changes: 13 additions & 1 deletion Clients/src/presentation/pages/AIGateway/shared.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -274,6 +274,7 @@ export async function streamPromptTest(
let content = "";
let tokens = 0;
let cost = 0;
let streamError = "";

try {
while (true) {
Expand All @@ -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;
Expand All @@ -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 };
}
Loading