get-matching-context.ts

import type { ContextItem, MentionQuery } from '@sourcegraph/cody-shared'

import { DEFAULT_FAST_MODEL_CHARS_LIMIT } from '@sourcegraph/cody-shared/src/prompt/constants'
import { getFileContextFiles, getSymbolContextFiles } from '../../editor/utils/editor-context'
import { getLabelForContextItem } from './utils'

const MAX_FUZZY_RESULTS = 20

interface FixupMatchingContext {
    /* Unique identifier for the context, shown in the input value but not necessarily in the quick pick selector */
    key: string
    /* If present, will override the key shown in the quick pick selector */
    shortLabel?: string
    file: ContextItem
}

export async function getMatchingContext(
    mentionQuery: MentionQuery
): Promise<FixupMatchingContext[] | null> {
    if (mentionQuery.type === 'symbol') {
        const symbolResults = await getSymbolContextFiles(mentionQuery.text, MAX_FUZZY_RESULTS)
        return symbolResults.map(result => ({
            key: getLabelForContextItem(result),
            file: result,
            shortLabel: `${result.kind === 'class' ? '$(symbol-structure)' : '$(symbol-method)'} ${
                result.symbolName
            }`,
        }))
    }

    if (mentionQuery.type === 'file') {
        const fileResults = await getFileContextFiles(
            mentionQuery.text,
            MAX_FUZZY_RESULTS,
            DEFAULT_FAST_MODEL_CHARS_LIMIT
        )
        return fileResults.map(result => ({
            key: getLabelForContextItem(result),
            file: result,
        }))
    }

    return null
}