Skip to content

grpc.md

Latest commit

 

History

History
277 lines (213 loc) · 11.6 KB

grpc.md

File metadata and controls

277 lines (213 loc) · 11.6 KB
Raw

rag-gateway gRPC reference (Step 3.2)

Overview

The gateway exposes the same Phase 2 pipeline composition over a gRPC contract defined in proto/rag.proto. The service mirrors the REST surface from Step 3.1 method-for-method so clients can move between protocols without re-shaping requests.

Wire package: rag.gateway.v1. Service: RagService.

RPC Streaming REST equivalent
Query server-streaming POST /v1/query
Retrieve unary POST /v1/retrieve
IngestDocument unary POST /v1/ingest/document
ListCorpora unary GET /v1/corpora
GetCorpus unary GET /v1/corpora/{id}
Converse server-streaming POST /v1/agent

The gRPC server also exposes the standard health + reflection services:

  • grpc.health.v1.Health — load-balancer + orchestrator probe.
  • grpc.reflection.v1alpha.ServerReflection — enables grpcurl and similar tools without distributing .proto files.

Usage

Python client

import grpc
from rag_gateway._grpc_gen import rag_pb2, rag_pb2_grpc

async with grpc.aio.insecure_channel("127.0.0.1:50051") as channel:
    stub = rag_pb2_grpc.RagServiceStub(channel)
    req = rag_pb2.QueryRequest(
        tenant_id="acme",
        principal_id="alice",
        query="how does retrieval work?",
        top_k=5,
        rerank=True,
        pack=True,
        generate=False,
    )
    async for event in stub.Query(req):
        kind = event.WhichOneof("event")
        if kind == "stage_started":
            print(f"→ {event.stage_started.stage}")
        elif kind == "answer_delta":
            print(event.answer_delta.text, end="")
        elif kind == "result":
            print(f"\nrequest_id={event.result.request_id}")

grpcurl

Reflection is enabled by default so grpcurl works without -proto flags:

# Health check
grpcurl -plaintext 127.0.0.1:50051 grpc.health.v1.Health/Check

# List corpora
grpcurl -plaintext \
  -H "x-tenant-id: acme" \
  -H "x-principal-id: alice" \
  127.0.0.1:50051 rag.gateway.v1.RagService/ListCorpora

# Server-streaming query
grpcurl -plaintext -d '{"tenant_id":"acme","principal_id":"alice","query":"hi","top_k":3}' \
  127.0.0.1:50051 rag.gateway.v1.RagService/Query

ragctl smoke command

uv run ragctl grpc-query "what is RAG?" --show-events

Drives an in-process server on an ephemeral loopback port — no infrastructure required.

RPC catalogue

Query (QueryRequest) → stream QueryStreamEvent

Server-streams pipeline events in execution order and terminates with a QueryStreamEvent.result carrying the full QueryResponse.

Event kinds (one-of on QueryStreamEvent.event):

Kind When emitted Payload
stage_started Just before each pipeline stage runs. StageStarted — stage enum + timestamp
stage_completed Right after each stage finishes. StageCompleted — stage + elapsed_ms + timestamp
answer_delta One or more times when generate=true. AnswerDelta — text fragment + optional tokens
result Always last; carries the final envelope. QueryResponse

v0 backends emit a single AnswerDelta carrying the full answer text; clients SHOULD concatenate text fields in arrival order. Token-by- token streaming is forward-compatible with this shape.

Request fields mirror rag_core.gateway_types.QueryRequest 1:1. Defaults match the REST handler: top_k=10, rerank=true, pack=true, pack_budget=2048, generate=false.

Retrieve (RetrieveRequest) → RetrieveResponse

Retrieval-only — runs understanding + router, returns ChunkRefs without rerank / pack / LLM. Use this when the caller brings its own reranker.

IngestDocument (IngestDocumentRequest) → IngestResult

Single-document unary ingest. Carries the raw bytes inline; for multi-MB payloads use a chunked upload (deferred to a later step).

ListCorpora (ListCorporaRequest) → CorpusList

Returns every corpus visible to the calling tenant. Identity is read from inbound metadata: x-tenant-id (required) and x-principal-id (optional, defaults to gateway-anon).

GetCorpus (GetCorpusRequest) → Corpus

Returns a single corpus or NOT_FOUND. Cross-tenant lookups also return NOT_FOUND deliberately, so attackers can't probe for the existence of other tenants' corpora.

Converse (ConverseRequest) → stream ConverseStreamEvent

The gRPC twin of POST /v1/agent (Step 3.6): drives the rag-agent AgentLoop over a goal, server-streaming one ConverseStreamEvent per state-machine transition and ending the stream with a terminal run_completed (or run_failed) frame.

ConverseStreamEvent is a flat message — kind (AgentEventKind) discriminates and the optional fields carry whatever that kind needs — mirroring rag_core.agent_types.AgentEvent field-for-field (enforced by tests/contract/test_grpc_proto_compat.py):

Kind When emitted Payload fields
AGENT_EVENT_KIND_RUN_STARTED First frame. run_id
AGENT_EVENT_KIND_PHASE_CHANGED On each planning/acting/observing/finalizing transition. phase, step_index
AGENT_EVENT_KIND_TOOL_STARTED Before a tool is invoked. tool_call (AgentToolCall, args as Struct)
AGENT_EVENT_KIND_TOOL_COMPLETED After a tool returns. tool_result (AgentToolResult + chunk_refs)
AGENT_EVENT_KIND_CHECKPOINT_SAVED After a step is checkpointed. step_index
AGENT_EVENT_KIND_ANSWER_DELTA Once the final answer is ready (post-egress). text
AGENT_EVENT_KIND_RUN_COMPLETED Always last on success. result (AgentRunResult)
AGENT_EVENT_KIND_RUN_FAILED Terminal on failure (e.g. resume of an unknown run). error_code, message

ConverseRequest mirrors AgentRequest 1:1 — goal, corpus_ids, top_k (default 8), max_steps (default 6), the four optional budgets (budget_tokens / budget_dollars / budget_wall_ms / budget_iter; unset = uncapped), request_id, and resume_run_id. Identity comes from the request body (tenant_id / principal_id), like the other body-driven RPCs.

Egress governance over the final answer runs at the boundary exactly as on REST: PolicyEngine.egress_text is consulted once before the answer_delta frame, so a denied answer never reaches the wire; the decision is replayed onto the terminal result. A resume of an unknown run_id is surfaced in-band as a run_failed frame, not a gRPC status abort.

Metadata contract

The server reads three header keys from inbound metadata, all lowercase per gRPC convention:

Key Purpose
traceparent W3C Trace Context — parsed into the per-request TraceContext. Optional; a fresh trace starts when absent or malformed.
x-request-id Caller-supplied request id. Generated as UUID4 when absent.
x-tenant-id Tenant for the header-driven RPCs (ListCorpora / GetCorpus). Body-driven RPCs (Query / Retrieve / IngestDocument) read the tenant from the request message.
x-principal-id Principal for header-driven RPCs. Defaults to gateway-anon.
authorization Reserved for production auth wiring (Phase 6 governance). Not consulted by the default NoopAuth.

Error envelope

Every non-OK response sets:

  1. The standard grpc.StatusCode (see table below).
  2. A short human-readable string on status.details().
  3. A binary-packed GatewayError on trailing metadata under rag-gateway-error-bin so clients can recover the full structured context.
err = rag_pb2.GatewayError()
err.ParseFromString(rpc_error.trailing_metadata()["rag-gateway-error-bin"])
print(err.code, err.message, dict(err.context.items()))

Status code mapping

The mapping is intentionally narrow and mirrors the REST handler at apps/gateway/src/rag_gateway/query.py:

gRPC code Domain trigger REST equivalent
PERMISSION_DENIED ACLDeniedError 403
UNAUTHENTICATED AuthError / missing identity metadata 401
RESOURCE_EXHAUSTED RateLimitError 429
UNAVAILABLE RetrievalError (downstream backend down) 502
INVALID_ARGUMENT ValueError / malformed request 400
NOT_FOUND GetCorpus returns None 404
INTERNAL fall-through unexpected exceptions 500

Observability

Canonical OTel spans wrap the streaming + unary flows:

  • grpc.query — full RagService.Query flow.
  • grpc.retrieve — Retrieve flow.
  • grpc.converse — Converse flow (wraps the loop's agent.run span).

Attributes mirror the REST spans (rag.tenant_id, rag.gateway.results_n, rag.gateway.elapsed_ms, …) so dashboards can union the two surfaces.

Structured-log events emitted via rag-observability:

  • grpc.query_complete — one per successful Query.
  • grpc.ingest_complete — one per successful IngestDocument.
  • grpc.converse_complete — one per successful Converse run.
  • grpc.answer.failed — LLM call failed under generate=true (non- fatal; chunks still returned).
  • grpc.answer.policy_denied / gateway.agent.policy_denied — PolicyEngine egress_text deny (query / agent answer respectively).

Codegen

Python stubs live at apps/gateway/src/rag_gateway/_grpc_gen/ and are committed for the proto-drift gate. Regenerate via task proto:gen (python scripts/gen_proto_py.py). The script uses grpcio-tools + mypy-protobuf — pure-Python install, no system protoc needed.

buf is only required for lint + breaking-change checks; CI installs it via bufbuild/buf-setup-action. Local devs can brew install bufbuild/buf/buf or winget install Buf.Buf if they want to run task proto:lint before pushing.

Internals

The implementation lives in two files:

  • grpc_service.py — GrpcRagService(RagServiceServicer) with the per-method handlers.
  • grpc_server.py — build_grpc_server(...) factory + serve(...) convenience wrapper.

build_grpc_server accepts the same kwargs as rag_gateway.app.build_app, so a single set of SPI implementations serves both REST + gRPC in one process. Defaults pull from the shared build_default_* helpers in rag_gateway.app.

Extension points

  • Auth: swap NoopAuth for a production impl via the auth kwarg to build_grpc_server. The interceptor honours the authorization: Bearer … metadata header for header-driven RPCs.
  • PolicyEngine: passed through to the same egress_text consult used by REST in _generate_answer — see architecture/policy-engine.md.
  • Per-call deadlines / cancellation: handled natively by grpc.aio — server handlers cancel cleanly when the client closes the stream.
  • Production deployment: bind a TLS port via server.add_secure_port(bind, server_credentials) instead of the insecure helper. The factory is unopinionated about transport security.