# Orpheus WebSocket Chat System - Technical Documentation

## Overview

This document provides a comprehensive technical overview of the Orpheus WebSocket chat system, a real-time streaming chat interface designed specifically for AI-powered music generation using the Suno Studio API. The system implements a sophisticated WebSocket protocol that handles streaming text responses, real-time song generation, and embedded audio playback.

## System Architecture

The system follows a three-tier architecture with real-time WebSocket communication:

```mermaid
graph TB
    subgraph "Client Layer"
        A["Browser<br/>(OrpheusWebSocketChat)"]
        B["LocalStorage<br/>(Token & Chat UUID)"]
    end
    
    subgraph "Server Layer"
        C["FastAPI WebSocket<br/>(/ws/{chat_uuid})"]
        D["ConnectionManager<br/>(Per-connection state)"]
        E["Chat History<br/>(JSONL files)"]
    end
    
    subgraph "External APIs"
        F["OpenAI API<br/>(GPT-4 + Tools)"]
        G["Suno Studio API<br/>(Song Generation)"]
    end
    
    A <--> C
    A <--> B
    C <--> D
    D <--> E
    D <--> F
    D <--> G
    
    C -.->|"Real-time polling"| G
```

### Core Components

- **Frontend Client**: JavaScript-based WebSocket client (`OrpheusWebSocketChat`)
- **Backend Server**: FastAPI with WebSocket support and connection management
- **External APIs**: OpenAI GPT-4 for conversation, Suno Studio for music generation
- **Persistence**: JSONL-based chat history storage

### Key Features

- **Streaming Text Responses**: Real-time character-by-character streaming from OpenAI
- **Embedded Audio Playback**: Direct integration of generated songs in chat messages
- **Real-time Song Generation**: Live updates during Suno API processing
- **Persistent Chat History**: Session-based storage with UUID-based identification
- **Debug Mode**: Comprehensive function call debugging and monitoring

## WebSocket Protocol Specification

The system implements a custom WebSocket protocol with four primary message types:

### Client → Server Messages

```javascript
{
  "content": "User message content",
  "token": "suno_studio_bearer_token"
}
```

### Server → Client Message Types

#### 1. `new_message`
Creates a new message in the chat history.

```javascript
{
  "type": "new_message",
  "data": {
    "message_id": "uuid",
    "chat_id": "chat_uuid", 
    "role": "user|assistant|system",
    "content": "Message content",
    "status": "pending|complete|error",
    "timestamp": "ISO timestamp",
    "metadata": {
      "clips": [...], // Song clips array
      "type": "song_generation_pending|function_call_debug"
    }
  }
}
```

#### 2. `update_content`
Streams content updates for real-time text generation.

```javascript
{
  "type": "update_content",
  "message_id": "target_message_uuid",
  "content_delta": "incremental_text_chunk"
}
```

#### 3. `update_status`
Updates message processing status.

```javascript
{
  "type": "update_status", 
  "message_id": "target_message_uuid",
  "status": "pending|complete|error|executing_tool"
}
```

#### 4. `update_clips`
Updates embedded audio clips without modifying message content.

```javascript
{
  "type": "update_clips",
  "message_id": "target_message_uuid", 
  "clips": [
    {
      "id": "clip_uuid",
      "title": "Song Title",
      "status": "pending|streaming|complete|error",
      "audio_url": "https://...",
      "image_url": "https://...",
      "error_message": "Error details if failed"
    }
  ]
}
```

## Message Flow Architecture

### Complete Conversation Flow

```mermaid
sequenceDiagram
    participant U as User
    participant C as Client
    participant S as Server
    participant O as OpenAI API
    participant SU as Suno API
    
    U->>C: Type message
    C->>S: WebSocket: {"content": "...", "token": "..."}
    S->>S: Store user message
    S->>C: new_message (user)
    
    S->>O: Stream chat completion
    S->>C: new_message (assistant, pending)
    
    loop Streaming Response
        O-->>S: Content delta
        S->>C: update_content (delta)
    end
    
    alt Tool Call Required
        O->>S: Tool call: generate_song
        S->>C: update_status (executing_tool)
        S->>C: new_message (system debug)
        
        S->>SU: Generate song request
        SU->>S: Clip IDs + metadata
        S->>C: new_message (song_generation_pending)
        
        loop Polling Song Status
            S->>SU: Check clip status
            SU-->>S: Clip progress
            S->>C: update_clips (real-time)
        end
        
        SU->>S: Song complete + audio_url
        S->>C: update_clips (complete)
    else Text Only
        S->>C: update_status (complete)
    end
```

## Song Generation Workflow

The song generation process involves multiple asynchronous steps with real-time status updates:

### Song Generation State Machine

### Polling Strategy

The server implements an intelligent polling mechanism for song generation:

- **Polling Interval**: 2 seconds (optimized for quick streaming detection)
- **Maximum Attempts**: 60 (2 minutes total timeout)
- **Early Termination**: Stops polling once `status === 'complete'`
- **Error Handling**: Captures and reports generation failures

### Clip Update Strategy

The system uses a sophisticated clip update mechanism that preserves audio player state:

1. **Targeted Updates**: Only update changed clips, preserve existing audio players
2. **Embedded Rendering**: Clips are rendered directly in message content HTML
3. **Progressive Enhancement**: Audio players appear once `audio_url` is available
4. **Status Tracking**: Visual status indicators update in real-time

## Client Implementation (`OrpheusWebSocketChat`)

### Core Architecture Patterns

#### 1. Message State Management

```javascript
class OrpheusWebSocketChat {
    constructor() {
        this.messages = [];  // Central message store
        this.websocket = null;
        this.chatUuid = this.getChatUuid();
    }
    
    handleWebSocketMessage(eventData) {
        // Route messages by type for targeted updates
        switch (eventData.type) {
            case 'new_message': 
                this.handleNewMessage(eventData.data);
                break;
            case 'update_content':
                this.handleContentUpdate(eventData);
                break;
            // ... additional types
        }
    }
}
```

#### 2. Streaming Text Updates

The client handles incremental content updates efficiently:

```javascript
handleContentUpdate(update) {
    const message = this.messages.find(m => m.message_id === update.message_id);
    if (message) {
        // Append streaming deltas to existing content
        message.content += update.content_delta;
        this.updateRenderedMessage(message);
    }
}
```

#### 3. Clip Management System

```javascript
updateClipsOnly(message) {
    // Smart update: only modify clips that changed
    message.metadata.clips.forEach((clip, index) => {
        const existingClip = contentDiv.querySelector(`[data-clip-id="${clip.id}"]`);
        if (existingClip) {
            this.updateSingleClip(existingClip, clip);  // Preserve audio state
        } else {
            // Add new clip without disrupting existing ones
            const clipHtml = this.createClipHtml(clip, index);
            contentDiv.insertAdjacentHTML('beforeend', clipHtml);
        }
    });
}
```

### Key Technical Patterns

#### 1. **Optimistic UI Updates**
- Messages appear immediately in the UI before server confirmation
- Status indicators provide real-time feedback during processing
- Error states are handled gracefully with rollback capability

#### 2. **Progressive Enhancement**
- Audio players only appear when `audio_url` becomes available
- Thumbnails load asynchronously without blocking the UI
- Fallback loading states for all external content

#### 3. **Connection Resilience**
- Automatic reconnection with exponential backoff
- Message queue persistence during disconnections
- Graceful degradation when WebSocket is unavailable

#### 4. **Memory Management**
- Efficient message rendering with targeted DOM updates
- Audio player preservation during clip updates
- Cleanup of event listeners on component destruction

## Data Persistence

### Chat History Format (JSONL)

Each message is stored as a single JSON line in `./chats_ws/{chat_uuid}.jsonl`:

```json
{"message_id": "uuid", "chat_id": "uuid", "role": "user", "content": "Hello", "status": "complete", "timestamp": "2024-01-01T00:00:00Z", "metadata": {}}
{"message_id": "uuid", "chat_id": "uuid", "role": "assistant", "content": "Hi there!", "status": "complete", "timestamp": "2024-01-01T00:00:01Z", "metadata": {"clips": [...]}}
```

### Message Metadata Structure

```javascript
{
  "clips": [
    {
      "id": "suno_clip_uuid",
      "title": "Generated Song Title", 
      "status": "complete",
      "audio_url": "https://...",
      "image_url": "https://...",
      "error_message": null
    }
  ],
  "type": "song_generation_pending|function_call_debug",
  "total_clips": 2,
  "completed_clips": 1
}
```

## Configuration & Environment

### Required Environment Variables

```bash
OPENAI_API_KEY=sk-...           # OpenAI API key for GPT-4
ROOT_PATH=/                     # Optional: API root path for reverse proxy
```

### Client Configuration

```javascript
// LocalStorage keys used by client
const CONFIG_KEYS = {
    CHAT_UUID: 'orpheus_ws_chat_uuid',
    SUNO_TOKEN: 'orpheus_suno_token'
};

// WebSocket connection settings
const WS_CONFIG = {
    RECONNECT_DELAY: 3000,      // 3 second reconnection delay
    MAX_RECONNECT_ATTEMPTS: -1   // Infinite reconnection attempts
};
```

## API Integration Details

### OpenAI Integration

The system uses OpenAI's function calling with streaming:

```python
SONGWRITING_TOOLS = [{
    "type": "function",
    "function": {
        "name": "generate_song",
        "description": "Generate a new song using Suno Studio API",
        "parameters": {
            "type": "object", 
            "properties": {
                "prompt": {"type": "string"},
                "tags": {"type": "string"},
                "title": {"type": "string"},
                "make_instrumental": {"type": "boolean"},
                "continue_clip_id": {"type": "string"},
                "continue_at": {"type": "number"}
            },
            "required": ["prompt"]
        }
    }
}]
```

### Suno Studio API Integration

```python
# API Endpoints Used
SUNO_STUDIO_GENERATE_SONG_URL = "https://studio-api.staging.suno.com/api/generate/v2-web"
SUNO_STUDIO_FEED_URL = "https://studio-api.staging.suno.com/api/feed/v2"

# Required Headers
headers = {
    "Authorization": f"Bearer {suno_token}",
    "Content-Type": "application/json"
}
```

## Error Handling & Resilience

### Connection Management

1. **WebSocket Disconnections**: Automatic reconnection with visual feedback
2. **API Timeouts**: Graceful degradation with user notification  
3. **Token Validation**: Real-time validation with settings modal prompts
4. **Song Generation Failures**: Detailed error reporting with retry options

### Client-Side Error Recovery

```javascript
// Automatic reconnection on disconnect
this.websocket.onclose = () => {
    this.updateStatus('Disconnected. Reconnecting...');
    this.setInputEnabled(false);
    setTimeout(() => this.connect(), 3000);
};

// Graceful error handling
this.websocket.onerror = (error) => {
    console.error('WebSocket Error:', error); 
    this.updateStatus('Connection error. Please refresh.');
    this.setInputEnabled(false);
};
```

## Performance Considerations

### Client Performance

- **DOM Updates**: Targeted updates prevent full re-renders
- **Audio Management**: Players persist across clip updates  
- **Memory Leaks**: Event listeners properly cleaned up
- **Rendering**: Markdown parsing cached where possible

### Server Performance

- **Connection Pooling**: Individual `ConnectionManager` per WebSocket
- **File I/O**: Async file operations with `aiofiles`
- **API Rate Limiting**: Built-in backoff for external API calls
- **Memory Usage**: Streaming responses prevent large memory accumulation

## Security Considerations

### Token Management

- **Client Storage**: Suno tokens stored in localStorage (consider more secure alternatives)
- **Server Security**: Tokens passed through WebSocket (consider authentication layers)
- **API Keys**: OpenAI keys stored server-side in environment variables

### Data Privacy

- **Chat Persistence**: All conversations stored locally in JSONL files
- **Audio URLs**: Direct links to Suno-hosted content (temporary URLs)
- **User Identification**: UUID-based session management without personal data

## Deployment Architecture

### Recommended Production Setup

```yaml
# docker-compose.yml example
version: '3.8'
services:
  orpheus-api:
    build: .
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - ROOT_PATH=/api
    volumes:
      - ./chats_ws:/app/chats_ws
    ports:
      - "8000:8000"
      
  nginx:
    image: nginx:alpine
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf
      - ./app/static:/usr/share/nginx/html
```

### WebSocket Proxy Configuration

```nginx
# nginx.conf WebSocket proxy
location /ws/ {
    proxy_pass http://orpheus-api:8000;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_set_header Host $host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_read_timeout 86400;
}
```

## Future Technical Improvements

### Scalability Enhancements

1. **Redis Integration**: Move from JSONL to Redis for message persistence
2. **Horizontal Scaling**: Support multiple server instances with shared state
3. **Connection Pooling**: Implement connection limits and queuing
4. **CDN Integration**: Cache static assets and audio files

### Feature Extensions  

1. **Message Reactions**: Add reaction/rating system for generated songs
2. **Collaboration**: Multi-user chat sessions with shared song generation
3. **Version Control**: Track song iterations and allow rollbacks
4. **Export Options**: Download conversations and audio in various formats

### Security Hardening

1. **JWT Authentication**: Replace simple token passing with proper auth
2. **Rate Limiting**: Implement per-user API rate limits  
3. **Content Validation**: Sanitize and validate all user inputs
4. **Audit Logging**: Track all API calls and user actions

## Monitoring & Observability

### Recommended Metrics

- **WebSocket Connections**: Active connections, connection duration
- **Message Throughput**: Messages per second, streaming latency
- **Song Generation**: Success/failure rates, generation time
- **API Integration**: OpenAI/Suno response times, error rates

### Logging Strategy

```python
# Structured logging example
logger.info(f"Song generation started", extra={
    "chat_uuid": chat_uuid,
    "prompt_length": len(prompt),
    "generation_type": generation_type,
    "user_tier": metadata.get("user_tier")
})
```

## Conclusion

This WebSocket chat system demonstrates a sophisticated approach to real-time AI interaction with integrated media generation. The architecture prioritizes user experience through streaming responses, progressive enhancement, and resilient connection management while maintaining clean separation between conversation flow and media generation workflows.

The modular design enables easy extension for additional AI tools and media types, making it suitable as a foundation for more complex creative collaboration platforms.