TypeScript LLM WebSocket Worker

A modular WebSocket worker that processes chat messages using multiple LLM providers (OpenAI, Anthropic, Google Gemini) with support for workflows and structured JSON responses.

Features

• 🔌 WebSocket API - Real-time communication
• 🤖 Multi-Provider Support - OpenAI, Anthropic (Claude), Google (Gemini)
• 🔄 Workflow Engine - Chain multiple LLM calls (Guardrails → Reasoning → Formatting)
• 📊 Structured Responses - Zod schema validation for JSON outputs
• 💾 Database Persistence - PostgreSQL with Prisma ORM
• ✅ Type-Safe - Full TypeScript implementation
• 🎨 Factory Pattern - Easy to extend with new providers

Tech Stack

• Runtime: Node.js 18+
• Language: TypeScript
• Database: PostgreSQL with Prisma
• LLM SDK: Vercel AI SDK
• WebSocket: ws library
• Validation: Zod

Project Structure

src/
├── index.ts                 # WebSocket server entry point
├── config/
│   └── schemas.ts          # Response schemas
├── adapters/
│   ├── LLMAdapter.ts       # Base adapter interface
│   ├── OpenAIAdapter.ts    # OpenAI implementation
│   ├── AnthropicAdapter.ts # Anthropic implementation
│   └── GeminiAdapter.ts    # Google Gemini implementation
├── services/
│   ├── AiFactory.ts        # Factory for LLM providers
│   ├── MessageService.ts   # Database operations
│   ├── PromptBuilder.ts    # System prompt construction
│   └── WorkflowEngine.ts   # Multi-step LLM workflows
├── validators/
│   └── schemas.ts          # Zod validation schemas
└── types/
    └── index.ts            # TypeScript types

Setup

1. Install Dependencies

npm install

2. Configure Environment Variables

Create a .env file:

DATABASE_URL="postgresql://user:password@localhost:5432/llm_worker?schema=public"
OPENAI_API_KEY="your-openai-api-key"
ANTHROPIC_API_KEY="your-anthropic-api-key"
GOOGLE_API_KEY="your-google-api-key"
PORT=8080

3. Setup Database

# Run migrations
npm run prisma:migrate

# Seed database with sample data
npm run db:seed

4. Start the Server

# Development mode with hot reload
npm run dev

# Production mode
npm run build
npm start

Usage

WebSocket Connection

Connect to ws://localhost:8080

Message Format

Send a JSON payload to process a chat:

{
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "threadId": "thread_123",
  "userMessage": "How do I ladder my FDs?",
  "prompt": "You are a helpful finance assistant.",
  "workflow": ["guardrails", "reason", "format"],
  "provider": "openai:gpt-4o-mini",
  "responseSchema": "financeAdviceV1"
}

Supported Providers

• OpenAI: openai:gpt-4o-mini, openai:gpt-4o, openai:gpt-4-turbo
• Anthropic: anthropic:claude-3-5-sonnet-20241022, anthropic:claude-3-opus-20240229
• Google: google:gemini-1.5-flash, google:gemini-1.5-pro

Response Schemas

Available schema names:

• financeAdviceV1 - Financial advice with reasoning and risks
• generic - General purpose response
• guardrails - Safety check response
• reasoning - Step-by-step reasoning

Workflow Steps

• guardrails - Content safety and policy checks
• reason - Detailed step-by-step reasoning
• format - Structure output according to schema

Response Format

The server sends multiple message types:

1. Connected

{
  "type": "connected",
  "message": "Connected to LLM Worker",
  "supportedProviders": ["openai", "anthropic", "google"]
}

2. Acknowledgment

{
  "type": "ack",
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "message": "Job received and processing started"
}

3. Progress Updates

{
  "type": "progress",
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "message": "Executing workflow: guardrails → reason → format"
}

4. Complete

{
  "type": "complete",
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "messageId": "cm1x2y3z4",
  "response": {
    "text": "Based on your query...",
    "json": {
      "advice": "...",
      "reasoning": "...",
      "risks": ["..."],
      "recommendations": ["..."]
    },
    "reasoning": "Step-by-step analysis..."
  },
  "usage": {
    "promptTokens": 150,
    "completionTokens": 300,
    "totalTokens": 450
  },
  "processingTime": 2500
}

5. Error

{
  "type": "error",
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "error": "Error message"
}

Testing with Postman

1. Open Postman and create a new WebSocket request
2. Connect to ws://localhost:8080
3. Send test messages using the format above
4. View real-time responses

Example Test Messages

Simple Query:

{
  "jobId": "550e8400-e29b-41d4-a716-446655440000",
  "threadId": "thread_test_001",
  "userMessage": "What is a mutual fund?",
  "prompt": "You are a helpful finance assistant.",
  "provider": "openai:gpt-4o-mini"
}

With Workflow:

{
  "jobId": "550e8400-e29b-41d4-a716-446655440001",
  "threadId": "thread_test_001",
  "userMessage": "Should I invest in stocks or bonds?",
  "prompt": "You are a helpful finance assistant.",
  "workflow": ["guardrails", "reason", "format"],
  "provider": "anthropic:claude-3-5-sonnet-20241022",
  "responseSchema": "financeAdviceV1"
}

Testing Different Providers:

{
  "jobId": "550e8400-e29b-41d4-a716-446655440002",
  "threadId": "thread_test_002",
  "userMessage": "Explain compound interest",
  "prompt": "You are a helpful finance assistant.",
  "provider": "google:gemini-1.5-flash",
  "responseSchema": "generic"
}

Development

Scripts

• npm run dev - Start development server with hot reload
• npm run build - Build for production
• npm start - Run production server
• npm run prisma:migrate - Run database migrations
• npm run prisma:generate - Generate Prisma client
• npm run db:seed - Seed database with sample data

Adding New Providers

1. Create adapter in src/adapters/NewProviderAdapter.ts
2. Implement LLMAdapter interface
3. Register in AiFactory.ts

Adding New Response Schemas

1. Define Zod schema in src/config/schemas.ts
2. Add to ResponseSchemas map
3. Use schema name in job payload

Architecture

Adapter Pattern

Each LLM provider has its own adapter implementing a common interface, making it easy to add new providers.

Factory Pattern

AiFactory creates the appropriate adapter based on provider string.

Workflow Pattern

WorkflowEngine chains multiple LLM calls for complex processing pipelines.

Service Layer

Separates business logic (services) from infrastructure (adapters) and presentation (WebSocket).

License

MIT

Contributing

Contributions welcome! Please feel free to submit a Pull Request.