diff --git a/agentic/agentic-server/README.md b/agentic/agentic-server/README.md index 0d5cdebf76..95eea42af1 100644 --- a/agentic/agentic-server/README.md +++ b/agentic/agentic-server/README.md @@ -46,6 +46,10 @@ app.listen(3001, () => { | GET | `/v1/providers` | List configured providers | | GET | `/healthz` | Health check | +## Streaming + +`POST /v1/chat/completions` with `stream: true` is relayed as a `text/event-stream`: the gateway forwards the provider's frames verbatim, so a client receives tokens as they are produced. It adds `stream_options: { include_usage: true }` unless the caller set `stream_options` itself, and meters the usage the final frame carries. Streaming is available for OpenAI-compatible providers; a provider whose stream would need translation (`ollama`, `anthropic`) is rejected `501` rather than answered with a non-streaming body. + ## Identity & tenancy Requests carry tenant identity via headers: `X-Database-Id` (**required** — requests without it are rejected `400`), `X-Entity-Id`, and `X-Actor-Id`. Routing to a specific provider can be forced with `X-LLM-Provider`. diff --git a/agentic/agentic-server/__tests__/streaming.test.ts b/agentic/agentic-server/__tests__/streaming.test.ts new file mode 100644 index 0000000000..2531723c9f --- /dev/null +++ b/agentic/agentic-server/__tests__/streaming.test.ts @@ -0,0 +1,187 @@ +/** + * Streaming tests for agentic-server. + * + * A harness that asks for `stream: true` needs its tokens as the provider emits + * them, so the gateway relays an OpenAI-compatible event stream verbatim rather + * than parsing the body as JSON. These tests assert the relayed frames, the + * usage the stream carries reaching the injected sink, the frame-splitting the + * scanner does across arbitrary chunk boundaries, and the loud rejection of a + * provider type whose stream the gateway cannot translate. + */ + +import express from 'express'; +import type { Server } from 'http'; + +import type { InferenceEntry } from '../src'; +import { createAgenticServer } from '../src'; +import { UsageScanner, withUsageStreamOptions } from '../src/streaming'; + +const CHUNKS = [ + 'data: {"id":"c1","choices":[{"delta":{"content":"Hel"}}]}\n\n', + 'data: {"id":"c1","choices":[{"delta":{"content":"lo"}}]}\n\n', + 'data: {"id":"c1","choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}\n\n', + 'data: [DONE]\n\n' +]; + +let mockLlmServer: Server; +let mockLlmPort: number; +let agenticServer: Server; +let agenticPort: number; +let ollamaServer: Server; +let ollamaPort: number; + +let llmRequests: Array<{ body: any }>; +let sinkEntries: InferenceEntry[]; + +beforeAll(async () => { + const llmApp = express(); + llmApp.use(express.json()); + + llmApp.post('/v1/chat/completions', (req: any, res: any) => { + llmRequests.push({ body: req.body }); + res.setHeader('Content-Type', 'text/event-stream'); + for (const chunk of CHUNKS) res.write(chunk); + res.end(); + }); + + await new Promise((resolve) => { + mockLlmServer = llmApp.listen(0, () => { + const addr = mockLlmServer.address(); + mockLlmPort = typeof addr === 'object' && addr ? addr.port : 0; + resolve(); + }); + }); + + const app = createAgenticServer({ + providerType: 'openai', + providerBaseUrl: `http://localhost:${mockLlmPort}`, + providerApiKey: 'test-key', + defaultModel: 'gpt-4o-mini', + inferenceSink: { logInference: (entry) => sinkEntries.push(entry) } + }); + + await new Promise((resolve) => { + agenticServer = app.listen(0, () => { + const addr = agenticServer.address(); + agenticPort = typeof addr === 'object' && addr ? addr.port : 0; + resolve(); + }); + }); + + const ollamaApp = createAgenticServer({ + providerType: 'ollama', + providerBaseUrl: `http://localhost:${mockLlmPort}`, + defaultModel: 'llama3' + }); + + await new Promise((resolve) => { + ollamaServer = ollamaApp.listen(0, () => { + const addr = ollamaServer.address(); + ollamaPort = typeof addr === 'object' && addr ? addr.port : 0; + resolve(); + }); + }); +}); + +afterAll(async () => { + await new Promise((r, e) => agenticServer.close((err) => (err ? e(err) : r()))); + await new Promise((r, e) => ollamaServer.close((err) => (err ? e(err) : r()))); + await new Promise((r, e) => mockLlmServer.close((err) => (err ? e(err) : r()))); +}); + +beforeEach(() => { + llmRequests = []; + sinkEntries = []; +}); + +// ─── SSE relay ──────────────────────────────────────────────────────────── + +describe('POST /v1/chat/completions with stream: true', () => { + it('relays the provider event stream and meters the usage it carried', async () => { + const res = await fetch(`http://localhost:${agenticPort}/v1/chat/completions`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Database-Id': 'db-stream-1', + 'X-Entity-Id': 'entity-stream-1' + }, + body: JSON.stringify({ + model: 'gpt-4o-mini', + stream: true, + messages: [{ role: 'user', content: 'Hello' }] + }) + }); + + expect(res.status).toBe(200); + expect(res.headers.get('content-type')).toContain('text/event-stream'); + expect(await res.text()).toBe(CHUNKS.join('')); + + expect(llmRequests).toHaveLength(1); + expect(llmRequests[0].body.stream).toBe(true); + expect(llmRequests[0].body.stream_options).toEqual({ include_usage: true }); + + await new Promise((r) => setTimeout(r, 100)); + expect(sinkEntries).toHaveLength(1); + expect(sinkEntries[0]).toMatchObject({ + databaseId: 'db-stream-1', + entityId: 'entity-stream-1', + service: 'chat', + status: 'ok', + inputTokens: 11, + outputTokens: 5, + totalTokens: 16 + }); + }); + + it('rejects (501) a provider type whose stream the gateway cannot translate', async () => { + const res = await fetch(`http://localhost:${ollamaPort}/v1/chat/completions`, { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Database-Id': 'db-stream-2' }, + body: JSON.stringify({ model: 'llama3', stream: true, messages: [] }) + }); + + expect(res.status).toBe(501); + expect(llmRequests).toHaveLength(0); + }); +}); + +// ─── Usage scanning ─────────────────────────────────────────────────────── + +describe('UsageScanner', () => { + it('finds usage split across arbitrary chunk boundaries', () => { + const scanner = new UsageScanner(); + const stream = CHUNKS.join(''); + for (let i = 0; i < stream.length; i += 7) scanner.push(stream.slice(i, i + 7)); + + expect(scanner.result()).toEqual({ prompt_tokens: 11, completion_tokens: 5, total_tokens: 16 }); + }); + + it('reports no usage for a stream that carried none', () => { + const scanner = new UsageScanner(); + scanner.push('data: {"choices":[{"delta":{"content":"hi"}}]}\n\ndata: [DONE]\n\n'); + + expect(scanner.result()).toBeUndefined(); + }); + + it('ignores a frame it cannot parse and keeps scanning', () => { + const scanner = new UsageScanner(); + scanner.push(': keep-alive\n\ndata: not-json\n\n'); + scanner.push('data: {"usage":{"prompt_tokens":3,"completion_tokens":4}}\n\n'); + + expect(scanner.result()).toEqual({ prompt_tokens: 3, completion_tokens: 4, total_tokens: 7 }); + }); +}); + +describe('withUsageStreamOptions', () => { + it('asks the provider for usage when the caller expressed no preference', () => { + expect(withUsageStreamOptions({ model: 'gpt-4o-mini' })).toEqual({ + model: 'gpt-4o-mini', + stream_options: { include_usage: true } + }); + }); + + it('keeps the stream options the caller chose', () => { + const body = { model: 'gpt-4o-mini', stream_options: { include_usage: false } }; + expect(withUsageStreamOptions(body)).toBe(body); + }); +}); diff --git a/agentic/agentic-server/src/router.ts b/agentic/agentic-server/src/router.ts index 64098fc589..fb13c9a948 100644 --- a/agentic/agentic-server/src/router.ts +++ b/agentic/agentic-server/src/router.ts @@ -6,6 +6,7 @@ * - Model-based routing: "anthropic/claude-3.5-sonnet" → anthropic provider * - Header-based routing: X-LLM-Provider: ollama → ollama provider * - Fire-and-forget inference metering via an injected InferenceSink + * - Streaming chat completions relayed as server-sent events * - /v1/usage reporting endpoint for external usage submission */ @@ -13,6 +14,7 @@ import { Logger } from '@pgpmjs/logger'; import { Router } from 'express'; import { buildProviderHeaders, resolveProvider, resolveUpstreamUrl } from './providers'; +import { relayEventStream, withUsageStreamOptions } from './streaming'; import { transformChatRequest, transformChatResponse, @@ -50,15 +52,41 @@ export const createRouter = (options: AgenticServerOptions): Router => { model: req.body?.model }); + // Only the OpenAI-compatible wire carries stream frames the gateway can + // relay; translating Ollama's or Anthropic's stream formats is separate + // work, and answering a stream request with a whole JSON body would break + // the client silently. + const streaming = req.body?.stream === true; + if (streaming && provider.type !== 'openai') { + res.status(501).json({ + error: { + message: `streaming is not supported for provider type '${provider.type}'` + } + }); + return; + } + try { const upstreamUrl = resolveUpstreamUrl(provider, '/v1/chat/completions'); - const body = transformChatRequest(provider, req.body || {}); + const transformed = transformChatRequest(provider, req.body || {}); + const body = streaming ? withUsageStreamOptions(transformed) : transformed; const headers = buildProviderHeaders(provider); + const abort = new AbortController(); + // A client that hangs up mid-stream stops the upstream read; the request + // stream's own 'close' fires as soon as the body is consumed, so the + // response is what reports the disconnect. + if (streaming) { + res.on('close', () => { + if (!res.writableEnded) abort.abort(); + }); + } + const upstream = await fetch(upstreamUrl, { method: 'POST', headers, - body: JSON.stringify(body) + body: JSON.stringify(body), + ...(streaming ? { signal: abort.signal } : {}) }); const latencyMs = Number(process.hrtime.bigint() - startTime) / 1e6; @@ -87,6 +115,37 @@ export const createRouter = (options: AgenticServerOptions): Router => { return; } + if (streaming) { + const streamUsage = await relayEventStream(upstream, res); + const streamLatencyMs = Number(process.hrtime.bigint() - startTime) / 1e6; + + log.info('inference complete', { + databaseId, + provider: provider.type, + model: req.body?.model, + promptTokens: streamUsage?.prompt_tokens, + completionTokens: streamUsage?.completion_tokens, + streamed: true + }); + + if (sink) { + sink.logInference({ + databaseId, entityId, actorId, + model: String(req.body?.model || body.model || ''), + provider: provider.type, + service: 'chat', + operation: 'chat/completions', + inputTokens: streamUsage?.prompt_tokens || 0, + outputTokens: streamUsage?.completion_tokens || 0, + totalTokens: streamUsage?.total_tokens || 0, + latencyMs: streamLatencyMs, + status: 'ok', + ...(streamUsage ? { rawUsage: streamUsage } : {}) + }); + } + return; + } + const data = await upstream.json() as Record; const { body: responseBody, usage } = transformChatResponse(data, provider); @@ -133,6 +192,13 @@ export const createRouter = (options: AgenticServerOptions): Router => { }); } + // A stream that failed mid-relay has already sent its status and frames; + // the client sees the truncated stream rather than a JSON error. + if (res.headersSent) { + res.end(); + return; + } + res.status(502).json({ error: { message: 'Failed to reach LLM provider', details: err.message } }); diff --git a/agentic/agentic-server/src/streaming.ts b/agentic/agentic-server/src/streaming.ts new file mode 100644 index 0000000000..9ac4113c02 --- /dev/null +++ b/agentic/agentic-server/src/streaming.ts @@ -0,0 +1,106 @@ +/** + * Server-sent-event passthrough for streaming chat completions. + * + * The gateway relays an OpenAI-compatible `text/event-stream` body verbatim so + * clients receive tokens as the provider emits them, while scanning the frames + * for the terminal usage object that metering needs. + */ + +import type { UsageResult } from './types'; + +/** Accumulates SSE frames, exposing the last usage object the stream carried. */ +export class UsageScanner { + private buffer = ''; + private usage: UsageResult | undefined; + + /** Feed one decoded chunk of the event stream. */ + push(chunk: string): void { + this.buffer += chunk; + + const frames = this.buffer.split('\n'); + this.buffer = frames.pop() ?? ''; + + for (const frame of frames) { + const line = frame.trim(); + if (!line.startsWith('data:')) continue; + + const payload = line.slice('data:'.length).trim(); + if (payload === '' || payload === '[DONE]') continue; + + let parsed: { usage?: Partial | null }; + try { + parsed = JSON.parse(payload); + } catch { + // A provider frame the gateway cannot parse carries no usage it could + // report; relaying is the router's job and happens regardless. + continue; + } + + const usage = parsed.usage; + if (!usage) continue; + + this.usage = { + prompt_tokens: usage.prompt_tokens ?? 0, + completion_tokens: usage.completion_tokens ?? 0, + total_tokens: usage.total_tokens ?? (usage.prompt_tokens ?? 0) + (usage.completion_tokens ?? 0) + }; + } + } + + /** The last usage object seen, or undefined when the provider sent none. */ + result(): UsageResult | undefined { + return this.usage; + } +} + +/** + * Relay an upstream event-stream response to the client byte for byte and + * answer the usage its final frame carried, or undefined when it carried none. + * + * @param upstream - the provider response whose body is being relayed. + * @param res - the client response the frames are written to. + * @returns the usage the stream reported, if any. + */ +export async function relayEventStream( + upstream: { body: ReadableStream | null; headers: { get(name: string): string | null } }, + res: { + setHeader(name: string, value: string): void; + flushHeaders?(): void; + write(chunk: string): void; + end(): void; + } +): Promise { + res.setHeader('Content-Type', upstream.headers.get('content-type') ?? 'text/event-stream'); + res.setHeader('Cache-Control', 'no-cache, no-transform'); + res.setHeader('Connection', 'keep-alive'); + res.flushHeaders?.(); + + const scanner = new UsageScanner(); + + if (upstream.body) { + const decoder = new TextDecoder(); + const reader = upstream.body.getReader(); + + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + const text = decoder.decode(value, { stream: true }); + scanner.push(text); + res.write(text); + } + } + + res.end(); + return scanner.result(); +} + +/** + * Ask an OpenAI-compatible provider to include usage in its final frame. + * Callers that already set `stream_options` keep their own choice. + */ +export function withUsageStreamOptions( + body: Record +): Record { + if (body.stream_options !== undefined) return body; + return { ...body, stream_options: { include_usage: true } }; +}