sh-mcp/packages/knowledge-base/src/client.ts

/**
 * Knowledge-base client interface.
 *
 * The production implementation calls Amazon Bedrock Knowledge Base Retrieve API.
 * Swap the implementation by injecting a different KnowledgeBaseClient at the call
 * site — tests pass a mock, production code uses BedrockKnowledgeBaseClient.
 *
 * Future note: when migrating to Salesforce Data Cloud, replace
 * BedrockKnowledgeBaseClient with a DataCloudKnowledgeBaseClient that satisfies
 * the same KnowledgeBaseClient interface.
 */

/** A single retrieval result returned by the knowledge base. */
export interface KnowledgeBaseResult {
  /** Source document URI or human-readable reference (e.g. S3 URI, Notion page title). */
  source: string;
  /** Relevance score in the range [0, 1] as returned by the underlying retriever. */
  score: number;
  /** Text passage extracted from the source document. */
  passage: string;
}

/** Options forwarded to the underlying retriever on each query. */
export interface RetrieveOptions {
  /** Free-text query string. */
  query: string;
  /**
   * Maximum number of results to return.
   * Defaults to 5 if omitted; callers should not exceed 20.
   */
  maxResults?: number;
}

/**
 * The interface every knowledge-base client must satisfy.
 * Production code uses BedrockKnowledgeBaseClient; tests supply a mock.
 */
export interface KnowledgeBaseClient {
  retrieve(options: RetrieveOptions): Promise<KnowledgeBaseResult[]>;
}

// ---------------------------------------------------------------------------
// Production implementation — Bedrock Knowledge Base Retrieve API
// ---------------------------------------------------------------------------

/**
 * Configuration for the production Bedrock KB client.
 * All values come from environment variables; no defaults are hard-coded so
 * that the module can be imported without triggering any AWS calls.
 */
export interface BedrockKnowledgeBaseClientConfig {
  /** Bedrock Knowledge Base ID (e.g. "ABCD1234EF"). */
  knowledgeBaseId: string;
  /** AWS region (e.g. "us-east-1"). */
  region: string;
}

/**
 * Production Bedrock KB client.
 *
 * The @aws-sdk/client-bedrock-agent-runtime package is imported LAZILY inside
 * retrieve() so that importing this module at the top level (e.g. in tests)
 * does NOT trigger any network activity or AWS credential resolution.
 *
 * TODO(auth-layer): once the deferred auth layer is in place, thread the
 * caller's AWS credentials / assumed role ARN through here if we want
 * per-user IAM audit trails on Bedrock calls.
 */
export class BedrockKnowledgeBaseClient implements KnowledgeBaseClient {
  private readonly config: BedrockKnowledgeBaseClientConfig;

  constructor(config: BedrockKnowledgeBaseClientConfig) {
    this.config = config;
  }

  /**
   * Calls the Bedrock Retrieve API.
   *
   * The AWS SDK import is deferred to keep module load side-effect-free.
   * If the environment lacks AWS credentials this will throw at call time,
   * not at import time — which is the desired behaviour for testing.
   */
  async retrieve(options: RetrieveOptions): Promise<KnowledgeBaseResult[]> {
    // @aws-sdk/client-bedrock-agent-runtime is intentionally absent from package.json
    // until the Lambda runtime bundle is assembled. The module specifier is stored in a
    // variable so TypeScript skips static module resolution at build time.
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const dynImport = (s: string): Promise<any> =>
      new Function('s', 'return import(s)')(s) as Promise<any>;

    interface BedrockRetrievalResult {
      location?: { s3Location?: { uri?: string }; type?: string };
      score?: number;
      content?: { text?: string };
    }
    interface BedrockRetrieveResponse {
      retrievalResults?: BedrockRetrievalResult[];
    }

    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const { BedrockAgentRuntimeClient, RetrieveCommand } = (await dynImport(
      '@aws-sdk/client-bedrock-agent-runtime',
    )) as {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      BedrockAgentRuntimeClient: new (cfg: { region: string }) => {
        send: (cmd: any) => Promise<BedrockRetrieveResponse>;
      };
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      RetrieveCommand: new (input: any) => unknown;
    };

    const client = new BedrockAgentRuntimeClient({ region: this.config.region });

    const command = new RetrieveCommand({
      knowledgeBaseId: this.config.knowledgeBaseId,
      retrievalQuery: { text: options.query },
      retrievalConfiguration: {
        vectorSearchConfiguration: {
          numberOfResults: options.maxResults ?? 5,
        },
      },
    });

    const response = await client.send(command);
    const rawResults = response.retrievalResults ?? [];

    return rawResults.map((r: BedrockRetrievalResult) => ({
      source: r.location?.s3Location?.uri ?? r.location?.type ?? 'unknown',
      score: r.score ?? 0,
      passage: r.content?.text ?? '',
    }));
  }
}

/**
 * Build a production BedrockKnowledgeBaseClient from environment variables.
 *
 * Expected env vars:
 *   KNOWLEDGE_BASE_ID  — Bedrock Knowledge Base ID
 *   AWS_REGION         — AWS region (falls back to 'us-east-1')
 *
 * This factory is intentionally NOT called at module import time.
 */
export function createBedrockClientFromEnv(): BedrockKnowledgeBaseClient {
  const knowledgeBaseId = process.env['KNOWLEDGE_BASE_ID'];
  if (!knowledgeBaseId) {
    throw new Error(
      'KNOWLEDGE_BASE_ID environment variable is required for the production KB client',
    );
  }
  return new BedrockKnowledgeBaseClient({
    knowledgeBaseId,
    region: process.env['AWS_REGION'] ?? 'us-east-1',
  });
}