def register_cg_rag_routes(app: FastAPI, service: Any, *, prefix: str = "/cg-rag") -> None:
install_request_id_middleware(app)
route_prefix = _clean_prefix(prefix)
@app.get(f"{route_prefix}/health")
def cg_rag_health() -> Any:
try:
return _with_request_id(service.health())
except Exception as exc:
return _service_error_response(exc)
@app.get(f"{route_prefix}/profiles")
def cg_rag_profiles() -> Any:
try:
return service.profiles()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/warmup")
async def cg_rag_warmup(request: Request) -> Any:
parsed, error = await _parse_post_request(CgWarmupRequest, request)
if error is not None:
return error
try:
return _with_request_id(await _run_service_with_cancellation(service.warmup, parsed.scopes))
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/retrieve-rerank")
async def cg_rag_retrieve_rerank(request: Request) -> Any:
parsed, error = await _parse_post_request(CgRetrieveRerankRequest, request)
if error is not None:
return error
try:
return _with_request_id(await _run_service_with_cancellation(service.retrieve_rerank, parsed))
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/constrained-generate")
async def cg_rag_constrained_generate(request: Request) -> Any:
parsed, error = await _parse_post_request(CgConstrainedGenerateRequest, request)
if error is not None:
return error
try:
return _with_request_id(await _run_service_with_cancellation(service.constrained_generate, parsed))
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/rag")
async def cg_rag(request: Request) -> Any:
parsed, error = await _parse_post_request(CgRagRequest, request)
if error is not None:
return error
try:
return _with_request_id(await _run_service_with_cancellation(service.rag, parsed))
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/rag/stream")
async def cg_rag_stream(request: Request) -> Any:
parsed, error = await _parse_post_request(CgRagRequest, request)
if error is not None:
return error
try:
cancel_token = GenerationToken(current_request_id())
stream_events_kwargs: dict[str, Any] = {}
if _supports_keyword(service.rag_stream_events, "cancel_token"):
stream_events_kwargs["cancel_token"] = cancel_token
event_iterable = _sse_events(
service.rag_stream_events(parsed, **stream_events_kwargs),
request=parsed,
cancel_token=cancel_token,
)
async_event_iterable = CloseableAsyncIterator(
event_iterable,
cleanup=lambda: cancel_token.cancel(),
)
return CleanupStreamingResponse(
async_event_iterable,
cleanup=async_event_iterable.aclose,
media_type="text/event-stream",
)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
def _vision_stream_response(
stream_events: Callable[..., Iterable[tuple[str, dict[str, Any]]]],
parsed: Any,
) -> Any:
cancel_token = GenerationToken(current_request_id())
stream_events_kwargs: dict[str, Any] = {}
if _supports_keyword(stream_events, "cancel_token"):
stream_events_kwargs["cancel_token"] = cancel_token
event_iterable = _sse_events(
stream_events(parsed, **stream_events_kwargs),
cancel_token=cancel_token,
)
async_event_iterable = CloseableAsyncIterator(
event_iterable,
cleanup=lambda: cancel_token.cancel(),
)
return CleanupStreamingResponse(
async_event_iterable,
cleanup=async_event_iterable.aclose,
media_type="text/event-stream",
)
@app.post(f"{route_prefix}/vision/answer")
async def cg_rag_vision_answer(request: Request) -> Any:
parsed, error = await _parse_post_request(CgVisionAnswerRequest, request)
if error is not None:
return error
try:
return _with_request_id(
await _run_service_with_cancellation(service.vision_answer, parsed)
)
except CgVisionError as exc:
return _vision_error_response(exc)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/vision/answer/stream")
async def cg_rag_vision_answer_stream(request: Request) -> Any:
parsed, error = await _parse_post_request(CgVisionAnswerRequest, request)
if error is not None:
return error
try:
return _vision_stream_response(service.vision_answer_stream_events, parsed)
except CgVisionError as exc:
return _vision_error_response(exc)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/vision/observe")
async def cg_rag_vision_observe(request: Request) -> Any:
parsed, error = await _parse_post_request(CgVisionAnswerRequest, request)
if error is not None:
return error
try:
return _with_request_id(
await _run_service_with_cancellation(service.vision_observe, parsed)
)
except CgVisionError as exc:
return _vision_error_response(exc)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/vision/agentic")
async def cg_rag_vision_agentic(request: Request) -> Any:
parsed, error = await _parse_post_request(CgVisionAgenticRequest, request)
if error is not None:
return error
try:
return _with_request_id(
await _run_service_with_cancellation(service.vision_agentic, parsed)
)
except CgVisionError as exc:
return _vision_error_response(exc)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/vision/agentic/stream")
async def cg_rag_vision_agentic_stream(request: Request) -> Any:
parsed, error = await _parse_post_request(CgVisionAgenticRequest, request)
if error is not None:
return error
try:
return _vision_stream_response(service.vision_agentic_stream_events, parsed)
except CgVisionError as exc:
return _vision_error_response(exc)
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)
@app.post(f"{route_prefix}/answer")
async def cg_rag_answer(request: Request) -> Any:
parsed, error = await _parse_post_request(CgRagRequest, request)
if error is not None:
return error
try:
return _with_request_id(await _run_service_with_cancellation(service.rag, parsed))
except CapacityExceededError as exc:
return _capacity_error_response(exc)
except GenerationCancelledError:
return _cancelled_response()
except Exception as exc:
return _service_error_response(exc)