LLM Proxy

import {
  wrapLLMClient,
  createRehydraFetch,
  createRehydraProxy,
  createRehydraProxyServer,
  SSEParser,
  OpenAIProvider,
  AnthropicProvider,
  detectProvider,
} from 'rehydra/proxy';

function wrapLLMClient<T>(
  client: T,
  config: Omit<RehydraFetchConfig, 'provider'>
): T

import OpenAI from 'openai';
import { wrapLLMClient } from 'rehydra/proxy';
import { InMemoryKeyProvider, SQLitePIIStorageProvider } from 'rehydra';

const storage = new SQLitePIIStorageProvider('./pii.db');
await storage.initialize();

const client = wrapLLMClient(new OpenAI(), {
  anonymizer: { ner: { mode: 'quantized' } },
  keyProvider: new InMemoryKeyProvider(),
  piiStorageProvider: storage,
  getSessionId: () => 'session-1',
});

function createRehydraFetch(config: RehydraFetchConfig): typeof globalThis.fetch

import { createRehydraFetch } from 'rehydra/proxy';
import { InMemoryKeyProvider, SQLitePIIStorageProvider } from 'rehydra';

const storage = new SQLitePIIStorageProvider('./pii.db');
await storage.initialize();

const fetch = createRehydraFetch({
  anonymizer: { ner: { mode: 'quantized' } },
  keyProvider: new InMemoryKeyProvider(),
  piiStorageProvider: storage,
  provider: 'openai',  // or 'anthropic', or auto-detected
  getSessionId: () => 'session-1',
});

const response = await fetch('https://api.openai.com/v1/chat/completions', {
  method: 'POST',
  body: JSON.stringify({ model: 'gpt-4', messages: [...] }),
});

function createRehydraProxy(
  config: RehydraProxyConfig
): (request: Request) => Promise<Response>

import { createRehydraProxy } from 'rehydra/proxy';
import { InMemoryKeyProvider, SQLitePIIStorageProvider } from 'rehydra';

const storage = new SQLitePIIStorageProvider('./pii.db');
await storage.initialize();

const proxy = createRehydraProxy({
  upstream: 'https://api.openai.com',
  anonymizer: { ner: { mode: 'quantized' } },
  keyProvider: new InMemoryKeyProvider(),
  piiStorageProvider: storage,
  getSessionId: () => 'session-1',
});

Bun.serve({ port: 8080, fetch: proxy });

function createRehydraProxyServer(
  config: RehydraProxyServerConfig
): Promise<RehydraProxyServer>

interface RehydraProxyServer {
  server: http.Server;
  port: number;
  host: string;
  close(): Promise<void>;
}

import { createRehydraProxyServer } from 'rehydra/proxy';
import { InMemoryKeyProvider, SQLitePIIStorageProvider } from 'rehydra';

const storage = new SQLitePIIStorageProvider('./pii.db');
await storage.initialize();

const server = await createRehydraProxyServer({
  upstream: 'https://api.openai.com',
  port: 8080,
  anonymizer: { ner: { mode: 'quantized' } },
  keyProvider: new InMemoryKeyProvider(),
  piiStorageProvider: storage,
  getSessionId: () => 'session-1',
});

// Point clients at http://localhost:8080/v1
await server.close();

interface RehydraFetchConfig {
  // Encryption key provider (required)
  keyProvider: KeyProvider;

  // Storage for session persistence (required)
  piiStorageProvider: PIIStorageProvider;

  // Anonymizer configuration (includes tagFormat)
  anonymizer?: AnonymizerConfig;

  // Policy override
  policy?: Partial<AnonymizationPolicy>;

  // Provider hint (auto-detected if omitted)
  provider?: 'openai' | 'anthropic' | 'auto';

  // Derive session ID from request
  getSessionId?: (request: Request) => string | Promise<string>;

  // Handle streaming responses (default: true)
  handleStreaming?: boolean;

  // Locale hint
  locale?: string;

  // Tool execution callback (non-streaming only)
  // Arguments are rehydrated before calling; return value is anonymized
  onToolCall?: (
    name: string,
    args: Record<string, unknown>,
    toolCallId: string,
  ) => unknown | Promise<unknown>;

  // Maximum tool execution rounds (default: 10)
  maxToolRounds?: number;

  // PII system instruction: string to override, false to disable, undefined for default
  systemInstruction?: string | false;
}

interface RehydraProxyConfig extends RehydraFetchConfig {
  // Upstream LLM API URL (e.g., 'https://api.openai.com')
  upstream: string;

  // Headers to forward to upstream
  // Default: ['authorization', 'content-type', 'x-api-key',
  //   'anthropic-version', 'openai-organization', 'openai-project']
  forwardHeaders?: string[];

  // Strip this prefix from the request path before forwarding
  stripPrefix?: string;

  // API key to inject into upstream requests
  // Adds Authorization header (OpenAI) or x-api-key header (Anthropic)
  apiKey?: string;
}

function detectProvider(
  url: string,
  headers: Headers,
  hint?: string
): LLMContentProvider

interface LLMContentProvider {
  readonly name: string;

  // Request detection
  matchesRequest(url: string, headers: Headers): boolean;
  isStreamingRequest(body: unknown): boolean;

  // Request handling
  extractRequestText(body: unknown): string[];
  rebuildRequestBody(body: unknown, anonymizedTexts: string[]): unknown;

  // Non-streaming response handling
  extractResponseText(body: unknown): string[];
  rebuildResponseBody(body: unknown, rehydratedTexts: string[]): unknown;

  // Streaming content handling
  extractSSEDelta(data: unknown): string | null;
  rebuildSSEDelta(data: unknown, rehydratedText: string): unknown;

  // Tool/function call handling (non-streaming)
  extractResponseToolCalls?(body: unknown): string[];
  rebuildResponseToolCalls?(body: unknown, rehydratedArgs: string[]): unknown;

  // Tool/function call handling (streaming)
  extractSSEToolCallDeltas?(data: unknown): ToolCallDelta[] | null;
  rebuildSSEToolCallDeltas?(data: unknown, rehydratedArgs: Map<number, string>): unknown;
  extractSSEToolCallStop?(data: unknown): number | null;
}

interface ToolCallDelta {
  index: number;
  arguments: string;
}

interface SSEEvent {
  event: string;   // defaults to 'message'
  data: string;
}

const parser = new SSEParser();

// Parse SSE text
const events = parser.parse(chunk);
// [{ event: 'message', data: '...' }, ...]

// Check for stream end
import { isSSEDone } from 'rehydra/proxy';
if (isSSEDone(event.data)) { /* stream complete */ }

// Serialize back to SSE format
import { serializeSSEEvent } from 'rehydra/proxy';
const text = serializeSSEEvent({ event: 'message', data: '...' });