agentcore-cli 0.31.1__py3-none-any.whl

agentcore/primitives/knowledge_base.py

@@ -0,0 +1,477 @@
+"""Knowledge base add/remove (port of src/cli/primitives/KnowledgeBasePrimitive.ts).
+
+`agentcore add knowledge-base` creates the KB and its data sources. The repeatable ``--source`` flag
+maps to entries in the KB's ``dataSources`` array; re-invoking ``add`` with an existing ``--name``
+appends new data sources (idempotent append). Non-S3 connector configs (``--data-source-type`` +
+``--connector-config <file>``) are materialized under ``app/<kbName>/`` and referenced by
+project-relative path. Binding a pre-existing KB is done via ``add gateway-target --type connector``.
+
+The project spec is handled as its camelCase JSON form (as upstream's plain objects) and validated
+on write by ``ConfigIO``.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import os
+import random
+import re
+import shutil
+import string
+import sys
+from dataclasses import dataclass, field
+from typing import Any
+
+import typer
+from pydantic import ValidationError as PydanticValidationError
+
+from agentcore.cli import interactive
+from agentcore.cli.errors import get_error_message
+from agentcore.cli.run_command import run_cli_command
+from agentcore.lib.constants import APP_DIR
+from agentcore.lib.errors.types import ValidationError
+from agentcore.lib.result import Err, Ok, serialize_result
+from agentcore.lib.schemas.io.config_io import to_jsonable
+from agentcore.lib.schemas.io.path_resolver import find_config_root
+from agentcore.operations.knowledge_base.connector_config import (
+    flag_to_wire_type,
+    is_connector_config_type,
+    read_connector_config,
+)
+from agentcore.primitives.base import BasePrimitive, RemovableResource, RemovalPreview, SchemaChange
+from agentcore.schema import CONNECTOR_ID, KnowledgeBase
+
+# Cheap shape check for early errors at the CLI boundary; S3DataSource in the schema is canonical.
+_S3_URI_PATTERN = re.compile(r"^s3://[^/]+(/.*)?$", re.DOTALL)
+
+
+def _is_s3_uri(uri: str) -> bool:
+    return bool(_S3_URI_PATTERN.match(uri))
+
+
+def _data_source_key(ds: dict[str, Any]) -> str:
+    """S3 sources are keyed by URI, connector sources by their config file path."""
+    return ds["uri"] if ds.get("type") == "S3" else ds["connectorConfigFile"]
+
+
+def _posix_rel(path: str, start: str) -> str:
+    return os.path.relpath(path, start).replace("\\", "/")
+
+
+def zod_error_message(exc: PydanticValidationError) -> str:
+    """Approximation of zod v4's ``ZodError.message`` (``JSON.stringify(issues, null, 2)``)."""
+    issues = [{"code": "custom", "path": list(e.get("loc") or ()), "message": e.get("msg", "")} for e in exc.errors()]
+    return json.dumps(issues, indent=2)
+
+
+@dataclass
+class AddKnowledgeBaseOptions:
+    name: str
+    description: str | None = None
+    source: list[str] = field(default_factory=list)
+    """Repeatable ``--source`` values (S3 URIs). Required for S3 data sources."""
+    connector_config: list[str] = field(default_factory=list)
+    """Repeatable ``--connector-config`` values. Required for non-S3 connectors."""
+    data_source_type: str | None = None
+    """``--data-source-type`` flag (s3 default, or web-crawler/confluence/...)."""
+    json: bool = False
+
+
+@dataclass
+class _PruneAction:
+    gateway_name: str
+    target_name: str
+    removed_target: bool
+
+
+class KnowledgeBasePrimitive(BasePrimitive):
+    kind = "knowledge-base"
+    label = "Knowledge Base"
+
+    # --- add ----------------------------------------------------------------------------
+
+    def add(self, options: AddKnowledgeBaseOptions) -> Ok[dict[str, Any]] | Err[Exception]:
+        try:
+            return Ok(self._add(options))
+        except Exception as exc:  # noqa: BLE001 - surfaced as a failed Result, as upstream
+            return Err(exc)
+
+    def _add(self, options: AddKnowledgeBaseOptions) -> dict[str, Any]:
+        sources = list(options.source or [])
+        connector_configs = list(options.connector_config or [])
+        wire_type = flag_to_wire_type(options.data_source_type or "s3")
+        warnings: list[str] = []
+
+        # Phase 1: pure validation, no side effects. File copies are deferred so a later validation
+        # failure never leaves a stray file behind.
+        if is_connector_config_type(wire_type):
+            if sources:
+                raise ValueError(f"--source is only valid for S3. For {wire_type}, use --connector-config.")
+            if not connector_configs:
+                raise ValueError(f"--connector-config is required for --data-source-type {options.data_source_type}.")
+            seen_sources: set[str] = set()
+            basename_to_source: dict[str, str] = {}
+            for cfg_path in connector_configs:
+                warnings.extend(read_connector_config(cfg_path, wire_type).warnings)
+                resolved_src = os.path.abspath(cfg_path)
+                if resolved_src in seen_sources:
+                    continue  # exact duplicate: the batch-dedup loop below reports it
+                seen_sources.add(resolved_src)
+                base = os.path.basename(resolved_src)
+                prior = basename_to_source.get(base)
+                if prior:
+                    raise ValueError(
+                        f"Connector config files '{prior}' and '{cfg_path}' would both be stored as "
+                        f"'app/{options.name}/{base}'. Rename one so their filenames differ."
+                    )
+                basename_to_source[base] = cfg_path
+
+            def build_data_sources() -> list[dict[str, Any]]:
+                return [
+                    {"type": wire_type, "connectorConfigFile": self._materialize_connector_config(options.name, p)}
+                    for p in connector_configs
+                ]
+
+        else:
+            if connector_configs:
+                raise ValueError("--connector-config is only valid for non-S3 data source types.")
+            if not sources:
+                raise ValueError("At least one --source is required for S3 data sources.")
+            for uri in sources:
+                if not _is_s3_uri(uri):
+                    raise ValueError(f"Invalid S3 URI: {uri}. Expected s3://bucket[/prefix].")
+
+            def build_data_sources() -> list[dict[str, Any]]:
+                return [{"type": "S3", "uri": uri} for uri in sources]
+
+        # Reject duplicates inside this batch up front with an actionable message.
+        if is_connector_config_type(wire_type):
+            project_root = os.path.dirname(self.config_io.get_config_root())
+            batch_keys = [_posix_rel(os.path.abspath(p), project_root) for p in connector_configs]
+        else:
+            batch_keys = sources
+        seen_in_batch: set[str] = set()
+        for key in batch_keys:
+            if key in seen_in_batch:
+                raise ValueError(f"Duplicate data source in this invocation: {key}")
+            seen_in_batch.add(key)
+
+        project = to_jsonable(self.read_project_spec())
+
+        # All validation passed: materialize connector configs and build the data sources.
+        new_data_sources = build_data_sources()
+
+        existing = next((kb for kb in project.get("knowledgeBases", []) if kb["name"] == options.name), None)
+        if existing is not None:
+            return self._append_to_existing(existing, project, new_data_sources, options, warnings)
+
+        kb_input: dict[str, Any] = {"name": options.name}
+        if options.description:
+            kb_input["description"] = options.description
+        kb_input["dataSources"] = new_data_sources
+        try:
+            kb = KnowledgeBase.model_validate(kb_input)
+        except PydanticValidationError as exc:
+            raise ValueError(zod_error_message(exc)) from None
+
+        project.setdefault("knowledgeBases", []).append(to_jsonable(kb))
+        self.write_project_spec(project)
+
+        if not options.json:
+            for w in warnings:
+                print(w, file=sys.stderr)
+
+        return {
+            "knowledgeBaseName": kb.name,
+            "appended": False,
+            "newDataSources": [_data_source_key(ds) for ds in new_data_sources],
+        }
+
+    def _append_to_existing(
+        self,
+        existing: dict[str, Any],
+        project: dict[str, Any],
+        new_data_sources: list[dict[str, Any]],
+        options: AddKnowledgeBaseOptions,
+        warnings: list[str],
+    ) -> dict[str, Any]:
+        # '' and None both mean "no description provided".
+        desc_changed = options.description not in (None, "") and options.description != existing.get("description")
+        if desc_changed:
+            raise ValueError(
+                f'Knowledge base "{options.name}" already exists. Update operations are not supported in Wave 1; '
+                "edit agentcore.json directly to change the description."
+            )
+        existing_keys = {_data_source_key(ds) for ds in existing.get("dataSources", [])}
+        for ds in new_data_sources:
+            key = _data_source_key(ds)
+            if key in existing_keys:
+                raise ValueError(f'Data source "{key}" already exists on knowledge-base "{options.name}".')
+
+        existing.setdefault("dataSources", []).extend(new_data_sources)
+        self.write_project_spec(project)
+
+        if not options.json:
+            for w in warnings:
+                print(w, file=sys.stderr)
+
+        return {
+            "knowledgeBaseName": existing["name"],
+            "appended": True,
+            "newDataSources": [_data_source_key(ds) for ds in new_data_sources],
+        }
+
+    def _materialize_connector_config(self, kb_name: str, cfg_path: str) -> str:
+        """Ensure the config lives under ``app/<kbName>/``; return its project-relative path."""
+        project_root = os.path.dirname(self.config_io.get_config_root())
+        src = os.path.abspath(cfg_path)
+        if not os.path.exists(src):
+            raise ValueError(f"Connector config file not found: {cfg_path}")
+        dest_dir = os.path.join(project_root, APP_DIR, kb_name)
+        dest = os.path.join(dest_dir, os.path.basename(src))
+
+        if os.path.abspath(src) == os.path.abspath(dest):
+            return _posix_rel(dest, project_root)
+        # Never silently overwrite a different file already at the destination.
+        if os.path.exists(dest):
+            raise ValueError(
+                f"Connector config '{cfg_path}' would overwrite the existing file at "
+                f"'{_posix_rel(dest, project_root)}'. Rename it so its filename differs."
+            )
+        os.makedirs(dest_dir, exist_ok=True)
+        shutil.copyfile(src, dest)
+        print(f"Copied connector config to {_posix_rel(dest, project_root)}", file=sys.stderr)
+        return _posix_rel(dest, project_root)
+
+    # --- remove -------------------------------------------------------------------------
+
+    def remove(self, name: str) -> Ok[None] | Err[Exception]:
+        try:
+            project = to_jsonable(self.read_project_spec())
+            kbs = project.get("knowledgeBases", [])
+            idx = next((i for i, kb in enumerate(kbs) if kb["name"] == name), -1)
+            if idx == -1:
+                raise ValueError(f'Knowledge base "{name}" not found.')
+            kb = kbs.pop(idx)
+            if kb.get("gateway"):
+                self._remove_connector_target(project, kb["gateway"], kb["name"])
+            # Without this the spec would be unwriteable: the cross-spec validator rejects
+            # AgenticRetrieveStream entries that reference a KB name no longer in knowledgeBases[].
+            self._prune_agentic_retrieve_references(project, name)
+            self.write_project_spec(project)
+            return Ok(None)
+        except Exception as exc:  # noqa: BLE001
+            return Err(exc)
+
+    def preview_remove(self, name: str) -> RemovalPreview:
+        project = to_jsonable(self.read_project_spec())
+        kb = next((k for k in project.get("knowledgeBases", []) if k["name"] == name), None)
+        if kb is None:
+            raise ValueError(f'Knowledge base "{name}" not found.')
+        summary = [
+            f"Removing knowledge base: {name}",
+            f"  Data sources ({len(kb['dataSources'])}):",
+            *[f"    - {_data_source_key(ds)}" for ds in kb["dataSources"]],
+        ]
+        if kb.get("gateway"):
+            summary.append(f"  Gateway target: '{name}' on '{kb['gateway']}' will be removed")
+
+        after = copy.deepcopy(project)
+        after["knowledgeBases"] = [k for k in after["knowledgeBases"] if k["name"] != name]
+        if kb.get("gateway"):
+            gw = next((g for g in after.get("agentCoreGateways", []) if g["name"] == kb["gateway"]), None)
+            if gw is not None:
+                gw["targets"] = [t for t in gw.get("targets", []) if t["name"] != name]
+        for action in self._prune_agentic_retrieve_references(after, name):
+            if action.removed_target:
+                summary.append(
+                    f"  Gateway \"{action.gateway_name}\" agentic-retrieve target '{action.target_name}' will be "
+                    "removed (was the last KB)"
+                )
+            else:
+                summary.append(
+                    f"  Gateway \"{action.gateway_name}\" agentic-retrieve target '{action.target_name}' will lose "
+                    f"KB '{name}'"
+                )
+
+        return RemovalPreview(
+            summary=summary,
+            directories_to_delete=[],
+            schema_changes=[SchemaChange(file="agentcore/agentcore.json", before=project, after=after)],
+        )
+
+    def _prune_agentic_retrieve_references(self, project: dict[str, Any], kb_reference: str) -> list[_PruneAction]:
+        """Drop ``kb_reference`` from every bedrock-knowledge-bases connector target's configurations.
+
+        Removes the target entirely once no configurations remain.
+        """
+        actions: list[_PruneAction] = []
+        for gw in project.get("agentCoreGateways", []):
+            targets = gw.get("targets", [])
+            for t_idx in range(len(targets) - 1, -1, -1):
+                target = targets[t_idx]
+                if (
+                    target.get("targetType") != "connector"
+                    or target.get("connectorId") != CONNECTOR_ID["BEDROCK_KNOWLEDGE_BASES"]
+                ):
+                    continue
+                configs = target.get("configurations")
+                if configs is None:
+                    configs = []
+
+                agentic_idx = next((i for i, c in enumerate(configs) if c.get("name") == "AgenticRetrieveStream"), -1)
+                if agentic_idx != -1:
+                    agentic_cfg = configs[agentic_idx]
+                    pv = agentic_cfg.get("parameterValues")
+                    retrievers = (pv or {}).get("retrievers") or []
+                    ids = [r["configuration"]["knowledgeBase"]["knowledgeBaseId"] for r in retrievers]
+                    if kb_reference in ids:
+                        remaining = [i for i in ids if i != kb_reference]
+                        if not remaining:
+                            configs.pop(agentic_idx)
+                        else:
+                            agentic_cfg["parameterValues"]["retrievers"] = [
+                                {"configuration": {"knowledgeBase": {"knowledgeBaseId": i}}} for i in remaining
+                            ]
+
+                retrieve_idx = next(
+                    (
+                        i
+                        for i, c in enumerate(configs)
+                        if c.get("name") == "Retrieve"
+                        and (c.get("parameterValues") or {}).get("knowledgeBaseId") == kb_reference
+                    ),
+                    -1,
+                )
+                if retrieve_idx != -1:
+                    configs.pop(retrieve_idx)
+
+                if not configs:
+                    targets.pop(t_idx)
+                    actions.append(_PruneAction(gw["name"], target["name"], True))
+                elif agentic_idx != -1 or retrieve_idx != -1:
+                    actions.append(_PruneAction(gw["name"], target["name"], False))
+        return actions
+
+    @staticmethod
+    def _remove_connector_target(project: dict[str, Any], gateway_name: str, target_name: str) -> None:
+        gateway = next((g for g in project.get("agentCoreGateways", []) if g["name"] == gateway_name), None)
+        if gateway is not None:
+            gateway["targets"] = [t for t in gateway.get("targets", []) if t["name"] != target_name]
+
+    def get_removable(self) -> list[RemovableResource]:
+        try:
+            return [RemovableResource(kb.name) for kb in self.read_project_spec().knowledge_bases or []]
+        except Exception:  # noqa: BLE001
+            return []
+
+    # --- CLI ----------------------------------------------------------------------------
+
+    def _random_name(self) -> str:
+        existing = {kb.name for kb in self.read_project_spec().knowledge_bases or []}
+        alphabet = string.digits + string.ascii_lowercase
+        while True:
+            candidate = "kb-quick-start-" + "".join(random.choice(alphabet) for _ in range(5))
+            if candidate not in existing:
+                return candidate
+
+    def register_commands(self, add_app: typer.Typer, remove_app: typer.Typer) -> None:
+        @add_app.command(
+            self.kind, help="Add a knowledge base (FMKB) to the project, optionally wiring it to a gateway."
+        )
+        def add_command(
+            name: str | None = typer.Option(None, "--name", help="Knowledge base name (default: kb-quick-start-xxxxx)"),
+            description: str | None = typer.Option(
+                None, "--description", help="Optional description (used for tool discovery)"
+            ),
+            source: list[str] = typer.Option(
+                [],
+                "--source",
+                help="S3 URI for a data source (s3://bucket[/prefix]). Repeatable for multiple data sources.",
+            ),
+            data_source_type: str = typer.Option(
+                "s3",
+                "--data-source-type",
+                help="Data source type: s3 (default), web-crawler, confluence, sharepoint, onedrive, google-drive",
+            ),
+            connector_config: list[str] = typer.Option(
+                [],
+                "--connector-config",
+                help="Path to a JSON connector-config file (required for non-S3 types). Repeatable.",
+            ),
+            as_json: bool = typer.Option(False, "--json", help="Output as JSON [non-interactive]"),
+        ) -> None:
+            if not find_config_root():
+                typer.echo("No agentcore project found. Run `agentcore create` first.", err=True)
+                raise typer.Exit(1)
+
+            # --data-source-type always has a default, so only user-supplied flags count.
+            if not (name or description or source or connector_config or as_json):
+                try:
+                    interactive.require_tty()
+                    self.add_interactive()
+                except typer.Exit:
+                    raise
+                except Exception as exc:  # noqa: BLE001
+                    typer.echo(get_error_message(exc), err=True)
+                    raise typer.Exit(1) from None
+                raise typer.Exit(0)
+
+            def action() -> None:
+                result = self.add(
+                    AddKnowledgeBaseOptions(
+                        name=name or self._random_name(),
+                        description=description,
+                        source=list(source or []),
+                        data_source_type=data_source_type,
+                        connector_config=list(connector_config or []),
+                        json=as_json,
+                    )
+                )
+                if not result.success:
+                    raise result.error
+                value = result.value
+                if as_json:
+                    typer.echo(json.dumps(serialize_result(result), separators=(",", ":")))
+                elif value["appended"]:
+                    for uri in value["newDataSources"]:
+                        typer.echo(f"Added data source '{uri}' to knowledge-base '{value['knowledgeBaseName']}'")
+                else:
+                    typer.echo(f"Added knowledge base '{value['knowledgeBaseName']}'")
+                    for uri in value["newDataSources"]:
+                        typer.echo(f"  with data source '{uri}'")
+
+            run_cli_command(as_json, action)
+
+        @remove_app.command(self.kind, help="Remove a knowledge base from the project")
+        def remove_command(
+            name: str | None = typer.Option(None, "--name", help="Knowledge base name"),
+            as_json: bool = typer.Option(False, "--json", help="Output as JSON [non-interactive]"),
+        ) -> None:
+            if not find_config_root():
+                typer.echo("No agentcore project found. Run `agentcore create` first.", err=True)
+                raise typer.Exit(1)
+
+            def action() -> None:
+                if not name:
+                    raise ValidationError("A --name is required for `agentcore remove knowledge-base`.")
+                result = self.remove(name)
+                if not result.success:
+                    raise result.error
+                if as_json:
+                    typer.echo(json.dumps(serialize_result(result), separators=(",", ":")))
+                else:
+                    typer.echo(f"Removed knowledge base '{name}'")
+
+            run_cli_command(as_json, action)
+
+    # --- interactive (port of tui/screens/knowledge-base AddKnowledgeBaseScreen + Flow) ----------
+
+    def add_interactive(self) -> None:
+        from agentcore.primitives._knowledge_base_wizard import run_add_knowledge_base_wizard
+
+        run_add_knowledge_base_wizard(self)
+
+
+PRIMITIVE = KnowledgeBasePrimitive()