# Stems and MIDI Architecture

## Overview

This document outlines the architecture for adding stem separation and MIDI generation capabilities to Suno Spaces. Users can request stems (isolated instrument tracks) and MIDI data for existing songs through natural language commands to the assistant.

## User Experience

### Command Examples

All commands must include an atom reference (attached song):

```
@suno get me all the stems for this song
@suno get me the drum stem and midi
@suno get me the vocal stem and midi
@suno get me the instrumental stem
```

### Error Handling

If a user requests stems/MIDI without attaching a song atom, the assistant will respond with an error message asking them to attach the song they want to process.

### Future Vision

These commands will be chained together for higher-level features:
```
@suno make me a karaoke game for this song
```
This would automatically:
1. Generate vocal + instrumental stems
2. Generate MIDI for the vocals
3. Create a karaoke UI component

## Data Model

### Schema Changes

**Stems are songs**: A stem is a subtype of a song atom, not a separate type. We use metadata fields to distinguish stems from regular songs.

**Provenance System**: For tracking how atoms were created, what inputs were used, and relationships between atoms, see [PROVENANCE.md](./PROVENANCE.md). The provenance system handles:
- Multi-input operations (e.g., cover with custom lyrics + style image)
- Operation parameters
- Querying all derivatives/children of an atom
- Full lineage trees

Add new MIDI atom type:

```typescript
// In schema.ts
atoms: defineTable({
  type: v.union(
    v.literal("song"),
    v.literal("video"),
    v.literal("image"),
    v.literal("lyrics"),
    v.literal("webview"),
    v.literal("midi")     // NEW - MIDI is its own type
  ),
  // ... existing fields ...
})
```

All atom relationships and provenance are tracked in the `atomProvenance` table (see PROVENANCE.md).

### Stem Song Metadata

Stems are `type: "song"` atoms with additional metadata fields:

```typescript
{
  type: "song",  // Stems are songs!
  // Relationship to source song tracked in atomProvenance table (see PROVENANCE.md)
  metadata: {
    sunoClipId: "<suno_stem_clip_id>",

    // Stem-specific fields
    isStem: true,  // Flag to identify this as a stem
    stemType: "vocals" | "drums" | "bass" | "other" | "instrumental" | ...,
    stemTypeId: 91,  // Suno's internal stem type ID
    stemTypeGroupName: "Twelve",  // Suno's stem generation method
    stemTask: "twelve",  // Task identifier

    // Standard song fields
    title: "Song Title - Vocals",
    artist: "Artist Name",
    audioUrl: "<url_to_stem_audio>",
    videoUrl: "<url_to_stem_video>",
    albumArtUrl: "<url_to_album_art>",
    duration: 180,
    sunoStatus: "completed" | "processing" | "failed"
  },
  status: "pending" | "processing" | "streaming" | "completed" | "failed",
  progress: 0-100
}
```

**Why stems are songs:**
- Stems have audio/video URLs and can be played like songs
- Stems support all song operations (play, like, add to queue)
- Stems can be used as inputs to other operations (e.g., MIDI generation from a stem)
- Frontend can reuse existing song playback components
- The `isStem` flag and `stemType` allow UI differentiation when needed

### MIDI Atom Metadata

MIDI is its own atom type since it's not playable audio:

```typescript
{
  type: "midi",
  // Relationship to source song/stem tracked in atomProvenance table (see PROVENANCE.md)
  metadata: {
    sunoClipId: "<source_suno_clip_id>",  // The clip ID MIDI was generated from
    midiUrl: "<url_to_midi_file>",
    midiStatus: "pending" | "processing" | "completed" | "failed",
    // MIDI-specific data from Suno API response
    midiData: {
      // Raw MIDI data or relevant metadata from Suno
    }
  },
  status: "pending" | "processing" | "completed" | "failed",
  progress: 0-100
}
```

## Suno API Integration

### Stem Generation API

**Endpoint**: `/api/generate/v2-web`

**Minimal Request Body** (only non-null fields):
```json
{
  "task": "gen_stem",
  "mv": "chirp-v3-0",
  "make_instrumental": true,
  "continue_clip_id": "<suno_clip_id_of_original_song>",
  "stem_type_id": 91,
  "stem_type_group_name": "Twelve",
  "stem_task": "twelve"
}
```

**Key Parameters**:
- `task`: **Required** - Must be `"gen_stem"` for stem generation
- `mv`: **Required** - Model version, use `"chirp-v3-0"` for stems (NOTE: different from regular song generation)
- `make_instrumental`: **Required** - Set to `true` for stems
- `continue_clip_id`: **Required** - Suno clip ID of the original song to extract stems from
- `stem_type_id`: **Required** - Identifier for which stem to extract (see mapping below)
- `stem_type_group_name`: **Required** - Set to `"Twelve"` for the 12-stem extraction method
- `stem_task`: **Required** - Set to `"twelve"` to generate all 12 stems at once

**Response**: Returns clip IDs for the generated stems (similar to song generation)

**Important Notes**:
- The `mv` field for stems is `"chirp-v3-0"`, not `"chirp-crow"` used for regular song generation
- All null/empty fields from the full API example can be omitted
- The `metadata` object is not needed for stem generation

### MIDI Generation API

**Endpoint**: `/api/gen/{suno_id}/midi`

**Method**: GET (polling endpoint)

**Response**:
```json
{
  "status": "pending" | "processing" | "completed" | "failed",
  "midi_url": "<url_to_midi_file>",  // Only present when completed
  // Additional MIDI metadata...
}
```

**Polling Strategy**: Poll every 5 seconds until status is "completed" or "failed"

### Stem Type Mapping

Need to determine the mapping between natural language stem names and Suno's `stem_type_id`:

```typescript
const STEM_TYPE_MAP = {
  "all": null,  // Special case: generate all stems using stem_task="twelve"
  "vocals": <id>,
  "instrumental": <id>,
  "drums": <id>,
  "bass": <id>,
  "guitar": <id>,
  "piano": <id>,
  "strings": <id>,
  // ... other stem types
}
```

## Implementation Components

### 1. SunoClient Extensions (lib/sunoClient.ts)

Add new methods to the `SunoClient` class:

```typescript
// Generate stems for a song
async generateStems(
  clipId: string,
  stemTypes: string[] | "all"
): Promise<string[]> {
  // Request body for stems:
  // {
  //   "task": "gen_stem",
  //   "mv": "chirp-v3-0",  // Important: different from regular songs!
  //   "make_instrumental": true,
  //   "continue_clip_id": clipId,
  //   "stem_type_id": 91,  // ID for the stem type (see STEM_TYPE_MAP)
  //   "stem_type_group_name": "Twelve",
  //   "stem_task": "twelve"
  // }
  //
  // If "all", use stem_task="twelve" to get all 12 stems at once
  // Otherwise, make individual requests for each stem type with appropriate stem_type_id
  // Returns array of stem clip IDs
}

// Poll MIDI generation status
async getMidiStatus(clipId: string): Promise<{
  status: "pending" | "processing" | "completed" | "failed",
  midiUrl?: string,
  midiData?: any
}> {
  // Poll /api/gen/{clipId}/midi endpoint
}
```

### 2. New Convex Actions (convex/atoms.ts)

```typescript
// Generate stems for a source song
export const generateStems = internalAction({
  args: {
    sourceAtomId: v.id("atoms"),
    stemTypes: v.array(v.string()),  // e.g., ["vocals", "drums"] or ["all"]
    spaceId: v.id("spaces"),
    userId: v.id("users"),
  },
  handler: async (ctx, args) => {
    // 1. Get source atom and extract Suno clip ID
    // 2. Call SunoClient.generateStems()
    // 3. Create stem atoms for each requested stem (type: "song", isStem: true)
    // 4. Create provenance records (see PROVENANCE.md)
    //    - operationType: "stem"
    //    - inputs: [{ atomId: sourceAtomId, role: "source" }]
    //    - parameters: { stemType, stemTypeId }
    // 5. Schedule polling for stem completion
    // 6. Return stem atom IDs
  }
});

// Poll stem generation status
export const pollStemStatus = internalAction({
  args: {
    stemAtomIds: v.array(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    // Similar to pollSongStatus but for stems
    // Poll Suno API for stem completion
    // Update stem atoms with audio URLs and status
    // Reschedule if not complete (2s interval)
  }
});

// Generate MIDI for a song or stem
export const generateMidi = internalAction({
  args: {
    sourceAtomId: v.id("atoms"),  // Can be song or stem atom
    spaceId: v.id("spaces"),
    userId: v.id("users"),
  },
  handler: async (ctx, args) => {
    // 1. Get source atom and extract Suno clip ID
    // 2. Create MIDI atom in "pending" state
    // 3. Create provenance record (see PROVENANCE.md)
    //    - operationType: "midi"
    //    - inputs: [{ atomId: sourceAtomId, role: "source" }]
    //    - parameters: {}
    // 4. Schedule MIDI polling (start immediately)
    // 5. Return MIDI atom ID
  }
});

// Poll MIDI generation status
export const pollMidiStatus = internalAction({
  args: {
    midiAtomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    // 1. Get MIDI atom
    // 2. Call SunoClient.getMidiStatus()
    // 3. Update MIDI atom with status and URL when complete
    // 4. Reschedule if not complete (5s interval - less frequent than stems)
  }
});
```

### 3. Query Functions (convex/atoms.ts)

```typescript
// Get all stems for a source song
// Uses provenance system to find children with operationType: "stem"
export const getStemsBySource = query({
  args: {
    sourceAtomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    // Get all provenance records where this atom is an input
    const allProvenance = await ctx.db.query("atomProvenance").collect();

    const stemProvenance = allProvenance.filter(prov =>
      prov.operationType === "stem" &&
      prov.inputs.some(input => input.atomId === args.sourceAtomId)
    );

    // Load the stem atoms
    const stems = await Promise.all(
      stemProvenance.map(async (prov) => {
        const atom = await ctx.db.get(prov.outputAtomId);
        return atom;
      })
    );

    return stems.filter(atom => atom !== null);
  }
});

// Get MIDI for a source song or stem
// Uses provenance system to find children with operationType: "midi"
export const getMidiBySource = query({
  args: {
    sourceAtomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    const allProvenance = await ctx.db.query("atomProvenance").collect();

    const midiProvenance = allProvenance.find(prov =>
      prov.operationType === "midi" &&
      prov.inputs.some(input => input.atomId === args.sourceAtomId)
    );

    if (!midiProvenance) return null;

    return await ctx.db.get(midiProvenance.outputAtomId);
  }
});
```

### 4. Assistant Integration (convex/assistant.ts)

Add new OpenAI function tools:

```typescript
const GENERATE_STEMS_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_stems",
    description: "Generate stem tracks (isolated instruments/vocals) from a song. Requires a source song atom ID.",
    parameters: {
      type: "object",
      properties: {
        sourceAtomId: {
          type: "string",
          description: "The atom ID of the source song to generate stems from. Look for atom IDs in the format '[Referenced atom IDs: ...]' in the user message."
        },
        stemTypes: {
          type: "array",
          items: { type: "string" },
          description: "Array of stem types to generate: ['vocals', 'instrumental', 'drums', 'bass', 'guitar', 'piano', 'strings', 'other'] or ['all'] to generate all stems."
        }
      },
      required: ["sourceAtomId", "stemTypes"]
    }
  }
};

const GENERATE_MIDI_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_midi",
    description: "Generate MIDI data from a song or stem. Requires a source atom ID.",
    parameters: {
      type: "object",
      properties: {
        sourceAtomId: {
          type: "string",
          description: "The atom ID of the source song or stem to generate MIDI from. Look for atom IDs in the format '[Referenced atom IDs: ...]' in the user message."
        }
      },
      required: ["sourceAtomId"]
    }
  }
};
```

Update the system prompt to include stem/MIDI instructions:

```typescript
const ASSISTANT_SYSTEM_PROMPT = `...

For stem generation requests:
- Parse which stems the user wants (vocals, instrumental, drums, bass, etc.)
- If user says "all stems", use stemTypes: ["all"]
- Validate that a song atom is referenced

For MIDI generation requests:
- Validate that a song or stem atom is referenced
- MIDI can be generated from either songs or stems

Examples:
- "get me all the stems for this song" → generate_stems with stemTypes: ["all"]
- "get me the drum stem and midi" → generate_stems with ["drums"], then generate_midi
- "get me the vocal stem and midi" → generate_stems with ["vocals"], then generate_midi
`;
```

Update `processAssistantRequest` to handle new tool types:

```typescript
export const processAssistantRequest = internalAction({
  // ... existing args ...
  handler: async (ctx, args) => {
    try {
      // Validate atom references for stem/MIDI requests
      if (!args.atomReferences || args.atomReferences.length === 0) {
        const needsAtom = args.content.toLowerCase().includes("stem") ||
                         args.content.toLowerCase().includes("midi");
        if (needsAtom) {
          await ctx.runMutation(internal.assistant.updateAssistantMessage, {
            messageId: args.responseMessageId,
            content: "⚠️ Please attach a song to generate stems or MIDI. You can attach a song by referencing it in your message.",
            assistantStatus: "failed",
          });
          return;
        }
      }

      const result = await ctx.runAction(internal.assistant.classifyIntent, {
        messageContent: args.content,
        atomReferences: args.atomReferences,
      });

      // ... existing generate_song and generate_image handling ...

      if (result.toolName === "generate_stems") {
        const params = result.params as {
          sourceAtomId: string;
          stemTypes: string[];
        };

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Generating ${params.stemTypes.join(", ")} stem(s)...`,
          assistantStatus: "processing",
        });

        const stemAtomIds = await ctx.runAction(internal.atoms.generateStems, {
          sourceAtomId: params.sourceAtomId as Id<"atoms">,
          stemTypes: params.stemTypes,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          atomReferences: stemAtomIds,
          assistantStatus: "completed",
        });
      }

      if (result.toolName === "generate_midi") {
        const params = result.params as {
          sourceAtomId: string;
        };

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎹 Generating MIDI...`,
          assistantStatus: "processing",
        });

        const midiAtomId = await ctx.runAction(internal.atoms.generateMidi, {
          sourceAtomId: params.sourceAtomId as Id<"atoms">,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          atomReferences: [midiAtomId],
          assistantStatus: "completed",
        });
      }

      // Handle combined stem + MIDI requests
      // OpenAI may return multiple tool calls for "get me the drum stem and midi"
      // Process them sequentially: stem first, then MIDI from the stem

    } catch (error) {
      // ... error handling ...
    }
  }
});
```

### 5. Frontend Components

#### Stem Song Rendering

Since stems are song atoms, they use the existing `SongAtom` component with conditional rendering:

```typescript
// In components/atoms/SongAtom.tsx (existing component)
// Add conditional rendering based on metadata.isStem

if (atom.metadata?.isStem) {
  // Display stem-specific UI:
  // - Stem type badge (vocals, drums, etc.)
  // - "Stem from: [Original Song Title]" link
  // - All standard playback controls (reused!)
  // - Add to queue, like, play count (all work automatically)
} else {
  // Standard song UI
}
```

**Benefits of reusing SongAtom:**
- All playback logic works immediately
- Stems can be added to room queues
- Like/play count tracking works automatically
- No duplicate code

#### MIDI Atom Rendering

```typescript
// In components/atoms/MidiAtom.tsx (new component)
- Download button for MIDI file
- Link to source song/stem
- Show source song/stem title and cover art
- Future: Inline MIDI visualizer
```

## Polling Strategy

### Stem Polling
- **Interval**: 2 seconds (same as song polling)
- **Max Duration**: ~5 minutes
- **Reschedule**: Continue polling until all stems are completed/failed

### MIDI Polling
- **Interval**: 5 seconds (less frequent than stems)
- **Max Duration**: ~5 minutes
- **Reschedule**: Continue polling until MIDI is completed/failed

## Edge Cases & Error Handling

1. **No atom reference**: Respond with error asking user to attach a song
2. **Invalid atom type**: Validate that referenced atom is a song (for stems) or song/stem (for MIDI)
3. **Stem type not found**: If user requests unsupported stem type, suggest valid options
4. **Generation failure**: Update atom status to "failed" and show user-friendly error message
5. **Suno API errors**: Retry with exponential backoff, max 3 retries
6. **Quota limits**: Detect rate limit errors and inform user

## Future Enhancements

### Phase 2: Chained Workflows
- "Make me a karaoke game" → stems + MIDI + UI component
- "Remix the drums from this song" → drum stem + remix generation

### Phase 3: MIDI Visualization
- Inline piano roll view in chat
- Interactive MIDI editor
- Synchronize MIDI visualization with audio playback

### Phase 4: Stem Manipulation
- Volume mixing between stems
- Apply effects to individual stems
- Create custom mixes by selecting which stems to play

## Implementation Checklist

### Backend - Schema
- [ ] Add `midi` atom type to schema
- [ ] Implement provenance system from PROVENANCE.md:
  - [ ] Add `atomProvenance` table to schema
  - [ ] Add `createProvenance` internal mutation
  - [ ] Add `getProvenanceByAtom` query
  - [ ] Add `getAtomChildren` query

### Backend - Stems & MIDI
- [ ] Add stem type ID mapping constants
- [ ] Extend SunoClient with `generateStems()` and `getMidiStatus()` methods
- [ ] Implement `generateStems` action (creates song atoms with `isStem: true`)
  - [ ] Create provenance records for each stem
- [ ] Implement `pollStemStatus` action (reuse song polling pattern)
- [ ] Implement `generateMidi` action
  - [ ] Create provenance record for MIDI
- [ ] Implement `pollMidiStatus` action (5s interval)
- [ ] Add `getStemsBySource` query function
- [ ] Add `getMidiBySource` query function

### Backend - Assistant
- [ ] Update assistant.ts with GENERATE_STEMS_TOOL
- [ ] Update assistant.ts with GENERATE_MIDI_TOOL
- [ ] Update ASSISTANT_SYSTEM_PROMPT with stem/MIDI instructions
- [ ] Update `processAssistantRequest` to handle stem generation
- [ ] Update `processAssistantRequest` to handle MIDI generation
- [ ] Add atom reference validation for stem/MIDI requests
- [ ] Run TypeScript checks: `npx tsc --noEmit --project convex/tsconfig.json`

### Frontend - Atoms
- [ ] Update SongAtom component with conditional rendering for `metadata.isStem`
- [ ] Add stem type badge UI to SongAtom
- [ ] Add "Stem from: [Original Song]" link in SongAtom (use provenance)
- [ ] Create new MidiAtom component
- [ ] Add MIDI download button
- [ ] Display source song/stem info in MidiAtom (use provenance)

### Frontend - Provenance (see PROVENANCE.md)
- [ ] Create `AtomProvenanceCard` component
- [ ] Create `AtomDerivativesCard` component
- [ ] Add provenance section to atom detail view
- [ ] Add derivatives section to atom detail view

### Testing
- [ ] Test end-to-end stem generation flow ("get me all stems")
- [ ] Test single stem generation ("get me the vocal stem")
- [ ] Test end-to-end MIDI generation flow
- [ ] Test combined stem + MIDI requests ("get me the drum stem and midi")
- [ ] Test error handling for missing atom references
- [ ] Test stem playback in room queue
- [ ] Test like/play count on stems