stum-ai-gateway 0.1.11

NatBus-to-LLM agent gateway (contract-validated, iterative planning)

AI Gateway

Interface-only bridge between NatBus and an external LLM. Maps human commands to service request/response subjects, asks an LLM to produce a JSON “plan”, publishes requests, and routes responses back to the human reply subject. No LLM SDKs bundled; you inject an async callback.

Features • LLM injection via async callback (prompt: str, system: Optional[str]) -> Awaitable[str] • Command registry for easy, extensible mapping of human commands → NatBus subjects • Correlation propagation using inbound x-correlation-id (or auto-generated UUID) • Optional LLM post-processing of service responses per-command or globally • Gzip compression for outbound messages with transparent inbound decompression • Supports PUSH consumers for human commands and service responses

Install Requirements

python3.11 -m pip install -r requirements.txt -v

Build

./build.sh

Concepts

Human command envelope (JSON on human_command_subject):

{"cmd":"<string>","args":{"...": "..."},"reply_subject":"<subject optional>"}

LLM plan schema (LLM must output a single JSON object):

{
"action": "send_request",
"subject": "<service.request.subject>",
"payload": { "service": "specific", "fields": "..." },
"await_response": true,
"response_subject": "<service.response.subject>" // optional; overrides mapping default
}

Service response envelope (typical):

{ "ok": true, "data": { "...": "..." }, "meta": { "..." : "..." } }

Human reply envelope (published to reply_subject):

{ "correlation_id": "<id>", "command": "<cmd>", "data": { "...": "..." } }

Quick Start

1. Provide an LLM callback

from typing import Optional

async def llm_call(prompt: str, system: Optional[str]) -> str:
# Must return a single JSON object string following the LLM plan schema
# Example: route "show active trades" to forex service
return (
'{"action":"send_request","subject":"forex.trades.list.req","payload":{},"await_response":true,'
'"response_subject":"forex.trades.list.resp"}'
)

2. Register command mappings

from ai_gateway import CommandRegistry, CommandMapping

registry = CommandRegistry()

registry.register(CommandMapping(
command="show active trades",
request_subject="forex.trades.list.req",
response_subject="forex.trades.list.resp",
llm_instructions="Use an empty payload; await_response true.",
llm_postprocess=True, # optional: run LLM to summarize service response
))

registry.register(CommandMapping(
command="get account info",
request_subject="forex.account.info.req",
response_subject="forex.account.info.resp",
))

3. Configure and run the agent

import asyncio
from typing import Optional
from natbus.config import NatsConfig
from natbus.client import NatsBus
from ai_gateway import LlmNatbusAgent, LlmAgentConfig

CFG = NatsConfig(
server="nats-nats-jetstream:4222",
username="nats-user",
password="changeme",
name="ai-gateway",
stream_create=True,
stream_name="AI_STREAM",
stream_subjects=("ai.human.commands","ai.human.replies","forex.trades.list.req","forex.trades.list.resp"),
queue_group="ai-gateway",
)

async def main():
bus = NatsBus(CFG)
await bus.connect()
agent = LlmNatbusAgent(
bus=bus,
llm_call=llm_call,
registry=registry,
cfg=LlmAgentConfig(
human_command_subject="ai.human.commands",
default_reply_subject="ai.human.replies",
compress_outbound=True,
pending_timeout_seconds=180,
),
)
await agent.start()
# keep running
while True:
await asyncio.sleep(60)

if name == "main":
asyncio.run(main())

4. Publish a human command (e.g., from UI/controller)

await bus.publish_json(
subject="ai.human.commands",
obj={"cmd": "show active trades", "args": {}, "reply_subject": "ai.human.replies"},
sender="ui",
)

Command Mapping Structure

CommandMapping fields:

• command: canonical human command string (lowercased for lookup) • request_subject: NatBus subject to publish the service request • response_subject: subject on which the service posts responses (optional) • llm_instructions: extra prompt hints for command-specific nuances (optional) • llm_postprocess: run LLM on service response before replying (optional)

CommandRegistry responsibilities:

• register(mapping): adds/overwrites a mapping keyed by command • get(command): returns the CommandMapping for a human command • all_response_subjects(): returns the set of response subjects for auto-subscription

Extending mappings:

Add a new command → subject pair
registry.register(CommandMapping(
command="get forex quote",
request_subject="forex.quote.req",
response_subject="forex.quote.resp",
llm_instructions="Payload must include symbol (e.g. EUR/USD). await_response true.",
))

LLM Integration Contract

The agent sends a prompt that includes:

• System prompt (schema and output constraints) • Context with command, args, default request_subject, and default response_subject • Any llm_instructions from the mapping

Your callback must:

• Return a single JSON document (no markdown, no commentary) • Include required keys: action, subject, payload, await_response • Optionally include response_subject (overrides mapping)

Retries:

• The agent retries invalid outputs up to llm_max_retries with a short JSON-only reminder • After retries, the agent publishes an error to the requester’s reply subject

Routing and Correlation • Inbound x-correlation-id is reused for all downstream messages and final replies • If missing, the agent generates a UUID and uses it consistently in request headers and the reply body and headers • The agent subscribes to all configured response subjects and matches responses by correlation ID

Compression Outbound (agent → bus):

• Controlled by LlmAgentConfig.compress_outbound • When enabled, JSON payloads are gzip-compressed and marked with content-encoding: gzip

Inbound (bus → agent):

• ReceivedMessage auto-decompresses if content-encoding: gzip • Your handlers and tests can read .as_json() or .as_text() regardless of compression

Subscriptions • Human commands: PUSH consumer on human_command_subject with a durable • Service responses: PUSH consumers per response subject derived from the registry and extra_response_subjects

Durables and queue groups:

• Use a durable name to preserve delivery cursor and acks across restarts • Use a queue group to load-balance multiple agent replicas

Error Paths

Unknown command:

{ "error": "unknown_command", "cmd": "<user text>" }

Invalid LLM output after retries:

{ "error": "invalid_llm_output", "detail": "<reason>" }

No response subject configured while await_response is true:

{ "error": "no_response_subject_configured" }

Unsupported plan action:

{ "error": "unsupported_action", "action": "<value>" }

Configuration Reference

LlmAgentConfig most relevant fields:

• human_command_subject: subject to receive human commands • default_reply_subject: fallback reply subject if requester did not specify one • compress_outbound: gzip-compress outbound messages when true • llm_max_retries: retry count for invalid LLM outputs • llm_system_prompt: schema and constraints for planning calls • llm_postprocess_system_prompt: schema for post-processing {result: ...} • pending_timeout_seconds: TTL for awaiting responses • extra_response_subjects: additional subjects to subscribe to

Example: Multi-service Registry

registry = CommandRegistry()

Forex
registry.register(CommandMapping(
command="show active trades",
request_subject="forex.trades.list.req",
response_subject="forex.trades.list.resp",
llm_instructions="Empty payload; await_response true.",
llm_postprocess=True,
))
registry.register(CommandMapping(
command="get forex quote",
request_subject="forex.quote.req",
response_subject="forex.quote.resp",
llm_instructions="Require args.symbol like 'EUR/USD'.",
))

Accounts
registry.register(CommandMapping(
command="get account info",
request_subject="acct.info.req",
response_subject="acct.info.resp",
))

Orders
registry.register(CommandMapping(
command="place order",
request_subject="orders.place.req",
response_subject="orders.place.resp",
llm_instructions="Payload must include side, symbol, qty, type.",
))

Testing

Unit tests use a FakeBus and stubbed llm_call. Tests cover compressed/uncompressed flows, correlation propagation, retries, unknown command, and non-JSON passthrough.

Run tests:

pytest -q

Key fixtures in tests/conftest.py:

• bus – fake NatBus • make_cfg – builds LlmAgentConfig with overrides • make_agent – starts LlmNatbusAgent with provided registry and callback • decode_bus_json – decompresses and decodes BusMessage bodies for assertions

Build and Distribute

Version is controlled in pyproject.toml (project.version). Build artifacts:

python -m pip install --upgrade build
python -m build
ls dist/

Notes for Service Authors

Subject naming:

• Requests: ...req • Responses: ...resp

Correlation:

• Echo inbound x-correlation-id on all responses • Respond on the subject specified by the plan or the documented default

Payloads:

• JSON only for requests and responses to maximize compatibility with the LLM plan schema • For large payloads, gzip; the gateway handles decompression automatically

Minimal API Surface (import points)

from ai_gateway import (
LlmAgentConfig,
CommandRegistry,
CommandMapping,
LlmNatbusAgent, # requires injected llm_call
)

This is sufficient to register commands, run the agent, and integrate your external LLM.