dify/web/features/agent-v2/agent-composer/store-modules/knowledge.ts

import type { AgentKnowledgeRetrievalItem } from '../form-state'
import type { DraftFieldUpdate } from './utils'
import { atom } from 'jotai'
import { syncKnowledgeReferenceLabels } from '../reference-labels'
import { agentComposerDraftAtom } from '../store'
import { resolveDraftFieldUpdate } from './utils'

export const agentComposerKnowledgeRetrievalsAtom = atom(
  get => get(agentComposerDraftAtom).knowledgeRetrievals,
  (get, set, knowledgeRetrievalsUpdate: DraftFieldUpdate<AgentKnowledgeRetrievalItem[]>) => {
    const draft = get(agentComposerDraftAtom)
    const knowledgeRetrievals = resolveDraftFieldUpdate(draft.knowledgeRetrievals, knowledgeRetrievalsUpdate)

    set(agentComposerDraftAtom, {
      ...draft,
      prompt: syncKnowledgeReferenceLabels({
        prompt: draft.prompt,
        currentRetrievals: draft.knowledgeRetrievals,
        nextRetrievals: knowledgeRetrievals,
      }),
      knowledgeRetrievals,
    })
  },
)

export const addKnowledgeRetrievalAtom = atom(null, (_get, set, retrieval: AgentKnowledgeRetrievalItem) => {
  set(agentComposerKnowledgeRetrievalsAtom, retrievals => [...retrievals, retrieval])
})

export const updateKnowledgeRetrievalAtom = atom(null, (_get, set, retrieval: AgentKnowledgeRetrievalItem) => {
  set(agentComposerKnowledgeRetrievalsAtom, retrievals => retrievals.map(currentRetrieval => (
    currentRetrieval.id === retrieval.id ? retrieval : currentRetrieval
  )))
})

export const removeKnowledgeRetrievalAtom = atom(null, (_get, set, retrievalId: string) => {
  set(agentComposerKnowledgeRetrievalsAtom, retrievals => retrievals.filter(retrieval => retrieval.id !== retrievalId))
})