diff --git a/README.md b/README.md index 5ae2f1f..1ad374b 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ The global context limit defaults to 256k tokens and can be changed with `MAX_CO Inference requests are limited per client IP to 2 requests per second, 100 requests per five hours, and $10 of reported upstream cost per five hours. Railway's `X-Real-IP` header is used to identify clients. -Set `DATABASE_URL` in `.env` (or the process environment) to a Postgres connection string. The proxy creates its `requests` table and index automatically on startup. Every incoming request and its complete response—including streaming responses and errors—is stored with headers, status, model, client IP, and timestamp. Successful inference requests also store a normalized training conversation containing the full chat history and generated assistant output, including reasoning, tool calls, and tool results. +Set `DATABASE_URL` in `.env` (or the process environment) to a Postgres connection string. The proxy creates its `requests` table and index automatically on startup. Only requests to the OpenAI Chat Completions and Responses endpoints and the Anthropic Messages endpoints are stored. Their complete responses—including streaming responses and errors—are stored with headers, status, model, client IP, and timestamp. On startup, stored requests for all other endpoints are deleted. Successful chat completion requests also store a normalized training conversation containing the full chat history and generated assistant output, including reasoning, tool calls, and tool results. ```sh npm start @@ -67,7 +67,9 @@ For Claude Code, either set `ANTHROPIC_BASE_URL=http://localhost:4005/ant` or us ## Exporting training data -Exports use JSONL: each line is one inference request containing its complete input history followed by the generated assistant message. The normalized OpenAI-style `messages` preserve system/user/assistant roles, `reasoning_content`, assistant `tool_calls`, and `tool` results. Non-chat routes and failed inference requests are excluded. +Each export creates two JSONL files. The regular `export-.jsonl` remains unchanged: each line is one inference request containing its complete input history followed by the generated assistant message. The additional `export--traces.jsonl` removes cumulative intermediate snapshots and keeps each complete conversation branch. If a conversation is rewound and continued in multiple ways, every branch leaf is retained as its own full trace. + +Chat Completions, Responses, and Anthropic Messages requests are all exported in the same normalized OpenAI-style `messages` format. It preserves system/user/assistant roles, `reasoning_content`, assistant `tool_calls`, and `tool` results. Non-chat routes and failed inference requests are excluded. Each line also has `metadata` with the exact tool definitions supplied on that request, tool choice, model, API format, endpoint, streaming mode, generation parameters, caller-supplied request metadata, response ID/model, finish reason, usage, HTTP status, duration, request ID, and timestamp. Tool metadata remains in its original OpenAI, Responses, or Anthropic format so no provider-specific schema information is lost. @@ -78,17 +80,17 @@ Each line also has `metadata` with the exact tool definitions supplied on that r Omit the limit to export every training request, newest first: ```sh -npm run --silent export > requests.jsonl +npm run --silent export ``` Pass a positive limit to export that many of the most recent training requests: ```sh -npm run --silent export -- 100 > recent-requests.jsonl +npm run --silent export -- 100 # Equivalent: npm run --silent export -- --limit 100 ``` -The export command uses the same `DATABASE_URL` as the server. `--silent` suppresses npm's banner, and the script's progress message is written to standard error, so redirected JSONL remains valid. Rows collected before normalized training storage was added are normalized from their saved raw request and response during export. +The export command uses the same `DATABASE_URL` as the server and saves both files in the project root. The limit applies to the regular per-request export; the traces file contains the complete branch leaves found among those requests. Rows collected before normalized training storage was added are normalized from their saved raw request and response during export. ## Tests diff --git a/src/chat-completion-endpoints.js b/src/chat-completion-endpoints.js new file mode 100644 index 0000000..60d5d9a --- /dev/null +++ b/src/chat-completion-endpoints.js @@ -0,0 +1,6 @@ +export const CHAT_COMPLETION_ENDPOINTS = Object.freeze([ + '/oai/v1/chat/completions', + '/oai/v1/responses', + '/ant/v1/messages', + '/ant/v1/v1/messages', +]); diff --git a/src/db.js b/src/db.js index c3560f7..cdd4f36 100644 --- a/src/db.js +++ b/src/db.js @@ -1,4 +1,5 @@ import pg from 'pg'; +import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js'; const { Pool } = pg; @@ -38,6 +39,14 @@ export async function ensureSchema(pool) { await pool.query('CREATE INDEX IF NOT EXISTS requests_created_at_idx ON requests (created_at DESC, id DESC)'); } +export async function removeNonChatCompletionRequests(pool) { + const result = await pool.query(` + DELETE FROM requests + WHERE split_part(endpoint, '?', 1) <> ALL($1::text[]) + `, [CHAT_COMPLETION_ENDPOINTS]); + return result.rowCount ?? 0; +} + export async function insertRequest(pool, entry) { await pool.query( `INSERT INTO requests ( diff --git a/src/export.js b/src/export.js index 8d10432..74d8613 100644 --- a/src/export.js +++ b/src/export.js @@ -1,7 +1,27 @@ import './env.js'; import { once } from 'node:events'; +import { createWriteStream } from 'node:fs'; +import { join, dirname } from 'node:path'; +import { fileURLToPath } from 'node:url'; import { createPool, ensureSchema } from './db.js'; import { createTrainingExample } from './training-data.js'; +import { createFullTraceSelector } from './traces.js'; +import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js'; + +const __dirname = dirname(fileURLToPath(import.meta.url)); + +function timestamp() { + const d = new Date(); + const pad = (n) => String(n).padStart(2, '0'); + return ( + `${d.getFullYear()}${pad(d.getMonth() + 1)}${pad(d.getDate())}` + + `-${pad(d.getHours())}${pad(d.getMinutes())}${pad(d.getSeconds())}` + ); +} + +const exportTimestamp = timestamp(); +const exportPath = join(__dirname, '..', `export-${exportTimestamp}.jsonl`); +const tracesPath = join(__dirname, '..', `export-${exportTimestamp}-traces.jsonl`); function usage() { return 'Usage: npm run export -- [limit]\n npm run export -- --limit '; @@ -28,26 +48,36 @@ try { const limit = parseLimit(process.argv.slice(2)); pool = createPool(); await ensureSchema(pool); + const stream = createWriteStream(exportPath, { flags: 'w' }); + const tracesStream = createWriteStream(tracesPath, { flags: 'w' }); + stream.on('error', (err) => { + console.error(`Error writing to ${exportPath}: ${err.message}`); + process.exitCode = 1; + }); + tracesStream.on('error', (err) => { + console.error(`Error writing to ${tracesPath}: ${err.message}`); + process.exitCode = 1; + }); + const isFullTrace = createFullTraceSelector(); let exported = 0; + let tracesExported = 0; let cursor = null; while (limit === null || exported < limit) { const remaining = limit === null ? 1_000 : limit - exported; const batchSize = Math.min(1_000, Math.max(100, remaining * 2)); - const params = cursor ? [cursor.cursor_created_at, cursor.id, batchSize] : [batchSize]; + const params = cursor + ? [cursor.cursor_created_at, cursor.id, CHAT_COMPLETION_ENDPOINTS, batchSize] + : [CHAT_COMPLETION_ENDPOINTS, batchSize]; const cursorClause = cursor ? 'AND (created_at, id) < ($1::timestamptz, $2::bigint)' : ''; - const batchParameter = cursor ? '$3' : '$1'; + const endpointParameter = cursor ? '$3' : '$1'; + const batchParameter = cursor ? '$4' : '$2'; const result = await pool.query( `SELECT id, created_at, created_at::text AS cursor_created_at, endpoint, request_body, response_body, response_status, duration_ms, stream, training_data FROM requests WHERE (method = 'POST' OR method IS NULL) AND response_status BETWEEN 200 AND 299 - AND ( - endpoint LIKE '/oai/v1/chat/completions%' - OR endpoint LIKE '/oai/v1/responses%' - OR endpoint LIKE '/ant/v1/messages%' - OR endpoint LIKE '/ant/v1/v1/messages%' - ) + AND split_part(endpoint, '?', 1) = ANY(${endpointParameter}::text[]) ${cursorClause} ORDER BY created_at DESC, id DESC LIMIT ${batchParameter}`, @@ -74,14 +104,23 @@ try { ...(row.duration_ms == null ? {} : { duration_ms: row.duration_ms }), }, }; - if (!process.stdout.write(`${JSON.stringify(exportedExample)}\n`)) await once(process.stdout, 'drain'); + if (!stream.write(`${JSON.stringify(exportedExample)}\n`)) await once(stream, 'drain'); + if (isFullTrace(exportedExample)) { + if (!tracesStream.write(`${JSON.stringify(exportedExample)}\n`)) await once(tracesStream, 'drain'); + tracesExported += 1; + } exported += 1; if (limit !== null && exported >= limit) break; } if (result.rows.length < batchSize || (limit !== null && exported >= limit)) break; cursor = result.rows.at(-1); } - console.error(`Exported ${exported} training request(s) as JSONL, newest first`); + const finished = Promise.all([once(stream, 'finish'), once(tracesStream, 'finish')]); + stream.end(); + tracesStream.end(); + await finished; + console.error(`Exported ${exported} training request(s) as JSONL to ${exportPath}, newest first`); + console.error(`Exported ${tracesExported} complete trace branch(es) as JSONL to ${tracesPath}, newest first`); } catch (error) { console.error(`Error: ${error.message}`); process.exitCode = 1; diff --git a/src/proxy.js b/src/proxy.js index d6d21fd..e631ab3 100644 --- a/src/proxy.js +++ b/src/proxy.js @@ -15,6 +15,7 @@ import { } from './anthropic.js'; import { translateResponsesRequest, translateResponsesResponse, translateResponsesChunk } from './responses.js'; import { createTrainingExample } from './training-data.js'; +import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js'; export const DEFAULT_UPSTREAM_BASE_URL = 'https://opencode.ai/zen/go/v1'; @@ -125,7 +126,7 @@ function withContextLimit(payload, maxContextTokens) { export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBaseUrl = DEFAULT_UPSTREAM_BASE_URL, maxContextTokens = maxContextFromEnv(), rateLimiter = createRateLimiter(), pagesDirectory = defaultPagesDirectory, requestLogger } = {}) { const rotator = createKeyRotator(keys); const app = express(); - if (requestLogger) app.use(collectRequestData(requestLogger)); + if (requestLogger) app.post(CHAT_COMPLETION_ENDPOINTS, collectRequestData(requestLogger)); app.use(express.json({ limit: '10mb' })); app.get('/', (_request, response) => { diff --git a/src/server.js b/src/server.js index 05fe25e..656f4f8 100644 --- a/src/server.js +++ b/src/server.js @@ -1,12 +1,13 @@ import path from 'node:path'; import { createProxyApp } from './proxy.js'; import { projectRoot } from './env.js'; -import { createPool, createRequestLogger, ensureSchema } from './db.js'; +import { createPool, createRequestLogger, ensureSchema, removeNonChatCompletionRequests } from './db.js'; import { loadConfiguredKeys } from './key-store.js'; const keys = await loadConfiguredKeys({ path: path.join(projectRoot, 'keys.txt') }); const pool = createPool(); await ensureSchema(pool); +await removeNonChatCompletionRequests(pool); const app = createProxyApp({ keys, requestLogger: createRequestLogger(pool) }); const port = Number(process.env.PORT || 4005); diff --git a/src/traces.js b/src/traces.js new file mode 100644 index 0000000..bd26af5 --- /dev/null +++ b/src/traces.js @@ -0,0 +1,54 @@ +function normalizedJson(value, parentKey = '') { + if (Array.isArray(value)) return value.map((item) => normalizedJson(item, parentKey)); + if (!value || typeof value !== 'object') return value; + + return Object.fromEntries(Object.keys(value).sort().flatMap((key) => { + const child = value[key]; + if (child === undefined) return []; + + // Tool arguments are JSON strings in normalized messages. Providers may + // reserialize them with different whitespace when they send the next turn. + if (parentKey === 'function' && key === 'arguments' && typeof child === 'string') { + try { + return [[key, normalizedJson(JSON.parse(child))]]; + } catch { + // Non-JSON arguments still compare as their original string. + } + } + return [[key, normalizedJson(child, key)]]; + })); +} + +function messageKey(message) { + return JSON.stringify(normalizedJson(message)); +} + +/** + * Selects maximal conversations from examples ordered newest first. + * + * Each trie node represents a message prefix already seen in a newer request. + * A conversation whose final node already exists is an intermediate snapshot + * (or duplicate) of that newer conversation. Divergent children remain + * independent, so every leaf of a rewound/branched conversation is retained. + */ +export function createFullTraceSelector() { + const root = new Map(); + + return (example) => { + if (!Array.isArray(example?.messages) || example.messages.length === 0) return true; + + let children = root; + let finalNodeExisted = true; + for (const message of example.messages) { + const key = messageKey(message); + let node = children.get(key); + if (!node) { + node = { children: new Map() }; + children.set(key, node); + finalNodeExisted = false; + } + children = node.children; + } + return !finalNodeExisted; + }; +} diff --git a/test/db.test.js b/test/db.test.js index 4f9ee5f..a23a0e3 100644 --- a/test/db.test.js +++ b/test/db.test.js @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest'; -import { createRequestLogger, ensureSchema, insertRequest } from '../src/db.js'; +import { createRequestLogger, ensureSchema, insertRequest, removeNonChatCompletionRequests } from '../src/db.js'; describe('request database', () => { it('creates the requests table, compatible columns, and recent-request index', async () => { @@ -37,6 +37,19 @@ describe('request database', () => { ]); }); + it('removes requests that are not for the chat completions endpoint', async () => { + const pool = { query: vi.fn().mockResolvedValue({ rowCount: 4 }) }; + + await expect(removeNonChatCompletionRequests(pool)).resolves.toBe(4); + + expect(pool.query).toHaveBeenCalledOnce(); + expect(pool.query.mock.calls[0][0]).toContain("split_part(endpoint, '?', 1) <> ALL($1::text[])"); + expect(pool.query.mock.calls[0][1]).toEqual([[ + '/oai/v1/chat/completions', '/oai/v1/responses', + '/ant/v1/messages', '/ant/v1/v1/messages', + ]]); + }); + it('reports storage failures without rejecting request handling', async () => { const onError = vi.fn(); const logger = createRequestLogger({ query: vi.fn().mockRejectedValue(new Error('database offline')) }, { onError }); diff --git a/test/proxy.test.js b/test/proxy.test.js index 262d218..c7e168c 100644 --- a/test/proxy.test.js +++ b/test/proxy.test.js @@ -126,17 +126,37 @@ describe('proxy', () => { }); }); - it('collects error and not-found responses too', async () => { + it('does not collect requests to non-chat-completion endpoints', async () => { + const fetchImpl = vi.fn().mockResolvedValue(response('{"object":"list","data":[]}')); const requestLogger = { log: vi.fn() }; - const app = createProxyApp({ keys: ['key'], requestLogger }); + const app = createProxyApp({ keys: ['key'], fetchImpl, requestLogger }); await request(app).get('/missing').expect(404); + await request(app).get('/oai/v1/models').expect(200); - expect(requestLogger.log).toHaveBeenCalledOnce(); - expect(requestLogger.log.mock.calls[0][0]).toMatchObject({ - method: 'GET', endpoint: '/missing', requestBody: null, responseStatus: 404, - }); - expect(requestLogger.log.mock.calls[0][0].responseBody).toContain('Cannot GET /missing'); + expect(requestLogger.log).not.toHaveBeenCalled(); + }); + + it('collects OpenAI Responses and Anthropic Messages requests', async () => { + const cases = [ + { endpoint: '/oai/v1/responses', body: { model: 'm', input: 'Hi' } }, + { endpoint: '/ant/v1/messages', body: { model: 'm', max_tokens: 10, messages: [{ role: 'user', content: 'Hi' }] } }, + { endpoint: '/ant/v1/v1/messages', body: { model: 'm', max_tokens: 10, messages: [{ role: 'user', content: 'Hi' }] } }, + ]; + + for (const { endpoint, body } of cases) { + const fetchImpl = vi.fn().mockResolvedValue(response(JSON.stringify({ + id: 'chatcmpl-response', model: 'm', + choices: [{ message: { role: 'assistant', content: 'Hello' }, finish_reason: 'stop' }], + }))); + const requestLogger = { log: vi.fn() }; + const app = createProxyApp({ keys: ['key'], fetchImpl, requestLogger }); + + await request(app).post(endpoint).send(body).expect(200); + + expect(requestLogger.log).toHaveBeenCalledOnce(); + expect(requestLogger.log.mock.calls[0][0]).toMatchObject({ method: 'POST', endpoint, requestBody: body }); + } }); it('translates Responses requests and returns a Responses object', async () => { diff --git a/test/traces.test.js b/test/traces.test.js new file mode 100644 index 0000000..3c06591 --- /dev/null +++ b/test/traces.test.js @@ -0,0 +1,55 @@ +import { describe, expect, it } from 'vitest'; +import { createFullTraceSelector } from '../src/traces.js'; + +const message = (role, content, extra = {}) => ({ role, content, ...extra }); +const example = (...messages) => ({ messages }); + +describe('full trace selection', () => { + it('keeps only the maximal snapshot of a growing conversation', () => { + const user = message('user', 'Start'); + const first = message('assistant', 'First'); + const next = message('user', 'Continue'); + const last = message('assistant', 'Done'); + const examples = [ + example(user, first, next, last), + example(user, first), + ]; + + expect(examples.filter(createFullTraceSelector())).toEqual([examples[0]]); + }); + + it('keeps every leaf when a conversation is rewound and branched', () => { + const shared = [message('user', 'Start'), message('assistant', 'Shared')]; + const branchC = example(...shared, message('user', 'Choose C'), message('assistant', 'C')); + const branchD = example(...shared, message('user', 'Choose D'), message('assistant', 'D')); + const intermediate = example(...shared); + + const examples = [branchD, branchC, intermediate]; + expect(examples.filter(createFullTraceSelector())).toEqual([branchD, branchC]); + }); + + it('matches JSON tool arguments despite provider whitespace changes', () => { + const requestToolCall = message('assistant', null, { + tool_calls: [{ id: 'call_1', type: 'function', function: { name: 'run', arguments: '{ "path": "a" }' } }], + }); + const echoedToolCall = message('assistant', null, { + tool_calls: [{ type: 'function', function: { arguments: '{"path":"a"}', name: 'run' }, id: 'call_1' }], + }); + const newest = example( + message('user', 'Run it'), echoedToolCall, + message('tool', 'ok', { tool_call_id: 'call_1' }), message('assistant', 'Done'), + ); + const intermediate = example(message('user', 'Run it'), requestToolCall); + + const examples = [newest, intermediate]; + expect(examples.filter(createFullTraceSelector())).toEqual([newest]); + }); + + it('does not combine unrelated conversations', () => { + const examples = [ + example(message('user', 'One'), message('assistant', '1')), + example(message('user', 'Two'), message('assistant', '2')), + ]; + expect(examples.filter(createFullTraceSelector())).toEqual(examples); + }); +});