semantic-search

/**
 * Semantic Search - Production Reference Agent
 *
 * Canon alignment: KB 105 (Agent SDK), KB 22 (HumanOS)
 * Demonstrates: ctx.llm.embed, ctx.memory.vector (persistent.setVector/search), ctx.files, ctx.provenance
 *
 * Real use case: Index documents into a vector store, then search by semantic
 * similarity. This is the foundation pattern for RAG (retrieval-augmented
 * generation). Useful for knowledge bases, support systems, and any
 * application where users search by meaning rather than keywords.
 */

import { handler, withProvenanceContext } from '@human/agent-sdk';
import type { ExecutionContext } from '@human/agent-sdk';

export const AGENT_ID = 'semantic-search';
export const VERSION = '1.0.0';
export const CAPABILITIES = ['research/semantic-search', 'research/index-documents'];

export interface SemanticSearchInput {
  /** Action: 'index' to store documents, 'search' to query */
  action: 'index' | 'search';
  /** Documents to index (for action='index') */
  documents?: Array<{ path: string; metadata?: Record<string, unknown> }>;
  /** Search query (for action='search') */
  query?: string;
  /** Number of results to return */
  top_k?: number;
  /** Minimum similarity threshold (0-1) */
  threshold?: number;
}

export interface SemanticSearchOutput {
  success: boolean;
  /** Number of documents indexed (for action='index') */
  indexed_count?: number;
  /** Search results (for action='search') */
  results?: Array<{
    key: string;
    score: number;
    metadata?: Record<string, unknown>;
  }>;
  provenance_id: string;
}

const execute = async (
  ctx: ExecutionContext,
  input: SemanticSearchInput
): Promise<SemanticSearchOutput> => {
  ctx.log.info('Semantic search agent invoked', { action: input.action });

  if (input.action === 'index') {
    return indexDocuments(ctx, input);
  } else {
    return searchDocuments(ctx, input);
  }
};

/**
 * Index documents: read files, generate embeddings, store in vector memory
 */
async function indexDocuments(
  ctx: ExecutionContext,
  input: SemanticSearchInput
): Promise<SemanticSearchOutput> {
  const documents = input.documents ?? [];
  let indexedCount = 0;

  for (const doc of documents) {
    // Read document content (ctx.files)
    const content = await ctx.files.readText(doc.path);

    // Generate embedding (ctx.llm.embed) - the key gap feature
    const embedding = await ctx.llm.embed(content);

    // Store in vector memory (ctx.memory.persistent.setVector) - the other key gap feature
    const vectorKey = `doc:${doc.path}`;
    await ctx.memory.persistent.setVector(vectorKey, embedding.vector, {
      path: doc.path,
      dimensions: embedding.dimensions,
      model: embedding.model,
      indexed_at: new Date().toISOString(),
      ...doc.metadata,
    });

    ctx.log.info('Indexed document', {
      path: doc.path,
      dimensions: embedding.dimensions,
      cost_usd: embedding.cost.usd,
    });

    indexedCount++;
  }

  const provenanceId = await ctx.provenance.log(
    withProvenanceContext(ctx, {
      action: 'semantic_search:indexed',
      status: 'success',
      input: { document_count: documents.length },
      output: { indexed_count: indexedCount },
    })
  );

  return {
    success: true,
    indexed_count: indexedCount,
    provenance_id: provenanceId,
  };
}

/**
 * Search documents: embed query, search vector memory, return ranked results
 */
async function searchDocuments(
  ctx: ExecutionContext,
  input: SemanticSearchInput
): Promise<SemanticSearchOutput> {
  const query = input.query ?? '';
  const topK = input.top_k ?? 5;
  const threshold = input.threshold ?? 0.5;

  if (!query) {
    ctx.log.warn('Empty search query');
    const provenanceId = await ctx.provenance.log(
      withProvenanceContext(ctx, {
        action: 'semantic_search:empty_query',
        status: 'success',
        input: { query: '' },
        output: { results: [] },
      })
    );
    return { success: true, results: [], provenance_id: provenanceId };
  }

  // Generate embedding for query (ctx.llm.embed)
  const queryEmbedding = await ctx.llm.embed(query);

  // Search vector memory (ctx.memory.persistent.search)
  const searchResults = await ctx.memory.persistent.search(
    queryEmbedding.vector,
    {
      topK,
      threshold,
      includeMetadata: true,
    }
  );

  const results = searchResults.map((r) => ({
    key: r.key,
    score: r.score,
    metadata: r.metadata,
  }));

  ctx.log.info('Search completed', {
    query,
    result_count: results.length,
    top_score: results[0]?.score ?? 0,
  });

  const provenanceId = await ctx.provenance.log(
    withProvenanceContext(ctx, {
      type: 'semantic_search:searched',
      status: 'success',
      metadata: { input: { query, top_k: topK, threshold }, output: { result_count: results.length } },
    })
  );

  return {
    success: true,
    results,
    provenance_id: provenanceId,
  };
}

export default handler({
  name: AGENT_ID,
  id: AGENT_ID,
  version: VERSION,
  capabilities: CAPABILITIES,
  description: 'Index documents and search by semantic similarity (RAG foundation)',
  manifest: {
    operations: [
      {
        name: 'index',
        description: 'Index documents into vector store for semantic search',
        paramsSchema: {
          action: { type: 'string', required: true, description: "Must be 'index'" },
          documents: { type: 'array', description: 'Documents with path and optional metadata' },
        },
        resultKind: 'agent.semantic-search.index-result',
      },
      {
        name: 'search',
        description: 'Search indexed documents by semantic similarity',
        paramsSchema: {
          action: { type: 'string', required: true, description: "Must be 'search'" },
          query: { type: 'string', description: 'Search query' },
          top_k: { type: 'number', description: 'Number of results' },
          threshold: { type: 'number', description: 'Minimum similarity (0-1)' },
        },
        resultKind: 'agent.semantic-search.search-result',
      },
    ],
  },
  execute,
});