insights/products/ai_observability/frontend/llmPersonsLazyLoaderLogic.ts

import { MakeLogicType, actions, connect, kea, listeners, path, reducers, selectors } from 'kea'

import api from 'lib/api'
import { teamLogic } from 'scenes/teamLogic'

import { LLMTracePerson } from '~/queries/schema/schema-general'

interface PersonAPIResponse {
    uuid: string
    name: string
    distinct_ids: string[]
    properties: Record<string, unknown>
    created_at: string
}

interface BatchByDistinctIdsResponse {
    results: Record<string, PersonAPIResponse>
}

function toTracePerson(apiPerson: PersonAPIResponse, distinctId: string): LLMTracePerson {
    return {
        uuid: apiPerson.uuid,
        created_at: apiPerson.created_at,
        properties: apiPerson.properties,
        distinct_id: distinctId,
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface llmPersonsLazyLoaderLogicValues {
    currentTeamId: number | null // teamLogic
    isDistinctIdLoading: (distinctId: string) => boolean
    loadingDistinctIds: Set<string>
    personsCache: Record<string, LLMTracePerson | null>
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface llmPersonsLazyLoaderLogicActions {
    ensurePersonLoaded: (distinctId: string) => {
        distinctId: string
    }
    loadPersonsBatchFailure: (requestedDistinctIds: string[]) => {
        requestedDistinctIds: string[]
    }
    loadPersonsBatchSuccess: (
        persons: Record<string, LLMTracePerson>,
        requestedDistinctIds: string[]
    ) => {
        persons: Record<string, LLMTracePerson>
        requestedDistinctIds: string[]
    }
    queuePersonLoad: (distinctId: string) => {
        distinctId: string
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface llmPersonsLazyLoaderLogicMeta {
    __keaTypeGenInternalSelectorTypes: {
        isDistinctIdLoading: (loadingDistinctIds: Set<string>) => (distinctId: string) => boolean
    }
}

export type llmPersonsLazyLoaderLogicType = MakeLogicType<
    llmPersonsLazyLoaderLogicValues,
    llmPersonsLazyLoaderLogicActions,
    Record<string, any>,
    llmPersonsLazyLoaderLogicMeta
>

export const llmPersonsLazyLoaderLogic = kea<llmPersonsLazyLoaderLogicType>([
    path(['products', 'ai_observability', 'frontend', 'llmPersonsLazyLoaderLogic']),

    connect({
        values: [teamLogic, ['currentTeamId']],
    }),

    actions({
        ensurePersonLoaded: (distinctId: string) => ({ distinctId }),
        queuePersonLoad: (distinctId: string) => ({ distinctId }),
        loadPersonsBatchSuccess: (persons: Record<string, LLMTracePerson>, requestedDistinctIds: string[]) => ({
            persons,
            requestedDistinctIds,
        }),
        loadPersonsBatchFailure: (requestedDistinctIds: string[]) => ({ requestedDistinctIds }),
    }),

    reducers({
        personsCache: [
            {} as Record<string, LLMTracePerson | null>,
            {
                loadPersonsBatchSuccess: (state, { persons, requestedDistinctIds }) => {
                    const newState = { ...state }

                    for (const distinctId of requestedDistinctIds) {
                        newState[distinctId] = persons[distinctId] ?? null
                    }

                    return newState
                },
                loadPersonsBatchFailure: (state, { requestedDistinctIds }) => {
                    const newState = { ...state }

                    for (const distinctId of requestedDistinctIds) {
                        newState[distinctId] = null
                    }

                    return newState
                },
            },
        ],

        loadingDistinctIds: [
            new Set<string>(),
            {
                queuePersonLoad: (state, { distinctId }) => {
                    if (state.has(distinctId)) {
                        return state
                    }

                    const newSet = new Set(state)
                    newSet.add(distinctId)
                    return newSet
                },
                loadPersonsBatchSuccess: (state, { requestedDistinctIds }) => {
                    const newSet = new Set(state)

                    for (const id of requestedDistinctIds) {
                        newSet.delete(id)
                    }

                    return newSet
                },
                loadPersonsBatchFailure: (state, { requestedDistinctIds }) => {
                    const newSet = new Set(state)

                    for (const id of requestedDistinctIds) {
                        newSet.delete(id)
                    }

                    return newSet
                },
            },
        ],
    }),

    selectors({
        isDistinctIdLoading: [
            (s) => [s.loadingDistinctIds],
            (loadingDistinctIds: Set<string>): ((distinctId: string) => boolean) => {
                return (distinctId: string) => loadingDistinctIds.has(distinctId)
            },
        ],
    }),

    listeners(({ values, actions }) => {
        let pendingDistinctIds = new Set<string>()
        let batchTimer: ReturnType<typeof setTimeout> | null = null
        const scheduleBatchLoad = (): void => {
            if (batchTimer) {
                return
            }

            batchTimer = setTimeout(async () => {
                const batch = Array.from(pendingDistinctIds)
                pendingDistinctIds = new Set()
                batchTimer = null

                if (batch.length === 0) {
                    return
                }

                const teamId = values.currentTeamId

                if (!teamId) {
                    actions.loadPersonsBatchFailure(batch)
                    return
                }

                try {
                    // nosemgrep: prefer-codegen-api
                    const response = await api.create<BatchByDistinctIdsResponse>(
                        `v1/environments/${teamId}/persons/batch_by_distinct_ids/`,
                        { distinct_ids: batch }
                    )

                    const persons: Record<string, LLMTracePerson> = {}

                    for (const [distinctId, personData] of Object.entries(response.results)) {
                        persons[distinctId] = toTracePerson(personData, distinctId)
                    }

                    actions.loadPersonsBatchSuccess(persons, batch)
                } catch {
                    actions.loadPersonsBatchFailure(batch)
                }
            }, 0)
        }

        return {
            ensurePersonLoaded: ({ distinctId }) => {
                if (values.personsCache[distinctId] !== undefined) {
                    return
                }

                actions.queuePersonLoad(distinctId)
                pendingDistinctIds.add(distinctId)
                scheduleBatchLoad()
            },
        }
    }),
])