@struktoai/mirage-core/dist/core/chroma/search.js

A Unified Virtual Filesystem For AI Agents

// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
import { mountedPath, mountPrefixOf } from "../../utils/key_prefix.js";
import { scoreFromDistance } from "../../utils/score.js";
import { rstripSlash, stripSlash } from "../../utils/slash.js";
import { compareCodePoints } from "../../utils/sort.js";
import { floatOption, intOption, textOption, validateOptions } from "../../vfs/search.js";
import { scalarString } from "../slug_tree/rows.js";
import { searchScope, targetEntries, validateQuery } from "../slug_tree/search.js";
import { PAGE_CHUNK_BATCH_SIZE } from "./client.js";
import { CHROMA_TREE } from "./tree.js";
export async function searchSegments(accessor, query, paths, topK = 10, mountPrefix = '', threshold = 0) {
    validateQuery(query, topK);
    if (mountPrefix === '' && paths.length > 0) {
        mountPrefix =
            (paths[0] === undefined ? undefined : mountPrefixOf(paths[0].virtual, paths[0].vfsPath)) ?? '';
    }
    let scopedSlugs = null;
    let where;
    if (paths.length > 0) {
        scopedSlugs = new Set((await targetEntries(CHROMA_TREE, accessor, paths)).keys());
        if (scopedSlugs.size === 0)
            return [];
        where = {
            [accessor.config.slugField]: { $in: [...scopedSlugs].sort(compareCodePoints) },
        };
    }
    const collection = await accessor.getCollection();
    const response = await collection.query({
        queryTexts: [query],
        nResults: topK,
        include: ['documents', 'metadatas', 'distances'],
        ...(where !== undefined ? { where } : {}),
    });
    return renderRecords(response, accessor.config.slugField, mountPrefix, scopedSlugs, threshold);
}
function renderRecords(response, slugField, mountPrefix, scopedSlugs = null, threshold = 0) {
    const documents = firstResultList(response.documents);
    const metadatas = firstResultList(response.metadatas);
    const distances = firstResultList(response.distances);
    const contents = [];
    for (let i = 0; i < documents.length; i++) {
        const document = documents[i];
        const metadata = metadatas[i];
        if (metadata === null || typeof metadata !== 'object' || Array.isArray(metadata))
            continue;
        const slug = scalarString(metadata[slugField]);
        if (slug === null)
            continue;
        const slugValue = stripSlash(slug);
        if (scopedSlugs !== null && !scopedSlugs.has(slugValue))
            continue;
        const score = scoreFromDistance(distances[i]);
        if (threshold > 0 && Number(score) < threshold)
            continue;
        let path = '/' + slugValue;
        const prefix = rstripSlash(mountPrefix);
        if (prefix !== '')
            path = prefix + path;
        const content = typeof document === 'string' ? document : '';
        contents.push([path, `${path}:${score}\n${content}`]);
    }
    return contents;
}
function firstResultList(value) {
    if (!Array.isArray(value))
        return [];
    if (value.length > 0 && Array.isArray(value[0]))
        return value[0];
    return value;
}
export async function searchMany(accessor, paths, query) {
    validateOptions(query, ['top_k', 'method', 'threshold']);
    const topK = intOption(query, 'top_k', 10);
    const threshold = floatOption(query, 'threshold', 0);
    if (textOption(query, 'method', 'semantic') !== 'semantic') {
        throw new Error("only the 'semantic' method is supported");
    }
    validateQuery(query.query, topK);
    const [targets, prefix] = await searchScope(CHROMA_TREE, accessor, paths);
    return searchSegments(accessor, query.query, targets, topK, prefix, threshold);
}
/**
 * The pages with a chunk holding `text`.
 *
 * A page is its chunks joined by line breaks, so a line, and so any match
 * grep makes, lies within one chunk, and Chroma's `$contains` over the chunks
 * answers exactly. It matches case, so -i reads every page. Below the mount's
 * root only the walked pages are asked. Mirrors Python's `files_containing`.
 */
export async function filesContaining(accessor, text, under, ignoreCase) {
    const [first] = under;
    if (ignoreCase || first === undefined)
        return null;
    const field = accessor.config.slugField;
    let where;
    if (under.every((scope) => stripSlash(scope.mountPath) !== '')) {
        const walked = await targetEntries(CHROMA_TREE, accessor, under);
        if (walked.size === 0)
            return [];
        where = { [field]: { $in: [...walked.keys()].sort(compareCodePoints) } };
    }
    const collection = await accessor.getCollection();
    const slugs = new Set();
    for (let offset = 0;; offset += PAGE_CHUNK_BATCH_SIZE) {
        const result = await collection.get({
            ...(where !== undefined ? { where } : {}),
            whereDocument: { $contains: text },
            include: ['metadatas'],
            limit: PAGE_CHUNK_BATCH_SIZE,
            offset,
        });
        for (const metadata of result.metadatas) {
            const slug = scalarString(metadata?.[field]);
            if (slug !== null)
                slugs.add(stripSlash(slug));
        }
        if (result.ids.length < PAGE_CHUNK_BATCH_SIZE)
            break;
    }
    return [...slugs].sort(compareCodePoints).map((slug) => mountedPath(first, `/${slug}`));
}
//# sourceMappingURL=search.js.map