discovery-by-llm/frontend/src/composables/useChatStream.ts

import { computed, onUnmounted, ref } from 'vue'
import {
  EVENT_LOG_MAX_ENTRIES,
  HISTORY_CONTENT_MAX_LENGTH,
  HISTORY_MAX_TURNS,
  PLAYLIST_NAME_MAX_LENGTH,
  PRIOR_RECOMMENDATIONS_MAX_TRACKS,
  QUERY_MAX_LENGTH,
} from '../lib/constants'
import { messages } from '../lib/messages'
import type { AssistantTurn, ChatTurn, EventLogEntry, TransportFailure } from '../lib/models'
import { EMPTY_PLAYLIST_STATE } from '../lib/models'
import {
  isAbortError,
  streamRecommendations,
  StreamTransportError,
} from '../lib/recommendationStream'
import type {
  HistoryTurn,
  PlaylistCreateRequest,
  PriorRecommendation,
  RecommendationRequest,
  StreamEvent,
} from '../lib/types'

type PlaylistCreator = (request: PlaylistCreateRequest) => Promise<{ url: string | null }>

interface ActiveRequest {
  controller: AbortController
  turnId: string
  terminalReceived: boolean
}

function assistantContent(turn: AssistantTurn): string {
  const trackLines = turn.tracks.map(
    (event) => `${event.rank}. ${event.track.title} by ${event.track.artists.join(', ')}`,
  )
  const content = [turn.intentSummary, ...trackLines].filter(Boolean).join('\n')
  return content.slice(0, HISTORY_CONTENT_MAX_LENGTH)
}

function buildHistory(turns: ChatTurn[]): HistoryTurn[] {
  return turns
    .map((turn): HistoryTurn | null => {
      if (turn.role === 'user') return { role: 'user', content: turn.text }
      const content = assistantContent(turn)
      return content ? { role: 'assistant', content } : null
    })
    .filter((turn): turn is HistoryTurn => turn !== null)
    .slice(-HISTORY_MAX_TURNS)
}

function findLatestCompletedAssistant(turns: ChatTurn[]): AssistantTurn | undefined {
  for (let index = turns.length - 1; index >= 0; index -= 1) {
    const turn = turns[index]
    if (turn?.role === 'assistant' && turn.status === 'done') return turn
  }
  return undefined
}

function buildPriorRecommendations(turns: ChatTurn[]): PriorRecommendation[] {
  const latest = findLatestCompletedAssistant(turns)
  if (!latest) return []
  return latest.tracks.slice(0, PRIOR_RECOMMENDATIONS_MAX_TRACKS).map((event) => ({
    rank: event.rank,
    track_id: event.track.id,
    title: event.track.title,
    artists: event.track.artists,
  }))
}

function reduceEvent(turn: AssistantTurn, event: StreamEvent): AssistantTurn {
  switch (event.type) {
    case 'metadata':
      return {
        ...turn,
        requestId: event.request_id,
        intentSummary: event.intent_summary,
        candidateCount: event.candidate_count,
      }
    case 'track':
      return { ...turn, tracks: [...turn.tracks, event] }
    case 'warning':
      return { ...turn, warnings: [...turn.warnings, event] }
    case 'error':
      return { ...turn, status: 'error', error: event }
    case 'done':
      if (event.track_count !== turn.tracks.length) {
        return {
          ...turn,
          status: 'error',
          completion: event,
          transportFailure: {
            kind: 'protocol',
            message: messages.useChatStreamTrackCountMismatch,
          },
        }
      }
      return { ...turn, status: 'done', completion: event }
  }
}

function playlistName(query: string): string {
  return query.slice(0, PLAYLIST_NAME_MAX_LENGTH).trim()
}

function createTurnId(): string {
  if (typeof crypto.randomUUID === 'function') return crypto.randomUUID()
  return `${Date.now()}-${Math.random().toString(36).slice(2)}`
}

function findLatestAssistant(turns: ChatTurn[]): AssistantTurn | undefined {
  for (let index = turns.length - 1; index >= 0; index -= 1) {
    const turn = turns[index]
    if (turn?.role === 'assistant') return turn
  }
  return undefined
}

/** Reduce recommendation streams into conversation view state. */
export function useChatStream(createPlaylist: PlaylistCreator) {
  const turns = ref<ChatTurn[]>([])
  const eventLog = ref<EventLogEntry[]>([])
  let activeRequest: ActiveRequest | null = null

  const isStreaming = computed(() =>
    turns.value.some((turn) => turn.role === 'assistant' && turn.status === 'streaming'),
  )
  const turnCount = computed(() => turns.value.filter((turn) => turn.role === 'user').length)
  const latestAssistant = computed(() => findLatestAssistant(turns.value))

  function updateAssistant(id: string, update: (turn: AssistantTurn) => AssistantTurn): void {
    turns.value = turns.value.map((turn) => {
      if (turn.role === 'assistant' && turn.id === id) return update(turn)
      return turn
    })
  }

  function addLog(type: string, detail: string): void {
    eventLog.value = [
      ...eventLog.value,
      { timestamp: new Date().toLocaleTimeString(), type, detail },
    ].slice(-EVENT_LOG_MAX_ENTRIES)
  }

  function cancelActive(shouldMarkTurn: boolean): void {
    const request = activeRequest
    if (!request) return
    request.controller.abort()
    if (shouldMarkTurn && !request.terminalReceived) {
      const failure: TransportFailure = {
        kind: 'cancelled',
        message: messages.useChatStreamRequestReplaced,
      }
      updateAssistant(request.turnId, (turn) => {
        if (turn.status !== 'streaming') return turn
        return { ...turn, status: 'error', transportFailure: failure }
      })
    }
    activeRequest = null
  }

  async function send(query: string): Promise<void> {
    const text = query.trim().slice(0, QUERY_MAX_LENGTH)
    if (!text) return
    cancelActive(true)

    const request: RecommendationRequest = {
      schema_version: 1,
      query: text,
      history: buildHistory(turns.value),
      prior_recommendations: buildPriorRecommendations(turns.value),
    }
    const userTurn: ChatTurn = { id: createTurnId(), role: 'user', text }
    const assistantTurn: AssistantTurn = {
      id: createTurnId(),
      role: 'assistant',
      query: text,
      status: 'streaming',
      requestId: null,
      intentSummary: '',
      candidateCount: null,
      tracks: [],
      warnings: [],
      error: null,
      transportFailure: null,
      completion: null,
      playlist: { ...EMPTY_PLAYLIST_STATE },
    }
    turns.value = [...turns.value, userTurn, assistantTurn]

    const controller = new AbortController()
    const requestState: ActiveRequest = {
      controller,
      turnId: assistantTurn.id,
      terminalReceived: false,
    }
    activeRequest = requestState

    try {
      await streamRecommendations(request, controller.signal, (event) => {
        if (activeRequest !== requestState || requestState.terminalReceived) return
        if (event.type === 'done' || event.type === 'error') {
          requestState.terminalReceived = true
        }
        updateAssistant(assistantTurn.id, (turn) => reduceEvent(turn, event))
        addLog(event.type, JSON.stringify(event))
      })
    } catch (error) {
      if (activeRequest !== requestState || requestState.terminalReceived) return
      if (isAbortError(error, controller.signal)) return
      let failure: TransportFailure
      if (error instanceof StreamTransportError) {
        failure = { kind: error.kind, message: error.message }
      } else {
        failure = { kind: 'network', message: messages.useChatStreamUnexpectedFailure }
      }
      updateAssistant(assistantTurn.id, (turn) => ({
        ...turn,
        status: 'error',
        transportFailure: failure,
      }))
      addLog(messages.useChatStreamTransportEvent, `${failure.kind}: ${failure.message}`)
    } finally {
      if (activeRequest === requestState) {
        activeRequest = null
      }
    }
  }

  async function savePlaylist(turnId: string): Promise<void> {
    const turn = turns.value.find(
      (candidate): candidate is AssistantTurn =>
        candidate.role === 'assistant' && candidate.id === turnId,
    )
    if (
      !turn ||
      turn.status !== 'done' ||
      turn.tracks.length === 0 ||
      turn.playlist.status === 'saving' ||
      turn.playlist.status === 'saved'
    ) {
      return
    }

    const name = playlistName(turn.query)
    updateAssistant(turnId, (current) => ({
      ...current,
      playlist: { status: 'saving', name, url: null, message: null },
    }))

    try {
      const response = await createPlaylist({
        schema_version: 1,
        name,
        track_uris: turn.tracks.map((event) => event.track.uri),
      })
      updateAssistant(turnId, (current) => ({
        ...current,
        playlist: { status: 'saved', name, url: response.url, message: null },
      }))
    } catch {
      updateAssistant(turnId, (current) => ({
        ...current,
        playlist: {
          status: 'error',
          name,
          url: null,
          message: messages.useChatStreamPlaylistFailure,
        },
      }))
    }
  }

  function reset(): void {
    cancelActive(false)
    turns.value = []
    eventLog.value = []
  }

  onUnmounted(() => cancelActive(false))

  return {
    turns,
    eventLog,
    isStreaming,
    turnCount,
    latestAssistant,
    send,
    savePlaylist,
    reset,
  }
}