playchat/frontend/src/stores/knowledgeBase.ts

import { defineStore } from 'pinia'
import { computed, ref } from 'vue'
import {
  fetchKbEmbeddingStatus,
  fetchKbFile,
  fetchKbFiles,
  reindexKbEmbeddings,
  testKbEmbedding,
  type KbEmbeddingProbe,
  type KbEmbeddingStatus,
  type KbFileContent,
  type KbFileInfo,
} from '../api/client'

export interface KbGroup {
  /** ナレッジベース直下は空文字。 */
  directory: string
  files: KbFileInfo[]
}

export const useKnowledgeBaseStore = defineStore('knowledgeBase', () => {
  const files = ref<KbFileInfo[]>([])
  /** 開いているファイルの全文。一覧に戻ると null に戻す。 */
  const openFile = ref<KbFileContent | null>(null)
  const loading = ref(false)
  const error = ref<string | null>(null)
  /** タイトル・パス・プレビューに対する絞り込み文字列。 */
  const filter = ref('')
  /** 折りたたまれているフォルダ名。web_search のようにファイル数が多いフォルダを畳むために使う。 */
  const collapsed = ref<Set<string>>(new Set())
  /** 埋め込みモデルの状態と進捗。 */
  const embedding = ref<KbEmbeddingStatus | null>(null)
  /** 「接続テスト」の結果。 */
  const probe = ref<KbEmbeddingProbe | null>(null)
  const testing = ref(false)

  /** 埋め込み実行中だけ状態をポーリングするタイマー。 */
  let pollTimer: number | null = null
  const POLL_INTERVAL_MS = 1500

  const filteredFiles = computed(() => {
    const q = filter.value.trim().toLowerCase()
    if (!q) return files.value
    return files.value.filter((f) =>
      `${f.title} ${f.path} ${f.preview}`.toLowerCase().includes(q),
    )
  })

  /** ディレクトリごとにまとめる。一覧はサーバー側でパス順に並んでいるので、そのままの順序でグループ化できる。 */
  const groups = computed<KbGroup[]>(() => {
    const result: KbGroup[] = []
    for (const file of filteredFiles.value) {
      const last = result[result.length - 1]
      if (last && last.directory === file.directory) last.files.push(file)
      else result.push({ directory: file.directory, files: [file] })
    }
    return result
  })

  async function load(): Promise<void> {
    loading.value = true
    error.value = null
    try {
      files.value = await fetchKbFiles()
    } catch (e) {
      error.value = e instanceof Error ? e.message : 'ナレッジベースの取得に失敗しました'
    } finally {
      loading.value = false
    }
  }

  async function open(path: string): Promise<void> {
    loading.value = true
    error.value = null
    try {
      openFile.value = await fetchKbFile(path)
    } catch (e) {
      error.value = e instanceof Error ? e.message : 'ファイルの読み込みに失敗しました'
    } finally {
      loading.value = false
    }
  }

  function close(): void {
    openFile.value = null
  }

  /** 埋め込みの状態を取り直す。実行中ならポーリングを始め、終わっていれば止める。 */
  async function loadEmbedding(): Promise<void> {
    try {
      embedding.value = await fetchKbEmbeddingStatus()
    } catch (e) {
      error.value = e instanceof Error ? e.message : '埋め込み状態の取得に失敗しました'
      stopPolling()
      return
    }
    if (embedding.value.running && pollTimer === null) {
      pollTimer = window.setInterval(() => void loadEmbedding(), POLL_INTERVAL_MS)
    } else if (!embedding.value.running && pollTimer !== null) {
      stopPolling()
      // 埋め込みが終わったので、一覧の「埋め込み済み」表示も更新する
      await load()
    }
  }

  function stopPolling(): void {
    if (pollTimer !== null) {
      clearInterval(pollTimer)
      pollTimer = null
    }
  }

  /** 埋め込みモデルに実際にリクエストして、応答するか確かめる。 */
  async function testEmbedding(): Promise<void> {
    testing.value = true
    probe.value = null
    try {
      probe.value = await testKbEmbedding()
    } catch (e) {
      error.value = e instanceof Error ? e.message : '埋め込みテストに失敗しました'
    } finally {
      testing.value = false
    }
  }

  /** 全ファイルの再埋め込みをバックグラウンドで開始する。 */
  async function reindex(): Promise<void> {
    error.value = null
    try {
      embedding.value = await reindexKbEmbeddings()
      if (embedding.value.running && pollTimer === null) {
        pollTimer = window.setInterval(() => void loadEmbedding(), POLL_INTERVAL_MS)
      }
    } catch (e) {
      error.value = e instanceof Error ? e.message : '再埋め込みの開始に失敗しました'
    }
  }

  function toggleGroup(directory: string): void {
    const next = new Set(collapsed.value)
    if (next.has(directory)) next.delete(directory)
    else next.add(directory)
    collapsed.value = next
  }

  return {
    files,
    openFile,
    loading,
    error,
    filter,
    collapsed,
    embedding,
    probe,
    testing,
    filteredFiles,
    groups,
    load,
    open,
    close,
    toggleGroup,
    loadEmbedding,
    stopPolling,
    testEmbedding,
    reindex,
  }
})