Function registerQueryExpansionFunction

src/functions/query-expansion.ts:69–130

(
  sdk: ISdk,
  provider: MemoryProvider,
)

}

export function registerQueryExpansionFunction(
  sdk: ISdk,
  provider: MemoryProvider,
): void {
  sdk.registerFunction("mem::expand-query", 
    async (data: { query: string; maxReformulations?: number } | undefined) => {
      if (!data || typeof data.query !== "string" || !data.query.trim()) {
        logger.warn("Invalid expand-query payload");
        return { success: false, error: "query must be a non-empty string" };
      }
      const rawMaxR = Number(data.maxReformulations);
      const maxR = Number.isFinite(rawMaxR)
        ? Math.max(1, Math.min(10, Math.floor(rawMaxR)))
        : 5;
      const query = data.query.trim();

      try {
        const response = await provider.compress(
          QUERY_EXPANSION_SYSTEM,
          `Expand this query for memory retrieval:\n\n"${query}"`,
        );

        const parsed = parseExpansionXml(response);
        if (!parsed) {
          logger.warn("Failed to parse query expansion");
          return {
            success: true,
            expansion: {
              original: query,
              reformulations: [],
              temporalConcretizations: [],
              entityExtractions: [],
            },
          };
        }

        parsed.original = query;
        parsed.reformulations = parsed.reformulations.slice(0, maxR);

        logger.info("Query expanded", {
          original: query,
          reformulations: parsed.reformulations.length,
          entities: parsed.entityExtractions.length,
        });

        return { success: true, expansion: parsed };
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        logger.error("Query expansion failed", { error: msg });
        return {
          success: true,
          expansion: {
            original: query,
            reformulations: [],
            temporalConcretizations: [],
            entityExtractions: [],
          },
        };