diff --git a/examples/ingest-plugin/README.md b/examples/ingest-plugin/README.md new file mode 100644 index 00000000..e8ee9b2e --- /dev/null +++ b/examples/ingest-plugin/README.md @@ -0,0 +1,27 @@ +# OpenKB Ingest Plugin Example + +This directory shows the minimal shape of an external package that contributes +bundle ingest components through Python entry points. + +Install a package like this in the same environment as OpenKB, then enable its +components in `.openkb/config.yaml`: + +```yaml +ingest: + pipeline: bundle + importers: + enabled: + - example_text + normalizers: + enabled: + - example_text +``` + +The entry point groups are: + +- `openkb.ingest.importers` +- `openkb.ingest.normalizers` +- `openkb.ingest.enrichers` + +Each entry point should resolve to a class or factory returning an object that +matches the corresponding OpenKB ingest protocol. diff --git a/examples/ingest-plugin/openkb_example_ingest.py b/examples/ingest-plugin/openkb_example_ingest.py new file mode 100644 index 00000000..afd3b732 --- /dev/null +++ b/examples/ingest-plugin/openkb_example_ingest.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord, TextBlock + + +class ExampleTextImporter: + name = "example_text" + + def can_handle(self, target: str, context) -> bool: + del context + return target.startswith("example-text:") + + def import_source(self, target: str, context) -> IngestInput: + path = context.staging_dir / "raw" / "example.txt" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(target.removeprefix("example-text:").strip() + "\n", encoding="utf-8") + return IngestInput( + target=target, + path=path, + source_uri=target, + media_type="text/x-example", + metadata={"display_name": "example.txt"}, + ) + + +class ExampleTextNormalizer: + name = "example_text" + + def supports(self, input_: IngestInput, context) -> bool: + del context + return input_.media_type == "text/x-example" + + def normalize(self, input_: IngestInput, context) -> DocumentBundle: + del context + if input_.path is None: + raise ValueError("Example normalizer requires a local path.") + source_uri = input_.source_uri or input_.target + text = input_.path.read_text(encoding="utf-8") + return DocumentBundle( + id=source_uri, + title="Example Text", + source_uri=source_uri, + blocks=[TextBlock(text)], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_identity": source_uri, + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) diff --git a/examples/ingest-plugin/pyproject.toml b/examples/ingest-plugin/pyproject.toml new file mode 100644 index 00000000..71840d20 --- /dev/null +++ b/examples/ingest-plugin/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "openkb-example-ingest-plugin" +version = "0.1.0" +requires-python = ">=3.10" +dependencies = ["openkb"] + +[project.entry-points."openkb.ingest.importers"] +example_text = "openkb_example_ingest:ExampleTextImporter" + +[project.entry-points."openkb.ingest.normalizers"] +example_text = "openkb_example_ingest:ExampleTextNormalizer" diff --git a/openkb/cli.py b/openkb/cli.py index f39d0b44..a70df64c 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -240,6 +240,17 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: ".csv", } +BUNDLE_ONLY_EXTENSIONS = { + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".bmp", +} + +BUNDLE_SUPPORTED_EXTENSIONS = SUPPORTED_EXTENSIONS | BUNDLE_ONLY_EXTENSIONS + # Map raw doc types to display types _TYPE_DISPLAY_MAP = { "long_pdf": "pageindex", @@ -987,9 +998,16 @@ def init(model, language): help="Import an already-indexed PageIndex Cloud document by its doc-id " "(no local file). Mutually exclusive with PATH.", ) +@click.option( + "--ingest-pipeline", + "ingest_pipeline", + type=click.Choice(["legacy", "bundle"]), + default=None, + help="Select the add ingest pipeline. Defaults to .openkb/config.yaml ingest.pipeline, then legacy.", +) @click.pass_context @_with_kb_lock(exclusive=True) -def add(ctx, path, from_pageindex_cloud): +def add(ctx, path, from_pageindex_cloud, ingest_pipeline): """Add a document or directory of documents at PATH to the knowledge base. PATH may be a local file, a local directory (which is walked @@ -1021,13 +1039,23 @@ def add(ctx, path, from_pageindex_cloud): click.echo("Provide a PATH or use --from-pageindex-cloud .") return - # URL ingest: download into raw/ first, then call add_single_file explicitly. - # Keep staged conversion enabled so converted source artifacts do not touch - # the live KB before the mutation snapshot exists. The tri-state outcome - # still lets us clean up the just-downloaded raw file on dedup. + from openkb.ingest.add import add_bundle_target, resolve_ingest_pipeline + + config = load_config(kb_dir / ".openkb" / "config.yaml") + pipeline = resolve_ingest_pipeline(config, ingest_pipeline) + supported_extensions = ( + BUNDLE_SUPPORTED_EXTENSIONS if pipeline == "bundle" else SUPPORTED_EXTENSIONS + ) + + # Legacy URL ingest downloads into raw/ first, then calls add_single_file. + # Bundle URL ingest keeps downloads inside the bundle staging directory and + # commits through the same mutation boundary as local bundle files. from openkb.url_ingest import looks_like_url, fetch_url_to_raw if looks_like_url(path): + if pipeline == "bundle": + add_bundle_target(path, kb_dir, legacy_fallback=add_single_file) + return fetched = fetch_url_to_raw(path, kb_dir) if fetched is None: return @@ -1049,7 +1077,7 @@ def add(ctx, path, from_pageindex_cloud): files = [ f for f in sorted(target.rglob("*")) - if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS + if f.is_file() and f.suffix.lower() in supported_extensions ] if not files: click.echo(f"No supported files found in {path}.") @@ -1058,15 +1086,21 @@ def add(ctx, path, from_pageindex_cloud): click.echo(f"Found {total} supported file(s) in {path}.") for i, f in enumerate(files, 1): click.echo(f"\n[{i}/{total}] ", nl=False) - add_single_file(f, kb_dir) + if pipeline == "bundle": + add_bundle_target(f, kb_dir, legacy_fallback=add_single_file) + else: + add_single_file(f, kb_dir) else: - if target.suffix.lower() not in SUPPORTED_EXTENSIONS: + if target.suffix.lower() not in supported_extensions: click.echo( f"Unsupported file type: {target.suffix}. " - f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}" + f"Supported: {', '.join(sorted(supported_extensions))}" ) return - add_single_file(target, kb_dir) + if pipeline == "bundle": + add_bundle_target(target, kb_dir, legacy_fallback=add_single_file) + else: + add_single_file(target, kb_dir) def _stream_to_tty() -> bool: @@ -1309,6 +1343,13 @@ def remove(ctx, identifier, keep_raw, keep_empty, dry_run, yes): ) ) + bundle_path = None + if meta.get("bundle_path"): + candidate = kb_dir / meta["bundle_path"] + if candidate.exists(): + bundle_path = candidate + actions.append(("DELETE", str(candidate.relative_to(kb_dir)))) + # Scan concept pages to predict which will be edited vs. deleted. # Only frontmatter ``sources:`` membership drives the plan — body-only # references (e.g. a stray ``See also:`` line a user added by hand @@ -1424,6 +1465,8 @@ def remove(ctx, identifier, keep_raw, keep_empty, dry_run, yes): source_json.unlink(missing_ok=True) if images_dir.is_dir(): shutil.rmtree(images_dir, ignore_errors=True) + if bundle_path is not None: + bundle_path.unlink(missing_ok=True) concept_result = remove_doc_from_concept_pages( wiki_dir, diff --git a/openkb/ingest/__init__.py b/openkb/ingest/__init__.py new file mode 100644 index 00000000..dd648e61 --- /dev/null +++ b/openkb/ingest/__init__.py @@ -0,0 +1,7 @@ +"""Pluggable ingest pipeline for OpenKB.""" + +from __future__ import annotations + +from openkb.ingest.add import add_bundle_target, resolve_ingest_pipeline + +__all__ = ["add_bundle_target", "resolve_ingest_pipeline"] diff --git a/openkb/ingest/add.py b/openkb/ingest/add.py new file mode 100644 index 00000000..3c809017 --- /dev/null +++ b/openkb/ingest/add.py @@ -0,0 +1,500 @@ +"""Bundle ingest orchestration and add integration.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Literal +from urllib.parse import unquote, urlparse + +import click + +from openkb.add_coordinator import AddMutationPlan, _cleanup_staging_dirs, run_add_mutation +from openkb.config import DEFAULT_CONFIG, load_config +from openkb.converter import _registry_path +from openkb.ingest.config import IngestOptions, resolve_ingest_options +from openkb.ingest.context import IngestContext +from openkb.ingest.enrichers import ImageVisionEnricher, LinkDiscoveryEnricher +from openkb.ingest.exceptions import LongDocumentFallback +from openkb.ingest.importers import FeishuImporter, FileImporter, UrlImporter +from openkb.ingest.importers.feishu import looks_like_feishu_url +from openkb.ingest.models import DiscoveredLink, DocumentBundle, ProvenanceRecord, RenderedBundle +from openkb.ingest.normalizers import ( + ImageNormalizer, + MarkdownNormalizer, + MarkItDownNormalizer, + PdfNormalizer, + PlainTextNormalizer, +) +from openkb.ingest.plugins import ( + ENRICHER_ENTRY_POINT_GROUP, + IMPORTER_ENTRY_POINT_GROUP, + NORMALIZER_ENTRY_POINT_GROUP, + load_ingest_entry_points, +) +from openkb.ingest.registry import ( + BundleEnricher, + BundleNormalizer, + BundleNormalizerRegistry, + SourceImporter, + SourceImporterRegistry, +) +from openkb.ingest.render import render_bundle_to_staging +from openkb.locks import kb_ingest_lock +from openkb.log import append_log +from openkb.mutation import publish_staged_tree +from openkb.state import HashRegistry + +logger = logging.getLogger(__name__) + +IngestPipeline = Literal["legacy", "bundle"] +AddOutcome = Literal["added", "skipped", "failed"] +LegacyFallback = Callable[[Path, Path], AddOutcome] +_BUILTIN_IMPORTERS = {"file", "feishu", "url"} +_BUILTIN_NORMALIZERS = {"markdown", "text", "pdf-local", "markitdown", "image"} +_BUILTIN_ENRICHERS = {"link_discovery", "image_vision"} + + +@dataclass +class _BundleAddState: + seen_targets: set[str] + processed_documents: int = 0 + + +def resolve_ingest_pipeline(config: dict, override: str | None = None) -> IngestPipeline: + """Resolve the requested ingest pipeline with legacy as the compatibility default.""" + raw = override + if raw is None: + ingest = config.get("ingest") + if isinstance(ingest, dict): + raw = ingest.get("pipeline") + if raw in (None, "", "legacy"): + return "legacy" + if raw == "bundle": + return "bundle" + raise click.BadParameter( + "ingest pipeline must be one of: legacy, bundle", + param_hint="'--ingest-pipeline'", + ) + + +def add_bundle_target( + target: str | Path, + kb_dir: Path, + *, + legacy_fallback: LegacyFallback | None = None, +) -> AddOutcome: + """Import one target through the bundle pipeline under the KB mutation lock.""" + with kb_ingest_lock(kb_dir / ".openkb"): + return _add_bundle_target_locked(str(target), kb_dir, legacy_fallback=legacy_fallback) + + +def _add_bundle_target_locked( + target: str, + kb_dir: Path, + *, + legacy_fallback: LegacyFallback | None = None, +) -> AddOutcome: + openkb_dir = kb_dir / ".openkb" + config = load_config(openkb_dir / "config.yaml") + _setup_ingest_llm(kb_dir) + model: str = config.get("model", DEFAULT_CONFIG["model"]) + options = resolve_ingest_options(config) + + state = _BundleAddState(seen_targets=set()) + queue: list[tuple[str, int, str | None]] = [(target, 0, None)] + outcomes: list[AddOutcome] = [] + while queue: + current_target, depth, parent_uri = queue.pop(0) + if state.processed_documents >= options.max_documents: + click.echo( + f" [SKIP] Bundle recursive import reached max_documents={options.max_documents}." + ) + break + result, bundle = _add_one_bundle_target_locked( + current_target, + kb_dir, + model, + options, + state, + legacy_fallback=legacy_fallback, + parent_uri=parent_uri, + ) + outcomes.append(result) + if result == "failed": + return "failed" + if bundle is None or not options.link_discovery or depth >= options.max_depth: + continue + for child in _allowed_recursive_targets(bundle, options): + key = _target_key(child) + if key in state.seen_targets: + continue + queue.append((str(child), depth + 1, bundle.source_uri)) + + return _summarize_outcomes(outcomes) + + +def _add_one_bundle_target_locked( + target: str, + kb_dir: Path, + model: str, + options: IngestOptions, + state: _BundleAddState, + *, + legacy_fallback: LegacyFallback | None = None, + parent_uri: str | None = None, +) -> tuple[AddOutcome, DocumentBundle | None]: + openkb_dir = kb_dir / ".openkb" + + context: IngestContext | None = None + target_key = _target_key(target) + if target_key in state.seen_targets: + return "skipped", None + state.seen_targets.add(target_key) + state.processed_documents += 1 + + click.echo(f"Adding: {_display_target(target)}") + try: + context = IngestContext.for_target(kb_dir, target) + importer = SourceImporterRegistry(_source_importers(options)).resolve(target, context) + input_ = importer.import_source(target, context) + if input_.path is None: + raise ValueError("Bundle file ingest requires a local file path.") + + file_hash = HashRegistry.hash_file(input_.path) + registry = HashRegistry(openkb_dir / "hashes.json") + if registry.is_known(file_hash): + stored = registry.get(file_hash) or {} + name = stored.get("name") or input_.path.name + click.echo(f" [SKIP] Already in knowledge base: {name}") + _cleanup_staging_dirs([context.staging_dir]) + return "skipped", None + + bundle = BundleNormalizerRegistry(_normalizers(options)).normalize(input_, context) + if parent_uri: + bundle = _with_parent_provenance(bundle, parent_uri) + bundle = _enrich_bundle(bundle, context, options) + rendered = render_bundle_to_staging(bundle, context) + except LongDocumentFallback as exc: + if legacy_fallback is None: + _cleanup_staging_dirs([context.staging_dir if context is not None else None]) + click.echo(f" [ERROR] Bundle ingest requires legacy fallback: {exc.reason}") + return "failed", None + click.echo(f" Falling back to legacy long-document ingest: {exc.reason}") + try: + outcome = legacy_fallback(exc.path, kb_dir) + finally: + _cleanup_staging_dirs([context.staging_dir if context is not None else None]) + return outcome, None + except Exception as exc: + click.echo(f" [ERROR] Bundle ingest failed: {exc}") + logger.debug("Bundle ingest traceback:", exc_info=True) + _cleanup_staging_dirs([context.staging_dir if context is not None else None]) + return "failed", None + + if rendered.is_long_doc: + click.echo(" [ERROR] Bundle long-document commit is not implemented yet.") + _cleanup_staging_dirs([context.staging_dir]) + return "failed", None + + return _commit_rendered_bundle(rendered, kb_dir, model, context.staging_dir), bundle + + +def _source_importers(options: IngestOptions) -> list[SourceImporter]: + importers: list[SourceImporter] = [] + if "file" in options.enabled_importers: + importers.append(FileImporter()) + if "feishu" in options.enabled_importers: + importers.append(FeishuImporter()) + if "url" in options.enabled_importers: + importers.append(UrlImporter()) + importers.extend( + load_ingest_entry_points( + IMPORTER_ENTRY_POINT_GROUP, + options.enabled_importers, + exclude=_BUILTIN_IMPORTERS, + ) + ) + return importers + + +def _normalizers(options: IngestOptions) -> list[BundleNormalizer]: + normalizers: list[BundleNormalizer] = [ + MarkdownNormalizer(), + PlainTextNormalizer(), + PdfNormalizer(), + MarkItDownNormalizer(), + ImageNormalizer(), + ] + normalizers.extend( + load_ingest_entry_points( + NORMALIZER_ENTRY_POINT_GROUP, + options.enabled_normalizers, + exclude=_BUILTIN_NORMALIZERS, + ) + ) + return normalizers + + +def _enrich_bundle( + bundle: DocumentBundle, + context: IngestContext, + options: IngestOptions, +) -> DocumentBundle: + enrichers: list[BundleEnricher] = [] + if options.link_discovery: + enrichers.append(LinkDiscoveryEnricher()) + if "image_vision" in options.enabled_enrichers: + enrichers.append(ImageVisionEnricher()) + enrichers.extend( + load_ingest_entry_points( + ENRICHER_ENTRY_POINT_GROUP, + options.enabled_enrichers, + exclude=_BUILTIN_ENRICHERS, + ) + ) + for enricher in enrichers: + if enricher.applies_to(bundle, context): + bundle = enricher.enrich(bundle, context) + return bundle + + +def _with_parent_provenance(bundle: DocumentBundle, parent_uri: str) -> DocumentBundle: + if not bundle.provenance: + return replace( + bundle, + provenance=[ + ProvenanceRecord( + source_uri=bundle.source_uri, + relationship="child", + parent_uri=parent_uri, + metadata={"parent_uri": parent_uri}, + ) + ], + ) + first = bundle.provenance[0] + metadata = {**first.metadata, "parent_uri": parent_uri} + provenance = [ + replace(first, relationship="child", parent_uri=parent_uri, metadata=metadata), + *bundle.provenance[1:], + ] + return replace(bundle, provenance=provenance) + + +def _allowed_recursive_targets(bundle: DocumentBundle, options: IngestOptions) -> list[str | Path]: + targets: list[str | Path] = [] + for link in bundle.links: + target = _resolve_recursive_target(link, bundle, options) + if target is not None: + targets.append(target) + return targets + + +def _resolve_recursive_target( + link: DiscoveredLink, + bundle: DocumentBundle, + options: IngestOptions, +) -> str | Path | None: + parsed = urlparse(link.uri) + if parsed.scheme in {"http", "https"}: + if not _network_importer_enabled(link.uri, options): + return None + if not _domain_allowed(parsed.netloc, options.allow_domains): + return None + return _strip_url_fragment(link.uri) + if parsed.scheme and parsed.scheme != "file": + return None + if "file" not in options.enabled_importers: + return None + + raw_path = unquote(parsed.path if parsed.scheme == "file" else _strip_fragment(link.uri)) + if not raw_path: + return None + path = Path(raw_path) + if not path.is_absolute(): + source_path = bundle.metadata.get("source_path") + if not isinstance(source_path, str) or not source_path: + return None + path = Path(source_path).parent / path + resolved = path.resolve() + return resolved if resolved.is_file() else None + + +def _strip_fragment(uri: str) -> str: + return uri.split("#", 1)[0].split("?", 1)[0].strip() + + +def _strip_url_fragment(uri: str) -> str: + return uri.split("#", 1)[0].strip() + + +def _network_importer_enabled(uri: str, options: IngestOptions) -> bool: + if looks_like_feishu_url(uri) and "feishu" in options.enabled_importers: + return True + return "url" in options.enabled_importers + + +def _domain_allowed(netloc: str, allow_domains: tuple[str, ...]) -> bool: + if not allow_domains: + return False + host = netloc.rsplit("@", 1)[-1].split(":", 1)[0].lower() + for domain in allow_domains: + cleaned = domain.lower().strip() + if host == cleaned or host.endswith(f".{cleaned}"): + return True + return False + + +def _target_key(target: str | Path) -> str: + path = Path(str(target)).expanduser() + if path.exists(): + return path.resolve().as_posix() + return str(target) + + +def _display_target(target: str) -> str: + parsed = urlparse(target) + if parsed.scheme in {"http", "https"}: + return target + return Path(target).name + + +def _summarize_outcomes(outcomes: list[AddOutcome]) -> AddOutcome: + if any(outcome == "added" for outcome in outcomes): + return "added" + if any(outcome == "failed" for outcome in outcomes): + return "failed" + return "skipped" + + +def _commit_rendered_bundle( + rendered: RenderedBundle, + kb_dir: Path, + model: str, + staging_dir: Path, +) -> AddOutcome: + final_raw, final_source, final_bundle = _final_artifact_paths(rendered, kb_dir) + + def commit_body(_snapshot) -> None: + publish_staged_tree(staging_dir, kb_dir) + if rendered.source_path is None or final_source is None: + raise RuntimeError(f"Rendered bundle has no source artifact: {rendered.display_name}") + + _compile_short_doc(rendered.doc_name, final_source, kb_dir, model) + + registry = HashRegistry(kb_dir / ".openkb" / "hashes.json") + meta = { + "name": rendered.display_name, + "doc_name": rendered.doc_name, + "type": rendered.doc_type, + "path": rendered.source_identity, + } + if final_raw is not None: + meta["raw_path"] = _registry_path(final_raw, kb_dir) + meta["source_path"] = _registry_path(final_source, kb_dir) + if final_bundle is not None: + meta["bundle_path"] = _registry_path(final_bundle, kb_dir) + + registry.remove_by_doc_name(rendered.doc_name) + for existing_hash, existing_meta in list(registry.all_entries().items()): + if ( + existing_hash != rendered.content_hash + and not existing_meta.get("doc_name") + and existing_meta.get("name") == rendered.display_name + ): + registry.remove_by_hash(existing_hash) + registry.add(rendered.content_hash, meta) + + def append_ingest_log() -> None: + append_log(kb_dir / "wiki", "ingest", rendered.display_name) + + plan = AddMutationPlan( + operation="add", + details={ + "file_hash": rendered.content_hash, + "name": rendered.display_name, + "doc_name": rendered.doc_name, + }, + touched_paths=_snapshot_add_paths( + kb_dir, + rendered.doc_name, + final_raw, + final_source, + final_bundle, + ), + body=commit_body, + post_commit_hooks=[append_ingest_log], + hardlink_dirs={ + kb_dir / "wiki" / "concepts", + kb_dir / "wiki" / "entities", + }, + staging_dirs=[staging_dir], + ) + if not run_add_mutation(kb_dir, plan): + return "failed" + click.echo(f" [OK] {rendered.display_name} added to knowledge base.") + return "added" + + +def _setup_ingest_llm(kb_dir: Path) -> None: + from openkb.cli import _setup_llm_key + + _setup_llm_key(kb_dir) + + +def _compile_short_doc(doc_name: str, source_path: Path, kb_dir: Path, model: str) -> None: + from openkb.agent.compiler import compile_short_doc + from openkb.cli import _run_compile_with_retry + + _run_compile_with_retry( + lambda: compile_short_doc(doc_name, source_path, kb_dir, model), + label="Compiling short doc", + ) + + +def _final_artifact_paths( + rendered: RenderedBundle, + kb_dir: Path, +) -> tuple[Path | None, Path | None, Path | None]: + final_raw = None + final_source = None + final_bundle = None + if rendered.raw_path is not None: + final_raw = kb_dir / "raw" / rendered.raw_path.name + if rendered.source_path is not None: + final_source = kb_dir / "wiki" / "sources" / rendered.source_path.name + if rendered.bundle_path is not None: + final_bundle = kb_dir / ".openkb" / "bundles" / rendered.bundle_path.name + return final_raw, final_source, final_bundle + + +def _snapshot_add_paths( + kb_dir: Path, + doc_name: str, + final_raw: Path | None, + final_source: Path | None, + final_bundle: Path | None, +) -> list[Path]: + paths = [ + kb_dir / ".openkb" / "hashes.json", + kb_dir / ".openkb" / "pageindex.db", + kb_dir / ".openkb" / "pageindex.db-wal", + kb_dir / ".openkb" / "pageindex.db-shm", + kb_dir / ".openkb" / "pageindex.db-journal", + kb_dir / "wiki" / "summaries" / f"{doc_name}.md", + kb_dir / "wiki" / "sources" / f"{doc_name}.json", + kb_dir / "wiki" / "sources" / "images" / doc_name, + kb_dir / "wiki" / "concepts", + kb_dir / "wiki" / "entities", + kb_dir / "wiki" / "index.md", + kb_dir / "wiki" / "log.md", + ] + if final_raw is not None: + paths.append(final_raw) + if final_source is not None: + paths.append(final_source) + if final_bundle is not None: + paths.append(final_bundle) + return paths diff --git a/openkb/ingest/config.py b/openkb/ingest/config.py new file mode 100644 index 00000000..cbe14acc --- /dev/null +++ b/openkb/ingest/config.py @@ -0,0 +1,71 @@ +"""Configuration parsing for bundle ingest.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class IngestOptions: + link_discovery: bool = False + max_depth: int = 0 + max_documents: int = 50 + allow_domains: tuple[str, ...] = () + enabled_importers: tuple[str, ...] = ("file",) + enabled_normalizers: tuple[str, ...] = () + enabled_enrichers: tuple[str, ...] = () + + +def resolve_ingest_options(config: dict[str, Any]) -> IngestOptions: + raw = config.get("ingest") + ingest = raw if isinstance(raw, dict) else {} + + enabled_enrichers = _enabled_names(ingest.get("enrichers"), default=()) + link_discovery = bool(ingest.get("link_discovery", False)) or ( + "link_discovery" in enabled_enrichers + ) + + return IngestOptions( + link_discovery=link_discovery, + max_depth=_int_at_least(ingest.get("max_depth"), default=0, minimum=0), + max_documents=_int_at_least(ingest.get("max_documents"), default=50, minimum=1), + allow_domains=tuple(_string_list(ingest.get("allow_domains"))), + enabled_importers=tuple(_enabled_names(ingest.get("importers"), default=("file",))), + enabled_normalizers=tuple(_enabled_names(ingest.get("normalizers"), default=())), + enabled_enrichers=tuple(enabled_enrichers), + ) + + +def _int_at_least(value: object, *, default: int, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, (int, str)): + return default + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return max(minimum, parsed) + + +def _enabled_names(value: object, *, default: tuple[str, ...]) -> list[str]: + if value is None: + return list(default) + if not isinstance(value, dict): + return list(default) + raw = value.get("enabled") + if raw is None: + return list(default) + return _string_list(raw) + + +def _string_list(value: object) -> list[str]: + if not isinstance(value, list): + return [] + names: list[str] = [] + for item in value: + if not isinstance(item, str): + continue + cleaned = item.strip() + if cleaned and cleaned not in names: + names.append(cleaned) + return names diff --git a/openkb/ingest/context.py b/openkb/ingest/context.py new file mode 100644 index 00000000..2a497ff3 --- /dev/null +++ b/openkb/ingest/context.py @@ -0,0 +1,28 @@ +"""Context object shared across ingest pipeline stages.""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from pathlib import Path + +from openkb.config import load_config +from openkb.converter import _sanitize_stem + + +@dataclass(frozen=True) +class IngestContext: + kb_dir: Path + config: dict + staging_dir: Path + + @classmethod + def for_target(cls, kb_dir: Path, target_name: str) -> "IngestContext": + safe = _sanitize_stem(Path(target_name).stem) + staging_dir = kb_dir / ".openkb" / "staging" / f"bundle-{safe}-{uuid.uuid4().hex[:8]}" + staging_dir.mkdir(parents=True, exist_ok=False) + return cls( + kb_dir=kb_dir, + config=load_config(kb_dir / ".openkb" / "config.yaml"), + staging_dir=staging_dir, + ) diff --git a/openkb/ingest/enrichers/__init__.py b/openkb/ingest/enrichers/__init__.py new file mode 100644 index 00000000..3e8451df --- /dev/null +++ b/openkb/ingest/enrichers/__init__.py @@ -0,0 +1,8 @@ +"""Built-in bundle enrichers.""" + +from __future__ import annotations + +from openkb.ingest.enrichers.image_vision import ImageVisionEnricher +from openkb.ingest.enrichers.link_discovery import LinkDiscoveryEnricher + +__all__ = ["ImageVisionEnricher", "LinkDiscoveryEnricher"] diff --git a/openkb/ingest/enrichers/image_vision.py b/openkb/ingest/enrichers/image_vision.py new file mode 100644 index 00000000..e9971c3f --- /dev/null +++ b/openkb/ingest/enrichers/image_vision.py @@ -0,0 +1,161 @@ +"""Image OCR and visual-description enricher.""" + +from __future__ import annotations + +import base64 +import json +from dataclasses import dataclass, replace +from typing import Any + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import Asset, Block, DocumentBundle, ImageBlock + + +@dataclass(frozen=True) +class ImageVisionResult: + visual_description: str | None = None + ocr_text: str | None = None + keywords: list[str] | None = None + uncertainty: str | None = None + + +class ImageVisionEnricher: + name = "image_vision" + + def applies_to(self, bundle: DocumentBundle, context: IngestContext) -> bool: + del context + asset_ids = {asset.id for asset in bundle.assets} + return any( + isinstance(block, ImageBlock) and block.asset_id in asset_ids for block in bundle.blocks + ) + + def enrich(self, bundle: DocumentBundle, context: IngestContext) -> DocumentBundle: + assets = {asset.id: asset for asset in bundle.assets} + model = _image_model(context) + blocks: list[Block] = [] + changed = False + for block in bundle.blocks: + if not isinstance(block, ImageBlock): + blocks.append(block) + continue + asset = assets.get(block.asset_id) + if asset is None: + blocks.append(block) + continue + result = _describe_image(asset, model) + metadata = { + **block.metadata, + "visual_description_derived": True, + "visual_description_model": model, + } + if result.keywords: + metadata["keywords"] = result.keywords + if result.uncertainty: + metadata["uncertainty"] = result.uncertainty + blocks.append( + replace( + block, + visual_description=result.visual_description, + ocr_text=result.ocr_text, + metadata=metadata, + ) + ) + changed = True + if not changed: + return bundle + return replace( + bundle, + blocks=blocks, + metadata={**bundle.metadata, "image_vision_model": model}, + ) + + +def _describe_image(asset: Asset, model: str) -> ImageVisionResult: + import litellm + + image_bytes = asset.path.read_bytes() + encoded = base64.b64encode(image_bytes).decode("ascii") + data_url = f"data:{asset.media_type};base64,{encoded}" + response = litellm.completion( + model=model, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Describe this image for a knowledge base. Return JSON with " + "visual_description, ocr_text, keywords, and uncertainty. " + "Use null when a field is not visible or not applicable." + ), + }, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + ) + return _parse_vision_response(_response_content(response)) + + +def _parse_vision_response(content: str) -> ImageVisionResult: + cleaned = _strip_code_fence(content.strip()) + try: + payload = json.loads(cleaned) + except json.JSONDecodeError: + return ImageVisionResult(visual_description=cleaned or None) + if not isinstance(payload, dict): + return ImageVisionResult(visual_description=cleaned or None) + keywords = payload.get("keywords") + parsed_keywords = ( + [item for item in keywords if isinstance(item, str)] if isinstance(keywords, list) else None + ) + return ImageVisionResult( + visual_description=_optional_string(payload.get("visual_description")), + ocr_text=_optional_string(payload.get("ocr_text")), + keywords=parsed_keywords, + uncertainty=_optional_string(payload.get("uncertainty")), + ) + + +def _response_content(response: Any) -> str: + choices = _get(response, "choices") + if isinstance(choices, list) and choices: + message = _get(choices[0], "message") + content = _get(message, "content") + if isinstance(content, str): + return content + content = _get(response, "content") + return content if isinstance(content, str) else "" + + +def _get(value: Any, key: str) -> Any: + if isinstance(value, dict): + return value.get(key) + return getattr(value, key, None) + + +def _strip_code_fence(content: str) -> str: + if not content.startswith("```"): + return content + lines = content.splitlines() + if len(lines) >= 2 and lines[-1].strip() == "```": + return "\n".join(lines[1:-1]).strip() + return content + + +def _optional_string(value: object) -> str | None: + if not isinstance(value, str): + return None + cleaned = value.strip() + return cleaned or None + + +def _image_model(context: IngestContext) -> str: + ingest = context.config.get("ingest") + if isinstance(ingest, dict): + image_model = ingest.get("image_model") + if isinstance(image_model, str) and image_model.strip(): + return image_model.strip() + model = context.config.get("model") + return str(model or "gpt-4o-mini") diff --git a/openkb/ingest/enrichers/link_discovery.py b/openkb/ingest/enrichers/link_discovery.py new file mode 100644 index 00000000..32062f0e --- /dev/null +++ b/openkb/ingest/enrichers/link_discovery.py @@ -0,0 +1,70 @@ +"""Deterministic link discovery for bundle ingest.""" + +from __future__ import annotations + +import re +from dataclasses import replace + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DiscoveredLink, DocumentBundle, TableBlock, TextBlock + +_MARKDOWN_LINK_RE = re.compile(r"(?)\]]+") + + +class LinkDiscoveryEnricher: + name = "link_discovery" + + def applies_to(self, bundle: DocumentBundle, context: IngestContext) -> bool: + del context + return any(isinstance(block, (TextBlock, TableBlock)) for block in bundle.blocks) + + def enrich(self, bundle: DocumentBundle, context: IngestContext) -> DocumentBundle: + del context + discovered: list[DiscoveredLink] = list(bundle.links) + seen = {link.uri for link in discovered} + for block_index, text in enumerate(_text_parts(bundle)): + block_id = f"block-{block_index}" + for label, uri in _MARKDOWN_LINK_RE.findall(text): + cleaned = _clean_link_uri(uri) + if not cleaned or cleaned in seen: + continue + seen.add(cleaned) + discovered.append( + DiscoveredLink( + uri=cleaned, + text=label.strip() or None, + source_block_id=block_id, + metadata={"source": "markdown"}, + ) + ) + for match in _BARE_URL_RE.finditer(text): + cleaned = _clean_link_uri(match.group(0)) + if not cleaned or cleaned in seen: + continue + seen.add(cleaned) + discovered.append( + DiscoveredLink( + uri=cleaned, + source_block_id=block_id, + metadata={"source": "bare_url"}, + ) + ) + return replace(bundle, links=discovered) + + +def _text_parts(bundle: DocumentBundle) -> list[str]: + parts: list[str] = [] + for block in bundle.blocks: + if isinstance(block, TextBlock): + parts.append(block.text) + elif isinstance(block, TableBlock): + parts.append(block.markdown) + return parts + + +def _clean_link_uri(uri: str) -> str: + cleaned = uri.strip().strip("<>").strip().rstrip(".,;:") + if not cleaned or cleaned.startswith("#"): + return "" + return cleaned diff --git a/openkb/ingest/exceptions.py b/openkb/ingest/exceptions.py new file mode 100644 index 00000000..5c837e13 --- /dev/null +++ b/openkb/ingest/exceptions.py @@ -0,0 +1,14 @@ +"""Exceptions used to control ingest pipeline routing.""" + +from __future__ import annotations + +from pathlib import Path + + +class LongDocumentFallback(Exception): + """Signal that the target should be routed to the legacy long-doc path.""" + + def __init__(self, path: Path, reason: str) -> None: + super().__init__(reason) + self.path = path + self.reason = reason diff --git a/openkb/ingest/importers/__init__.py b/openkb/ingest/importers/__init__.py new file mode 100644 index 00000000..b4b20d95 --- /dev/null +++ b/openkb/ingest/importers/__init__.py @@ -0,0 +1,9 @@ +"""Built-in source importers.""" + +from __future__ import annotations + +from openkb.ingest.importers.feishu import FeishuImporter +from openkb.ingest.importers.file import FileImporter +from openkb.ingest.importers.url import UrlImporter + +__all__ = ["FeishuImporter", "FileImporter", "UrlImporter"] diff --git a/openkb/ingest/importers/feishu.py b/openkb/ingest/importers/feishu.py new file mode 100644 index 00000000..14eac5b5 --- /dev/null +++ b/openkb/ingest/importers/feishu.py @@ -0,0 +1,154 @@ +"""Feishu/Lark document source importer for bundle ingest.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import IngestInput +from openkb.locks import atomic_write_text +from openkb.url_ingest import _sanitize_filename, _unique_path, looks_like_url + +_FEISHU_DOMAINS = ("feishu.cn", "larksuite.com") +_FEISHU_DOC_PREFIXES = ("/wiki/", "/docx/", "/docs/") + + +class FeishuImporter: + name = "feishu" + + def can_handle(self, target: str, context: IngestContext) -> bool: + del context + return looks_like_feishu_url(target) + + def import_source(self, target: str, context: IngestContext) -> IngestInput: + markdown, title = _fetch_markdown_with_lark_cli(target, context) + raw_dir = context.staging_dir / "raw" + raw_dir.mkdir(parents=True, exist_ok=True) + filename = _sanitize_filename(title or _fallback_title(target), ".md") + path = _unique_path(raw_dir / filename) + atomic_write_text(path, markdown.rstrip() + "\n") + return IngestInput( + target=target, + path=path, + source_uri=target, + media_type="text/markdown", + metadata={ + "display_name": path.name, + "title": title, + "source_system": "feishu", + "auth_adapter": "lark-cli", + "permission_boundary": "current lark-cli identity", + }, + ) + + +def looks_like_feishu_url(target: str) -> bool: + if not looks_like_url(target): + return False + parsed = urlparse(target) + if not _host_matches(parsed.netloc, _FEISHU_DOMAINS): + return False + return any(parsed.path.startswith(prefix) for prefix in _FEISHU_DOC_PREFIXES) + + +def _fetch_markdown_with_lark_cli(target: str, context: IngestContext) -> tuple[str, str | None]: + config = _feishu_config(context) + cli = str(config.get("cli") or "lark-cli") + timeout = _int_config(config.get("timeout"), default=120) + command = [ + cli, + "docs", + "+fetch", + "--api-version", + "v2", + "--doc", + target, + "--doc-format", + "markdown", + "--format", + "json", + ] + try: + completed = subprocess.run( + command, + check=False, + capture_output=True, + text=True, + timeout=timeout, + ) + except FileNotFoundError as exc: + raise ValueError("Feishu importer requires lark-cli on PATH.") from exc + except subprocess.TimeoutExpired as exc: + raise ValueError(f"Feishu fetch timed out after {timeout}s: {target}") from exc + + if completed.returncode != 0: + detail = ( + completed.stderr.strip() or completed.stdout.strip() or f"exit {completed.returncode}" + ) + raise ValueError(f"Feishu fetch failed: {detail}") + + markdown, title = _parse_lark_cli_output(completed.stdout) + if not markdown.strip(): + raise ValueError("Feishu fetch returned empty markdown.") + return markdown, title + + +def _parse_lark_cli_output(stdout: str) -> tuple[str, str | None]: + try: + payload = json.loads(stdout) + except json.JSONDecodeError: + return stdout, None + markdown = _find_first_string(payload, ("markdown", "content", "text", "body")) or "" + title = _find_first_string(payload, ("title", "name")) + return markdown, title + + +def _find_first_string(value: Any, keys: tuple[str, ...]) -> str | None: + if isinstance(value, dict): + for key in keys: + item = value.get(key) + if isinstance(item, str) and item.strip(): + return item + for item in value.values(): + found = _find_first_string(item, keys) + if found is not None: + return found + elif isinstance(value, list): + for item in value: + found = _find_first_string(item, keys) + if found is not None: + return found + return None + + +def _feishu_config(context: IngestContext) -> dict[str, Any]: + ingest = context.config.get("ingest") + if not isinstance(ingest, dict): + return {} + feishu = ingest.get("feishu") + return feishu if isinstance(feishu, dict) else {} + + +def _int_config(value: object, *, default: int) -> int: + if isinstance(value, bool) or not isinstance(value, (int, str)): + return default + try: + parsed = int(value) + except ValueError: + return default + return max(1, parsed) + + +def _fallback_title(target: str) -> str: + parsed = urlparse(target) + name = Path(parsed.path.rstrip("/")).name + return name or parsed.netloc or "feishu-document" + + +def _host_matches(netloc: str, domains: tuple[str, ...]) -> bool: + host = netloc.rsplit("@", 1)[-1].split(":", 1)[0].lower() + return any(host == domain or host.endswith(f".{domain}") for domain in domains) diff --git a/openkb/ingest/importers/file.py b/openkb/ingest/importers/file.py new file mode 100644 index 00000000..644097c3 --- /dev/null +++ b/openkb/ingest/importers/file.py @@ -0,0 +1,54 @@ +"""Local filesystem source importer.""" + +from __future__ import annotations + +from pathlib import Path + +from openkb.converter import _registry_path +from openkb.ingest.context import IngestContext +from openkb.ingest.models import IngestInput + + +class FileImporter: + name = "file" + + def can_handle(self, target: str, context: IngestContext) -> bool: + del context + return Path(target).expanduser().exists() + + def import_source(self, target: str, context: IngestContext) -> IngestInput: + path = Path(target).expanduser().resolve() + if not path.is_file(): + raise ValueError(f"Bundle file importer requires a file: {target}") + return IngestInput( + target=target, + path=path, + source_uri=_registry_path(path, context.kb_dir), + media_type=_media_type_for(path), + metadata={"display_name": path.name}, + ) + + +def _media_type_for(path: Path) -> str: + suffix = path.suffix.lower() + if suffix in {".md", ".markdown"}: + return "text/markdown" + if suffix == ".txt": + return "text/plain" + if suffix == ".pdf": + return "application/pdf" + if suffix in {".html", ".htm"}: + return "text/html" + if suffix == ".csv": + return "text/csv" + if suffix == ".png": + return "image/png" + if suffix in {".jpg", ".jpeg"}: + return "image/jpeg" + if suffix == ".gif": + return "image/gif" + if suffix == ".webp": + return "image/webp" + if suffix == ".bmp": + return "image/bmp" + return "application/octet-stream" diff --git a/openkb/ingest/importers/url.py b/openkb/ingest/importers/url.py new file mode 100644 index 00000000..b38b4b86 --- /dev/null +++ b/openkb/ingest/importers/url.py @@ -0,0 +1,41 @@ +"""HTTP(S) source importer for bundle ingest.""" + +from __future__ import annotations + +from pathlib import Path + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import IngestInput +from openkb.url_ingest import fetch_url_to_dir, looks_like_url + + +class UrlImporter: + name = "url" + + def can_handle(self, target: str, context: IngestContext) -> bool: + del context + return looks_like_url(target) + + def import_source(self, target: str, context: IngestContext) -> IngestInput: + raw_dir = context.staging_dir / "raw" + path = fetch_url_to_dir(target, raw_dir) + if path is None: + raise ValueError(f"URL fetch failed: {target}") + return IngestInput( + target=target, + path=path, + source_uri=target, + media_type=_media_type_for(path), + metadata={"display_name": path.name}, + ) + + +def _media_type_for(path: Path) -> str: + suffix = path.suffix.lower() + if suffix in {".md", ".markdown"}: + return "text/markdown" + if suffix == ".pdf": + return "application/pdf" + if suffix in {".html", ".htm"}: + return "text/html" + return "application/octet-stream" diff --git a/openkb/ingest/models.py b/openkb/ingest/models.py new file mode 100644 index 00000000..59e77b48 --- /dev/null +++ b/openkb/ingest/models.py @@ -0,0 +1,108 @@ +"""Data models for the bundle ingest pipeline.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + + +@dataclass(frozen=True) +class IngestInput: + """Raw source payload produced by a source importer.""" + + target: str + path: Path | None = None + source_uri: str | None = None + media_type: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class TextBlock: + text: str + page: int | None = None + anchor: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ImageBlock: + asset_id: str + caption: str | None = None + ocr_text: str | None = None + visual_description: str | None = None + page: int | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class TableBlock: + markdown: str + page: int | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class EmbedBlock: + uri: str + title: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +Block = TextBlock | ImageBlock | TableBlock | EmbedBlock + + +@dataclass(frozen=True) +class Asset: + id: str + path: Path + media_type: str + sha256: str + source_uri: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DiscoveredLink: + uri: str + text: str | None = None + source_block_id: str | None = None + confidence: float = 1.0 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ProvenanceRecord: + source_uri: str + relationship: str = "root" + parent_uri: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DocumentBundle: + id: str + title: str | None + source_uri: str + blocks: list[Block] + assets: list[Asset] = field(default_factory=list) + links: list[DiscoveredLink] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + provenance: list[ProvenanceRecord] = field(default_factory=list) + + +@dataclass(frozen=True) +class RenderedBundle: + """Staged artifacts compatible with OpenKB's existing add commit path.""" + + doc_name: str + display_name: str + source_identity: str + content_hash: str + raw_path: Path | None + source_path: Path | None + doc_type: str + bundle_path: Path | None = None + is_long_doc: bool = False + kind: Literal["short", "long"] = "short" diff --git a/openkb/ingest/normalizers/__init__.py b/openkb/ingest/normalizers/__init__.py new file mode 100644 index 00000000..d62c808d --- /dev/null +++ b/openkb/ingest/normalizers/__init__.py @@ -0,0 +1,18 @@ +"""Built-in bundle normalizers.""" + +from __future__ import annotations + +from openkb.ingest.normalizers.image import IMAGE_EXTENSIONS, ImageNormalizer +from openkb.ingest.normalizers.markdown import MarkdownNormalizer +from openkb.ingest.normalizers.markitdown import MarkItDownNormalizer +from openkb.ingest.normalizers.pdf import PdfNormalizer +from openkb.ingest.normalizers.text import PlainTextNormalizer + +__all__ = [ + "IMAGE_EXTENSIONS", + "ImageNormalizer", + "MarkdownNormalizer", + "MarkItDownNormalizer", + "PdfNormalizer", + "PlainTextNormalizer", +] diff --git a/openkb/ingest/normalizers/image.py b/openkb/ingest/normalizers/image.py new file mode 100644 index 00000000..83e0880e --- /dev/null +++ b/openkb/ingest/normalizers/image.py @@ -0,0 +1,62 @@ +"""Image normalizer for local image files.""" + +from __future__ import annotations + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import Asset, DocumentBundle, ImageBlock, IngestInput, ProvenanceRecord +from openkb.state import HashRegistry + +IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"} + + +class ImageNormalizer: + name = "image" + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: + del context + return input_.path is not None and input_.path.suffix.lower() in IMAGE_EXTENSIONS + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + del context + if input_.path is None: + raise ValueError("Image normalizer requires a local file path.") + source_uri = input_.source_uri or input_.path.as_posix() + media_type = input_.media_type or _media_type_for_suffix(input_.path.suffix) + image_hash = HashRegistry.hash_file(input_.path) + asset = Asset( + id="source-image", + path=input_.path, + media_type=media_type, + sha256=image_hash, + source_uri=source_uri, + ) + return DocumentBundle( + id=source_uri, + title=input_.path.stem, + source_uri=source_uri, + blocks=[ImageBlock(asset_id=asset.id, caption=input_.path.stem)], + assets=[asset], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_suffix": input_.path.suffix.lower(), + "source_identity": source_uri, + "render_strategy": "image", + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) + + +def _media_type_for_suffix(suffix: str) -> str: + normalized = suffix.lower() + if normalized in {".jpg", ".jpeg"}: + return "image/jpeg" + if normalized == ".png": + return "image/png" + if normalized == ".gif": + return "image/gif" + if normalized == ".webp": + return "image/webp" + if normalized == ".bmp": + return "image/bmp" + return "application/octet-stream" diff --git a/openkb/ingest/normalizers/markdown.py b/openkb/ingest/normalizers/markdown.py new file mode 100644 index 00000000..c8be9ee7 --- /dev/null +++ b/openkb/ingest/normalizers/markdown.py @@ -0,0 +1,45 @@ +"""Markdown normalizer for local files.""" + +from __future__ import annotations + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord, TextBlock + + +class MarkdownNormalizer: + name = "markdown" + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: + del context + if input_.path is None: + return False + return input_.path.suffix.lower() in {".md", ".markdown"} + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + del context + if input_.path is None: + raise ValueError("Markdown normalizer requires a local file path.") + source_uri = input_.source_uri or input_.path.as_posix() + input_metadata = dict(input_.metadata) + title = input_metadata.get("title") + if not isinstance(title, str) or not title.strip(): + title = input_.path.stem + return DocumentBundle( + id=source_uri, + title=title.strip(), + source_uri=source_uri, + blocks=[ + TextBlock( + text=input_.path.read_text(encoding="utf-8"), + metadata={"source_dir": input_.path.parent.as_posix()}, + ) + ], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_suffix": input_.path.suffix.lower(), + "source_identity": source_uri, + "source_metadata": input_metadata, + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) diff --git a/openkb/ingest/normalizers/markitdown.py b/openkb/ingest/normalizers/markitdown.py new file mode 100644 index 00000000..3ac29eab --- /dev/null +++ b/openkb/ingest/normalizers/markitdown.py @@ -0,0 +1,44 @@ +"""MarkItDown-backed normalizer for Office, HTML, and CSV-like files.""" + +from __future__ import annotations + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord + +_MARKITDOWN_EXTENSIONS = { + ".docx", + ".pptx", + ".xlsx", + ".xls", + ".html", + ".htm", + ".csv", +} + + +class MarkItDownNormalizer: + name = "markitdown" + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: + del context + return input_.path is not None and input_.path.suffix.lower() in _MARKITDOWN_EXTENSIONS + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + del context + if input_.path is None: + raise ValueError("MarkItDown normalizer requires a local file path.") + source_uri = input_.source_uri or input_.path.as_posix() + return DocumentBundle( + id=source_uri, + title=input_.path.stem, + source_uri=source_uri, + blocks=[], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_suffix": input_.path.suffix.lower(), + "source_identity": source_uri, + "render_strategy": "markitdown", + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) diff --git a/openkb/ingest/normalizers/pdf.py b/openkb/ingest/normalizers/pdf.py new file mode 100644 index 00000000..4a14a94e --- /dev/null +++ b/openkb/ingest/normalizers/pdf.py @@ -0,0 +1,42 @@ +"""PDF normalizer for short local PDFs.""" + +from __future__ import annotations + +from openkb.converter import get_pdf_page_count +from openkb.ingest.context import IngestContext +from openkb.ingest.exceptions import LongDocumentFallback +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord + + +class PdfNormalizer: + name = "pdf-local" + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: + del context + return input_.path is not None and input_.path.suffix.lower() == ".pdf" + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + if input_.path is None: + raise ValueError("PDF normalizer requires a local file path.") + threshold = int(context.config.get("pageindex_threshold", 20)) + page_count = get_pdf_page_count(input_.path) + if page_count >= threshold: + raise LongDocumentFallback( + input_.path, + f"PDF has {page_count} pages >= pageindex_threshold {threshold}", + ) + source_uri = input_.source_uri or input_.path.as_posix() + return DocumentBundle( + id=source_uri, + title=input_.path.stem, + source_uri=source_uri, + blocks=[], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_suffix": ".pdf", + "source_identity": source_uri, + "render_strategy": "pdf", + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) diff --git a/openkb/ingest/normalizers/text.py b/openkb/ingest/normalizers/text.py new file mode 100644 index 00000000..41362c25 --- /dev/null +++ b/openkb/ingest/normalizers/text.py @@ -0,0 +1,33 @@ +"""Plain text normalizer for local files.""" + +from __future__ import annotations + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord, TextBlock + + +class PlainTextNormalizer: + name = "text" + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: + del context + return input_.path is not None and input_.path.suffix.lower() == ".txt" + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + del context + if input_.path is None: + raise ValueError("Text normalizer requires a local file path.") + source_uri = input_.source_uri or input_.path.as_posix() + return DocumentBundle( + id=source_uri, + title=input_.path.stem, + source_uri=source_uri, + blocks=[TextBlock(text=input_.path.read_text(encoding="utf-8"))], + metadata={ + "display_name": input_.metadata.get("display_name", input_.path.name), + "source_path": input_.path.as_posix(), + "source_suffix": input_.path.suffix.lower(), + "source_identity": source_uri, + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) diff --git a/openkb/ingest/plugins.py b/openkb/ingest/plugins.py new file mode 100644 index 00000000..d9b761c9 --- /dev/null +++ b/openkb/ingest/plugins.py @@ -0,0 +1,60 @@ +"""Entry point loading for external ingest plugins.""" + +from __future__ import annotations + +from importlib.metadata import entry_points +from typing import Any + +IMPORTER_ENTRY_POINT_GROUP = "openkb.ingest.importers" +NORMALIZER_ENTRY_POINT_GROUP = "openkb.ingest.normalizers" +ENRICHER_ENTRY_POINT_GROUP = "openkb.ingest.enrichers" + + +def load_ingest_entry_points( + group: str, + enabled_names: tuple[str, ...], + *, + exclude: set[str] | None = None, +) -> list[Any]: + """Load configured external ingest components from Python entry points.""" + enabled = set(enabled_names) + if not enabled: + return [] + excluded = exclude or set() + components: list[Any] = [] + for entry_point in _entry_points_for(group): + if entry_point.name not in enabled or entry_point.name in excluded: + continue + try: + loaded = entry_point.load() + components.append(_component_instance(loaded, group, entry_point.name)) + except Exception as exc: + raise ValueError( + f"Could not load ingest plugin {entry_point.name!r} from {group}: {exc}" + ) from exc + return components + + +def _entry_points_for(group: str) -> list[Any]: + points = entry_points() + if hasattr(points, "select"): + return list(points.select(group=group)) + return list(points.get(group, [])) + + +def _component_instance(loaded: Any, group: str, name: str) -> Any: + if isinstance(loaded, type): + instance = loaded() + if _looks_like_component(instance): + return instance + if _looks_like_component(loaded): + return loaded + if callable(loaded): + instance = loaded() + if _looks_like_component(instance): + return instance + raise TypeError(f"entry point {name!r} in {group} did not produce an ingest component") + + +def _looks_like_component(value: Any) -> bool: + return hasattr(value, "name") diff --git a/openkb/ingest/registry.py b/openkb/ingest/registry.py new file mode 100644 index 00000000..07e7a5e4 --- /dev/null +++ b/openkb/ingest/registry.py @@ -0,0 +1,52 @@ +"""Small in-process registries for built-in ingest components.""" + +from __future__ import annotations + +from typing import Protocol + +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DocumentBundle, IngestInput + + +class SourceImporter(Protocol): + name: str + + def can_handle(self, target: str, context: IngestContext) -> bool: ... + def import_source(self, target: str, context: IngestContext) -> IngestInput: ... + + +class BundleNormalizer(Protocol): + name: str + + def supports(self, input_: IngestInput, context: IngestContext) -> bool: ... + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: ... + + +class BundleEnricher(Protocol): + name: str + + def applies_to(self, bundle: DocumentBundle, context: IngestContext) -> bool: ... + def enrich(self, bundle: DocumentBundle, context: IngestContext) -> DocumentBundle: ... + + +class SourceImporterRegistry: + def __init__(self, importers: list[SourceImporter]) -> None: + self._importers = importers + + def resolve(self, target: str, context: IngestContext) -> SourceImporter: + for importer in self._importers: + if importer.can_handle(target, context): + return importer + raise ValueError(f"No bundle source importer can handle: {target}") + + +class BundleNormalizerRegistry: + def __init__(self, normalizers: list[BundleNormalizer]) -> None: + self._normalizers = normalizers + + def normalize(self, input_: IngestInput, context: IngestContext) -> DocumentBundle: + for normalizer in self._normalizers: + if normalizer.supports(input_, context): + return normalizer.normalize(input_, context) + source = input_.path or input_.source_uri or input_.target + raise ValueError(f"No bundle normalizer can handle: {source}") diff --git a/openkb/ingest/render.py b/openkb/ingest/render.py new file mode 100644 index 00000000..4a16ada5 --- /dev/null +++ b/openkb/ingest/render.py @@ -0,0 +1,121 @@ +"""Render document bundles into staged OpenKB-compatible artifacts.""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +from markitdown import MarkItDown + +from openkb.converter import _registry_path, resolve_doc_name +from openkb.images import convert_pdf_with_images, copy_relative_images, extract_base64_images +from openkb.ingest.context import IngestContext +from openkb.ingest.models import DocumentBundle, ImageBlock, RenderedBundle, TableBlock, TextBlock +from openkb.ingest.serialization import write_bundle_json +from openkb.locks import atomic_write_text +from openkb.state import HashRegistry + + +def render_bundle_to_staging(bundle: DocumentBundle, context: IngestContext) -> RenderedBundle: + """Render a short-document bundle into staged raw/source artifacts.""" + source_path = _bundle_source_path(bundle) + registry = HashRegistry(context.kb_dir / ".openkb" / "hashes.json") + file_hash = HashRegistry.hash_file(source_path) + doc_name = resolve_doc_name( + source_path, + context.kb_dir, + registry, + persist_legacy=False, + ) + + raw_dir = context.staging_dir / "raw" + raw_dir.mkdir(parents=True, exist_ok=True) + raw_dest = raw_dir / f"{doc_name}{source_path.suffix.lower()}" + if source_path.resolve() != raw_dest.resolve(): + shutil.copy2(source_path, raw_dest) + + images_dir = context.staging_dir / "wiki" / "sources" / "images" / doc_name + strategy = bundle.metadata.get("render_strategy") + if strategy == "pdf": + markdown = convert_pdf_with_images(source_path, doc_name, images_dir) + elif strategy == "markitdown": + mid = MarkItDown() + result = mid.convert(str(source_path), keep_data_uris=True) + markdown = extract_base64_images(result.text_content, doc_name, images_dir) + elif strategy == "image": + markdown = _render_image_markdown(bundle, source_path, doc_name, images_dir) + else: + markdown = _render_short_markdown(bundle) + markdown = copy_relative_images(markdown, source_path.parent, doc_name, images_dir) + + sources_dir = context.staging_dir / "wiki" / "sources" + sources_dir.mkdir(parents=True, exist_ok=True) + source_dest = sources_dir / f"{doc_name}.md" + atomic_write_text(source_dest, markdown) + + bundles_dir = context.staging_dir / ".openkb" / "bundles" + bundles_dir.mkdir(parents=True, exist_ok=True) + bundle_dest = bundles_dir / f"{doc_name}.json" + write_bundle_json(bundle, bundle_dest) + + return RenderedBundle( + doc_name=doc_name, + display_name=str(bundle.metadata.get("display_name") or source_path.name), + source_identity=str( + bundle.metadata.get("source_identity") or _registry_path(source_path, context.kb_dir) + ), + content_hash=file_hash, + raw_path=raw_dest, + source_path=source_dest, + doc_type=source_path.suffix.lower().lstrip(".") or "md", + bundle_path=bundle_dest, + ) + + +def _bundle_source_path(bundle: DocumentBundle) -> Path: + raw = bundle.metadata.get("source_path") + if not isinstance(raw, str) or not raw: + raise ValueError("Bundle renderer requires metadata['source_path'].") + path = Path(raw) + if not path.is_file(): + raise ValueError(f"Bundle source path does not exist: {path}") + return path + + +def _render_short_markdown(bundle: DocumentBundle) -> str: + parts: list[str] = [] + for block in bundle.blocks: + if isinstance(block, TextBlock): + parts.append(block.text) + elif isinstance(block, TableBlock): + parts.append(block.markdown) + return "\n\n".join(part.strip("\n") for part in parts if part).rstrip() + "\n" + + +def _render_image_markdown( + bundle: DocumentBundle, + source_path: Path, + doc_name: str, + images_dir: Path, +) -> str: + images_dir.mkdir(parents=True, exist_ok=True) + image_dest = images_dir / source_path.name + if source_path.resolve() != image_dest.resolve(): + shutil.copy2(source_path, image_dest) + title = bundle.title or source_path.stem + parts = [ + f"# {title}", + f"![source image](sources/images/{doc_name}/{image_dest.name})", + ] + for block in bundle.blocks: + if not isinstance(block, ImageBlock): + continue + if block.visual_description: + parts.append( + "## 视觉描述\n\n" + "> 模型派生信息，可能不完整或不准确。\n\n" + f"{block.visual_description.strip()}" + ) + if block.ocr_text: + parts.append(f"## 可见文字\n\n{block.ocr_text.strip()}") + return "\n\n".join(parts).rstrip() + "\n" diff --git a/openkb/ingest/serialization.py b/openkb/ingest/serialization.py new file mode 100644 index 00000000..ba816049 --- /dev/null +++ b/openkb/ingest/serialization.py @@ -0,0 +1,116 @@ +"""Serialization helpers for persisted ingest bundle sidecars.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from openkb.ingest.models import ( + Asset, + DiscoveredLink, + DocumentBundle, + EmbedBlock, + ImageBlock, + ProvenanceRecord, + TableBlock, + TextBlock, +) +from openkb.locks import atomic_write_json + +_SCHEMA_VERSION = 1 + + +def write_bundle_json(bundle: DocumentBundle, path: Path) -> None: + atomic_write_json(path, bundle_to_dict(bundle)) + + +def bundle_to_dict(bundle: DocumentBundle) -> dict[str, Any]: + return { + "schema_version": _SCHEMA_VERSION, + "id": bundle.id, + "title": bundle.title, + "source_uri": bundle.source_uri, + "blocks": [_block_to_dict(block) for block in bundle.blocks], + "assets": [_asset_to_dict(asset) for asset in bundle.assets], + "links": [_link_to_dict(link) for link in bundle.links], + "metadata": _jsonable(bundle.metadata), + "provenance": [_provenance_to_dict(record) for record in bundle.provenance], + } + + +def _block_to_dict(block) -> dict[str, Any]: + if isinstance(block, TextBlock): + return { + "type": "text", + "text": block.text, + "page": block.page, + "anchor": block.anchor, + "metadata": _jsonable(block.metadata), + } + if isinstance(block, ImageBlock): + return { + "type": "image", + "asset_id": block.asset_id, + "caption": block.caption, + "ocr_text": block.ocr_text, + "visual_description": block.visual_description, + "page": block.page, + "metadata": _jsonable(block.metadata), + } + if isinstance(block, TableBlock): + return { + "type": "table", + "markdown": block.markdown, + "page": block.page, + "metadata": _jsonable(block.metadata), + } + if isinstance(block, EmbedBlock): + return { + "type": "embed", + "uri": block.uri, + "title": block.title, + "metadata": _jsonable(block.metadata), + } + raise TypeError(f"Unsupported bundle block: {type(block).__name__}") + + +def _asset_to_dict(asset: Asset) -> dict[str, Any]: + return { + "id": asset.id, + "path": asset.path.as_posix(), + "media_type": asset.media_type, + "sha256": asset.sha256, + "source_uri": asset.source_uri, + "metadata": _jsonable(asset.metadata), + } + + +def _link_to_dict(link: DiscoveredLink) -> dict[str, Any]: + return { + "uri": link.uri, + "text": link.text, + "source_block_id": link.source_block_id, + "confidence": link.confidence, + "metadata": _jsonable(link.metadata), + } + + +def _provenance_to_dict(record: ProvenanceRecord) -> dict[str, Any]: + return { + "source_uri": record.source_uri, + "relationship": record.relationship, + "parent_uri": record.parent_uri, + "metadata": _jsonable(record.metadata), + } + + +def _jsonable(value: Any) -> Any: + if isinstance(value, Path): + return value.as_posix() + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) diff --git a/openkb/mutation.py b/openkb/mutation.py index 24c91957..02167202 100644 --- a/openkb/mutation.py +++ b/openkb/mutation.py @@ -448,7 +448,7 @@ def publish_staged_tree(staging_dir: Path | None, kb_dir: Path) -> None: """Move staged raw/source artifacts into their final KB locations.""" if staging_dir is None or not staging_dir.exists(): return - for rel in ("raw", "wiki/sources"): + for rel in ("raw", "wiki/sources", ".openkb/bundles"): src_root = staging_dir / rel if not src_root.exists(): continue diff --git a/openkb/url_ingest.py b/openkb/url_ingest.py index ccd56d94..6250c2d2 100644 --- a/openkb/url_ingest.py +++ b/openkb/url_ingest.py @@ -213,8 +213,8 @@ def _extract_html(url: str, raw_dir: Path) -> Path | None: return target -def fetch_url_to_raw(url: str, kb_dir: Path) -> Path | None: - """Fetch ``url`` into ``/raw/`` and return the local path. +def fetch_url_to_dir(url: str, raw_dir: Path) -> Path | None: + """Fetch ``url`` into ``raw_dir`` and return the local path. Routing is decided by HTTP ``Content-Type`` validated against magic bytes (in case the server lies): @@ -223,11 +223,9 @@ def fetch_url_to_raw(url: str, kb_dir: Path) -> Path | None: - HTML → trafilatura main-content extract → ``raw/.md`` - anything else → error, returns None - The caller then hands the saved path to ``add_single_file``, so the - existing PageIndex / markitdown routing by file extension and page - count takes over from there. + The caller decides whether this raw directory is the live KB + ``raw/`` path or a mutation staging directory. """ - raw_dir = kb_dir / "raw" raw_dir.mkdir(parents=True, exist_ok=True) click.echo(f"Downloading: {url}") @@ -280,3 +278,13 @@ def fetch_url_to_raw(url: str, kb_dir: Path) -> Path | None: err=True, ) return None + + +def fetch_url_to_raw(url: str, kb_dir: Path) -> Path | None: + """Fetch ``url`` into ``/raw/`` and return the local path. + + The caller then hands the saved path to ``add_single_file``, so the + existing PageIndex / markitdown routing by file extension and page + count takes over from there. + """ + return fetch_url_to_dir(url, kb_dir / "raw") diff --git a/tests/test_ingest_boundaries.py b/tests/test_ingest_boundaries.py new file mode 100644 index 00000000..ea4aaf10 --- /dev/null +++ b/tests/test_ingest_boundaries.py @@ -0,0 +1,520 @@ +"""Boundary coverage for the bundle ingest internals.""" + +from __future__ import annotations + +import json +import subprocess +from dataclasses import dataclass +from pathlib import Path +from unittest.mock import patch + +import click +import pytest + +from openkb.ingest.add import ( + _add_one_bundle_target_locked, + _allowed_recursive_targets, + _BundleAddState, + _domain_allowed, + _summarize_outcomes, + _with_parent_provenance, + resolve_ingest_pipeline, +) +from openkb.ingest.config import resolve_ingest_options +from openkb.ingest.context import IngestContext +from openkb.ingest.enrichers.image_vision import ( + ImageVisionEnricher, + _image_model, + _parse_vision_response, + _response_content, +) +from openkb.ingest.enrichers.link_discovery import LinkDiscoveryEnricher +from openkb.ingest.importers.feishu import ( + FeishuImporter, + _fallback_title, + _fetch_markdown_with_lark_cli, + _find_first_string, + _int_config, + _parse_lark_cli_output, + looks_like_feishu_url, +) +from openkb.ingest.importers.file import FileImporter +from openkb.ingest.importers.file import _media_type_for as file_media_type_for +from openkb.ingest.importers.url import UrlImporter +from openkb.ingest.importers.url import _media_type_for as url_media_type_for +from openkb.ingest.models import ( + Asset, + DiscoveredLink, + DocumentBundle, + EmbedBlock, + ImageBlock, + IngestInput, + ProvenanceRecord, + TableBlock, + TextBlock, +) +from openkb.ingest.normalizers.image import ImageNormalizer, _media_type_for_suffix +from openkb.ingest.normalizers.markdown import MarkdownNormalizer +from openkb.ingest.normalizers.markitdown import MarkItDownNormalizer +from openkb.ingest.normalizers.pdf import PdfNormalizer +from openkb.ingest.normalizers.text import PlainTextNormalizer +from openkb.ingest.plugins import load_ingest_entry_points +from openkb.ingest.registry import BundleNormalizerRegistry, SourceImporterRegistry +from openkb.ingest.render import _bundle_source_path, _render_image_markdown +from openkb.ingest.serialization import bundle_to_dict + + +def _context(tmp_path: Path, config: dict | None = None) -> IngestContext: + return IngestContext(kb_dir=tmp_path, config=config or {}, staging_dir=tmp_path / "stage") + + +def _completed(stdout: str = "", stderr: str = "", returncode: int = 0): + return type( + "Completed", + (), + {"stdout": stdout, "stderr": stderr, "returncode": returncode}, + )() + + +def test_resolve_ingest_pipeline_rejects_invalid_value(): + with pytest.raises(click.BadParameter): + resolve_ingest_pipeline({"ingest": {"pipeline": "bogus"}}) + + +def test_resolve_ingest_options_handles_invalid_shapes(): + options = resolve_ingest_options( + { + "ingest": { + "max_depth": True, + "max_documents": "0", + "allow_domains": "example.com", + "importers": {"enabled": ["file", 7, "file", "url"]}, + "normalizers": {"enabled": ["plugin"]}, + "enrichers": {"enabled": ["link_discovery"]}, + } + } + ) + + assert options.max_depth == 0 + assert options.max_documents == 1 + assert options.allow_domains == () + assert options.enabled_importers == ("file", "url") + assert options.enabled_normalizers == ("plugin",) + assert options.link_discovery is True + + +def test_source_and_normalizer_registries_error_when_unhandled(tmp_path): + context = _context(tmp_path) + with pytest.raises(ValueError, match="No bundle source importer"): + SourceImporterRegistry([]).resolve("missing", context) + with pytest.raises(ValueError, match="No bundle normalizer"): + BundleNormalizerRegistry([]).normalize(IngestInput(target="x"), context) + + +def test_add_one_skips_duplicate_target(tmp_path): + state = _BundleAddState(seen_targets={str(tmp_path / "doc.md")}) + outcome, bundle = _add_one_bundle_target_locked( + str(tmp_path / "doc.md"), + tmp_path, + "model", + resolve_ingest_options({}), + state, + ) + assert outcome == "skipped" + assert bundle is None + + +def test_add_one_errors_when_importer_returns_no_path(tmp_path): + (tmp_path / ".openkb").mkdir() + (tmp_path / ".openkb" / "hashes.json").write_text("{}", encoding="utf-8") + (tmp_path / ".openkb" / "config.yaml").write_text( + "ingest:\n importers:\n enabled:\n - virtual\n", + encoding="utf-8", + ) + + class VirtualImporter: + name = "virtual" + + def can_handle(self, target, context): + del target, context + return True + + def import_source(self, target, context): + del context + return IngestInput(target=target, source_uri=target) + + entry_points = _FakeEntryPoints( + [_FakeEntryPoint("openkb.ingest.importers", "virtual", VirtualImporter)] + ) + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.plugins.entry_points", return_value=entry_points), + ): + outcome, bundle = _add_one_bundle_target_locked( + "virtual://doc", + tmp_path, + "model", + resolve_ingest_options({"ingest": {"importers": {"enabled": ["virtual"]}}}), + _BundleAddState(seen_targets=set()), + ) + + assert outcome == "failed" + assert bundle is None + + +def test_add_one_skips_known_hash(tmp_path): + kb = tmp_path + (kb / ".openkb").mkdir() + (kb / ".openkb" / "hashes.json").write_text("{}", encoding="utf-8") + doc = kb / "known.md" + doc.write_text("# Known\n", encoding="utf-8") + + from openkb.state import HashRegistry + + registry = HashRegistry(kb / ".openkb" / "hashes.json") + registry.add(HashRegistry.hash_file(doc), {"name": "known.md"}) + + outcome, bundle = _add_one_bundle_target_locked( + str(doc), + kb, + "model", + resolve_ingest_options({}), + _BundleAddState(seen_targets=set()), + ) + assert outcome == "skipped" + assert bundle is None + + +def test_with_parent_provenance_adds_record_when_missing(): + bundle = DocumentBundle(id="child", title=None, source_uri="child", blocks=[]) + out = _with_parent_provenance(bundle, "parent") + assert out.provenance[0].relationship == "child" + assert out.provenance[0].parent_uri == "parent" + + +def test_recursive_target_boundary_cases(tmp_path): + bundle = DocumentBundle( + id="root", + title=None, + source_uri="root", + blocks=[], + metadata={"source_path": str(tmp_path / "root.md")}, + links=[ + DiscoveredLink("mailto:a@example.com"), + DiscoveredLink("missing.md"), + DiscoveredLink("https://example.com/page#frag"), + ], + ) + options = resolve_ingest_options( + { + "ingest": { + "allow_domains": ["example.com"], + "importers": {"enabled": ["file", "url"]}, + } + } + ) + assert _allowed_recursive_targets(bundle, options) == ["https://example.com/page"] + + no_file_options = resolve_ingest_options({"ingest": {"importers": {"enabled": []}}}) + assert ( + _allowed_recursive_targets( + DocumentBundle( + id="root", + title=None, + source_uri="root", + blocks=[], + links=[DiscoveredLink("")], + ), + no_file_options, + ) + == [] + ) + assert _domain_allowed("", ()) is False + assert _summarize_outcomes(["skipped", "failed"]) == "failed" + + +def test_file_importer_directory_error_and_media_types(tmp_path): + context = _context(tmp_path) + assert FileImporter().can_handle(str(tmp_path), context) + with pytest.raises(ValueError, match="requires a file"): + FileImporter().import_source(str(tmp_path), context) + + expected = { + "a.markdown": "text/markdown", + "a.txt": "text/plain", + "a.pdf": "application/pdf", + "a.htm": "text/html", + "a.csv": "text/csv", + "a.png": "image/png", + "a.jpeg": "image/jpeg", + "a.gif": "image/gif", + "a.webp": "image/webp", + "a.bmp": "image/bmp", + "a.bin": "application/octet-stream", + } + for name, media_type in expected.items(): + assert file_media_type_for(Path(name)) == media_type + + +def test_url_importer_failure_and_media_types(tmp_path): + context = _context(tmp_path) + assert UrlImporter().can_handle("https://example.com", context) + with patch("openkb.ingest.importers.url.fetch_url_to_dir", return_value=None): + with pytest.raises(ValueError, match="URL fetch failed"): + UrlImporter().import_source("https://example.com", context) + + assert url_media_type_for(Path("x.markdown")) == "text/markdown" + assert url_media_type_for(Path("x.pdf")) == "application/pdf" + assert url_media_type_for(Path("x.htm")) == "text/html" + assert url_media_type_for(Path("x.bin")) == "application/octet-stream" + + +def test_image_normalizer_boundaries(tmp_path): + context = _context(tmp_path) + normalizer = ImageNormalizer() + assert not normalizer.supports(IngestInput(target="x", path=None), context) + with pytest.raises(ValueError, match="requires a local file path"): + normalizer.normalize(IngestInput(target="x"), context) + + for suffix, media_type in { + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", + ".tif": "application/octet-stream", + }.items(): + assert _media_type_for_suffix(suffix) == media_type + + +def test_other_normalizer_error_branches(tmp_path): + context = _context(tmp_path) + for normalizer in [ + MarkdownNormalizer(), + PlainTextNormalizer(), + PdfNormalizer(), + MarkItDownNormalizer(), + ]: + with pytest.raises(ValueError, match="requires a local file path"): + normalizer.normalize(IngestInput(target="x"), context) + + +def test_link_discovery_noop_and_trailing_punctuation(tmp_path): + context = _context(tmp_path) + bundle = DocumentBundle( + id="b", + title=None, + source_uri="b", + blocks=[ImageBlock(asset_id="missing")], + ) + assert not LinkDiscoveryEnricher().applies_to(bundle, context) + + enriched = LinkDiscoveryEnricher().enrich( + DocumentBundle( + id="b", + title=None, + source_uri="b", + blocks=[TableBlock("See https://example.com/a).")], + ), + context, + ) + assert enriched.links[0].uri == "https://example.com/a" + + +def test_feishu_cli_error_boundaries(tmp_path): + context = _context(tmp_path) + assert not looks_like_feishu_url("not-a-url") + + with patch("openkb.ingest.importers.feishu.subprocess.run", side_effect=FileNotFoundError): + with pytest.raises(ValueError, match="requires lark-cli"): + _fetch_markdown_with_lark_cli("https://x.feishu.cn/wiki/a", context) + + with patch( + "openkb.ingest.importers.feishu.subprocess.run", + side_effect=subprocess.TimeoutExpired("lark-cli", 1), + ): + with pytest.raises(ValueError, match="timed out"): + _fetch_markdown_with_lark_cli("https://x.feishu.cn/wiki/a", context) + + with patch( + "openkb.ingest.importers.feishu.subprocess.run", + return_value=_completed(stderr="denied", returncode=2), + ): + with pytest.raises(ValueError, match="denied"): + _fetch_markdown_with_lark_cli("https://x.feishu.cn/wiki/a", context) + + with patch( + "openkb.ingest.importers.feishu.subprocess.run", + return_value=_completed(stdout='{"data": {"title": "Empty"}}'), + ): + with pytest.raises(ValueError, match="empty markdown"): + _fetch_markdown_with_lark_cli("https://x.feishu.cn/wiki/a", context) + + +def test_feishu_parse_helpers_and_config(tmp_path): + assert _parse_lark_cli_output("plain markdown") == ("plain markdown", None) + nested = {"data": [{"node": {"body": "# Body", "name": "Nested"}}]} + assert _parse_lark_cli_output(json.dumps(nested)) == ("# Body", "Nested") + assert _find_first_string([{"x": ""}, {"title": "T"}], ("title",)) == "T" + assert _int_config(True, default=5) == 5 + assert _int_config("bad", default=5) == 5 + assert _int_config("0", default=5) == 1 + assert _fallback_title("https://x.feishu.cn/") == "x.feishu.cn" + + context = _context(tmp_path, {"ingest": {"feishu": {"timeout": "2", "cli": "custom"}}}) + with patch( + "openkb.ingest.importers.feishu.subprocess.run", + return_value=_completed(stdout="# Body"), + ) as run: + assert _fetch_markdown_with_lark_cli("https://x.feishu.cn/wiki/a", context) == ( + "# Body", + None, + ) + assert run.call_args.kwargs["timeout"] == 2 + assert run.call_args.args[0][0] == "custom" + + +def test_feishu_importer_unique_fallback_name(tmp_path): + context = _context(tmp_path) + existing = context.staging_dir / "raw" / "abc.md" + existing.parent.mkdir(parents=True) + existing.write_text("old", encoding="utf-8") + with patch( + "openkb.ingest.importers.feishu._fetch_markdown_with_lark_cli", + return_value=("# Body", None), + ): + input_ = FeishuImporter().import_source("https://x.feishu.cn/wiki/abc", context) + assert input_.path is not None + assert input_.path.name == "abc_2.md" + + +def test_image_vision_parse_and_noop_boundaries(tmp_path): + asset = tmp_path / "image.png" + asset.write_bytes(b"img") + bundle = DocumentBundle( + id="b", + title=None, + source_uri="b", + blocks=[TextBlock("text"), ImageBlock(asset_id="missing")], + assets=[ + Asset( + id="source", + path=asset, + media_type="image/png", + sha256="hash", + ) + ], + ) + context = _context(tmp_path, {}) + enricher = ImageVisionEnricher() + assert not enricher.applies_to(bundle, context) + assert enricher.enrich(bundle, context) is bundle + assert _image_model(context) == "gpt-4o-mini" + + parsed = _parse_vision_response("not json") + assert parsed.visual_description == "not json" + assert _parse_vision_response("[1]").visual_description == "[1]" + fenced = _parse_vision_response( + '```json\n{"visual_description": "desc", "ocr_text": 7, "keywords": ["a", 1]}\n```' + ) + assert fenced.visual_description == "desc" + assert fenced.ocr_text is None + assert fenced.keywords == ["a"] + assert _response_content({"content": "top"}) == "top" + + +def test_plugins_loader_boundaries(): + assert load_ingest_entry_points("group", ()) == [] + + class Component: + name = "component" + + class Factory: + def __call__(self): + return Component() + + entry_points = _FakeEntryPoints( + [ + _FakeEntryPoint("group", "skip", Component), + _FakeEntryPoint("group", "excluded", Component), + _FakeEntryPoint("group", "factory", Factory()), + _FakeEntryPoint("group", "bad", object()), + _FakeEntryPoint("group", "boom", RuntimeError("boom"), raises=True), + ] + ) + with patch("openkb.ingest.plugins.entry_points", return_value=entry_points): + loaded = load_ingest_entry_points("group", ("factory", "excluded"), exclude={"excluded"}) + assert [component.name for component in loaded] == ["component"] + with pytest.raises(ValueError, match="did not produce"): + load_ingest_entry_points("group", ("bad",)) + with pytest.raises(ValueError, match="boom"): + load_ingest_entry_points("group", ("boom",)) + + class LegacyEntryPoints(dict): + pass + + with patch("openkb.ingest.plugins.entry_points", return_value=LegacyEntryPoints(group=[])): + assert load_ingest_entry_points("group", ("x",)) == [] + + +def test_render_and_serialization_boundaries(tmp_path): + source = tmp_path / "image.png" + source.write_bytes(b"img") + images_dir = tmp_path / "images" + bundle = DocumentBundle( + id="b", + title="Title", + source_uri="b", + blocks=[ + ImageBlock( + asset_id="source", + visual_description="desc", + ocr_text="text", + ) + ], + ) + rendered = _render_image_markdown(bundle, source, "doc", images_dir) + assert "## 视觉描述" in rendered + assert "## 可见文字" in rendered + + with pytest.raises(ValueError, match="source_path"): + _bundle_source_path(DocumentBundle(id="x", title=None, source_uri="x", blocks=[])) + + sidecar = bundle_to_dict( + DocumentBundle( + id="x", + title=None, + source_uri="x", + blocks=[ + ImageBlock(asset_id="a"), + TableBlock("| a |"), + EmbedBlock("https://example.com", title="Example"), + ], + assets=[Asset("a", source, "image/png", "sha", metadata={"p": source})], + links=[DiscoveredLink("u", metadata={"items": (source, object())})], + metadata={1: object()}, + provenance=[ProvenanceRecord("x")], + ) + ) + assert [block["type"] for block in sidecar["blocks"]] == ["image", "table", "embed"] + assert sidecar["assets"][0]["metadata"]["p"].endswith("image.png") + assert sidecar["links"][0]["metadata"]["items"][0].endswith("image.png") + + +@dataclass +class _FakeEntryPoint: + group: str + name: str + loaded: object + raises: bool = False + + def load(self): + if self.raises: + raise self.loaded + return self.loaded + + +class _FakeEntryPoints(list): + def select(self, *, group: str): + return [entry_point for entry_point in self if entry_point.group == group] diff --git a/tests/test_ingest_bundle.py b/tests/test_ingest_bundle.py new file mode 100644 index 00000000..d6b0c68e --- /dev/null +++ b/tests/test_ingest_bundle.py @@ -0,0 +1,939 @@ +"""Tests for the bundle ingest add path.""" + +from __future__ import annotations + +import json +from dataclasses import replace +from unittest.mock import patch + +from click.testing import CliRunner + +from openkb.cli import cli +from openkb.converter import _registry_path +from openkb.ingest.add import add_bundle_target, resolve_ingest_pipeline +from openkb.ingest.context import IngestContext +from openkb.ingest.enrichers import LinkDiscoveryEnricher +from openkb.ingest.importers.feishu import looks_like_feishu_url +from openkb.ingest.models import DocumentBundle, IngestInput, ProvenanceRecord, TextBlock +from openkb.state import HashRegistry + + +def _setup_kb(tmp_path): + (tmp_path / "raw").mkdir() + (tmp_path / "wiki" / "sources" / "images").mkdir(parents=True) + (tmp_path / "wiki" / "summaries").mkdir(parents=True) + (tmp_path / "wiki" / "concepts").mkdir(parents=True) + (tmp_path / "wiki" / "entities").mkdir(parents=True) + (tmp_path / "wiki" / "index.md").write_text("# Index\n", encoding="utf-8") + (tmp_path / "wiki" / "log.md").write_text("# Operations Log\n\n", encoding="utf-8") + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir() + (openkb_dir / "config.yaml").write_text("model: gpt-4o-mini\n", encoding="utf-8") + (openkb_dir / "hashes.json").write_text(json.dumps({}), encoding="utf-8") + return tmp_path + + +def _write_config(kb_dir, body: str) -> None: + (kb_dir / ".openkb" / "config.yaml").write_text(body, encoding="utf-8") + + +def _write_url_markdown(raw_dir, name: str, body: str): + raw_dir.mkdir(parents=True, exist_ok=True) + path = raw_dir / name + path.write_text(body, encoding="utf-8") + return path + + +def _completed_process(stdout: str, *, returncode: int = 0, stderr: str = ""): + return type( + "Completed", + (), + {"stdout": stdout, "stderr": stderr, "returncode": returncode}, + )() + + +class _FakeEntryPoint: + def __init__(self, group: str, name: str, loaded): + self.group = group + self.name = name + self._loaded = loaded + + def load(self): + return self._loaded + + +class _FakeEntryPoints(list): + def select(self, *, group: str): + return [entry_point for entry_point in self if entry_point.group == group] + + +def test_resolve_ingest_pipeline_defaults_to_legacy(): + assert resolve_ingest_pipeline({}) == "legacy" + assert resolve_ingest_pipeline({"ingest": {"pipeline": "bundle"}}) == "bundle" + assert resolve_ingest_pipeline({"ingest": {"pipeline": "bundle"}}, "legacy") == "legacy" + + +def test_external_importer_and_normalizer_entry_points_can_add_custom_source(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n importers:\n" + " enabled:\n - mocksource\n normalizers:\n enabled:\n" + " - mocknormalizer\n", + ) + + class MockImporter: + name = "mocksource" + + def can_handle(self, target, context): + del context + return target == "mock://doc" + + def import_source(self, target, context): + path = context.staging_dir / "raw" / "mock.mock" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("raw mock", encoding="utf-8") + return IngestInput( + target=target, + path=path, + source_uri=target, + media_type="application/x-mock", + metadata={"display_name": "mock.mock"}, + ) + + class MockNormalizer: + name = "mocknormalizer" + + def supports(self, input_, context): + del context + return input_.media_type == "application/x-mock" + + def normalize(self, input_, context): + del context + source_uri = input_.source_uri or input_.target + return DocumentBundle( + id=source_uri, + title="Mock", + source_uri=source_uri, + blocks=[TextBlock("# Mock\n\nPlugin body.")], + metadata={ + "display_name": "mock.mock", + "source_path": input_.path.as_posix(), + "source_identity": source_uri, + }, + provenance=[ProvenanceRecord(source_uri=source_uri)], + ) + + async def compile_noop(*args, **kwargs): + return None + + entry_points = _FakeEntryPoints( + [ + _FakeEntryPoint("openkb.ingest.importers", "mocksource", MockImporter), + _FakeEntryPoint("openkb.ingest.normalizers", "mocknormalizer", MockNormalizer), + ] + ) + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.plugins.entry_points", return_value=entry_points), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target("mock://doc", kb_dir) + + assert outcome == "added" + assert (kb_dir / "wiki" / "sources" / "mock.md").read_text(encoding="utf-8") == ( + "# Mock\n\nPlugin body.\n" + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["path"] == "mock://doc" + + +def test_external_enricher_entry_point_can_modify_bundle(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n enrichers:\n" + " enabled:\n - append_text\n", + ) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n", encoding="utf-8") + + class AppendTextEnricher: + name = "append_text" + + def applies_to(self, bundle, context): + del bundle, context + return True + + def enrich(self, bundle, context): + del context + return replace(bundle, blocks=[*bundle.blocks, TextBlock("Plugin enrichment.")]) + + async def compile_noop(*args, **kwargs): + return None + + entry_points = _FakeEntryPoints( + [_FakeEntryPoint("openkb.ingest.enrichers", "append_text", AppendTextEnricher)] + ) + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.plugins.entry_points", return_value=entry_points), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "added" + rendered = (kb_dir / "wiki" / "sources" / "notes.md").read_text(encoding="utf-8") + assert rendered == "# Notes\n\nPlugin enrichment.\n" + + +def test_cli_bundle_pipeline_dispatches_single_file(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n", encoding="utf-8") + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.ingest.add.add_bundle_target", return_value="added") as mock_add, + patch("openkb.cli.add_single_file") as legacy_add, + ): + result = runner.invoke(cli, ["add", str(doc), "--ingest-pipeline", "bundle"]) + + assert result.exit_code == 0, result.output + mock_add.assert_called_once() + assert mock_add.call_args.args == (doc, kb_dir) + assert mock_add.call_args.kwargs["legacy_fallback"] is legacy_add + legacy_add.assert_not_called() + + +def test_cli_config_bundle_pipeline_dispatches_directory(tmp_path): + kb_dir = _setup_kb(tmp_path) + (kb_dir / ".openkb" / "config.yaml").write_text( + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n", + encoding="utf-8", + ) + docs = tmp_path / "docs" + docs.mkdir() + (docs / "a.md").write_text("# A\n", encoding="utf-8") + (docs / "b.md").write_text("# B\n", encoding="utf-8") + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.ingest.add.add_bundle_target", return_value="added") as mock_add, + ): + result = runner.invoke(cli, ["add", str(docs)]) + + assert result.exit_code == 0, result.output + assert [call.args[0].name for call in mock_add.call_args_list] == ["a.md", "b.md"] + assert all( + call.kwargs["legacy_fallback"].__name__ == "add_single_file" + for call in mock_add.call_args_list + ) + + +def test_cli_bundle_pipeline_dispatches_url_to_bundle(tmp_path): + kb_dir = _setup_kb(tmp_path) + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.ingest.add.add_bundle_target", return_value="added") as mock_add, + patch("openkb.cli.add_single_file") as legacy_add, + patch("openkb.url_ingest.fetch_url_to_raw") as legacy_fetch, + ): + result = runner.invoke( + cli, + ["add", "https://example.com/article", "--ingest-pipeline", "bundle"], + ) + + assert result.exit_code == 0, result.output + mock_add.assert_called_once_with( + "https://example.com/article", + kb_dir, + legacy_fallback=legacy_add, + ) + legacy_fetch.assert_not_called() + legacy_add.assert_not_called() + + +def test_feishu_importer_recognizes_supported_doc_urls(): + assert looks_like_feishu_url("https://acme.feishu.cn/wiki/D3l4w0Y") + assert looks_like_feishu_url("https://acme.feishu.cn/docx/AbCdEf") + assert looks_like_feishu_url("https://acme.larksuite.com/docs/doccn123") + assert not looks_like_feishu_url("https://acme.feishu.cn/sheets/abc") + assert not looks_like_feishu_url("https://example.com/wiki/D3l4w0Y") + + +def test_cli_bundle_pipeline_accepts_image_file(tmp_path): + kb_dir = _setup_kb(tmp_path) + image = tmp_path / "diagram.png" + image.write_bytes(b"\x89PNG\r\n\x1a\nfake") + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.ingest.add.add_bundle_target", return_value="added") as mock_add, + ): + result = runner.invoke(cli, ["add", str(image), "--ingest-pipeline", "bundle"]) + + assert result.exit_code == 0, result.output + mock_add.assert_called_once() + assert mock_add.call_args.args == (image, kb_dir) + + +def test_cli_legacy_pipeline_still_rejects_image_file(tmp_path): + kb_dir = _setup_kb(tmp_path) + image = tmp_path / "diagram.png" + image.write_bytes(b"\x89PNG\r\n\x1a\nfake") + + runner = CliRunner() + with patch("openkb.cli._find_kb_dir", return_value=kb_dir): + result = runner.invoke(cli, ["add", str(image)]) + + assert result.exit_code == 0, result.output + assert "Unsupported file type: .png" in result.output + + +def test_bundle_markdown_add_writes_compatible_artifacts_and_registry(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n\nBody.\n", encoding="utf-8") + compile_calls = [] + + async def compile_noop(*args, **kwargs): + compile_calls.append((args, kwargs)) + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "added" + assert len(compile_calls) == 1 + assert (kb_dir / "raw" / "notes.md").read_text(encoding="utf-8") == "# Notes\n\nBody.\n" + assert (kb_dir / "wiki" / "sources" / "notes.md").read_text(encoding="utf-8") == ( + "# Notes\n\nBody.\n" + ) + bundle_sidecar = kb_dir / ".openkb" / "bundles" / "notes.json" + bundle_json = json.loads(bundle_sidecar.read_text(encoding="utf-8")) + assert bundle_json["schema_version"] == 1 + assert bundle_json["source_uri"] == _registry_path(doc, kb_dir) + assert bundle_json["blocks"][0]["type"] == "text" + assert bundle_json["provenance"][0]["source_uri"] == _registry_path(doc, kb_dir) + + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert len(entries) == 1 + ((_, meta),) = entries.items() + + assert meta == { + "name": "notes.md", + "doc_name": "notes", + "type": "md", + "path": _registry_path(doc, kb_dir), + "raw_path": "raw/notes.md", + "source_path": "wiki/sources/notes.md", + "bundle_path": ".openkb/bundles/notes.json", + } + + +def test_bundle_markdown_add_rolls_back_on_compile_failure(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n", encoding="utf-8") + + async def fail_compile(*args, **kwargs): + raise RuntimeError("LLM 503") + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=fail_compile), + patch("openkb.cli.time.sleep"), + ): + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "failed" + assert not (kb_dir / "raw" / "notes.md").exists() + assert not (kb_dir / "wiki" / "sources" / "notes.md").exists() + assert HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() == {} + + +def test_bundle_text_add_writes_markdown_source(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "notes.txt" + doc.write_text("Plain notes\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "added" + assert (kb_dir / "raw" / "notes.txt").read_text(encoding="utf-8") == "Plain notes\n" + assert (kb_dir / "wiki" / "sources" / "notes.md").read_text(encoding="utf-8") == ( + "Plain notes\n" + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["type"] == "txt" + assert meta["raw_path"] == "raw/notes.txt" + assert meta["source_path"] == "wiki/sources/notes.md" + + +def test_bundle_short_pdf_uses_existing_pdf_renderer(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "paper.pdf" + doc.write_bytes(b"%PDF-1.4 fake") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.normalizers.pdf.get_pdf_page_count", return_value=3), + patch( + "openkb.ingest.render.convert_pdf_with_images", + return_value="# Paper\n\nConverted.", + ) as convert_pdf, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "added" + convert_pdf.assert_called_once() + assert (kb_dir / "wiki" / "sources" / "paper.md").read_text(encoding="utf-8") == ( + "# Paper\n\nConverted." + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["type"] == "pdf" + + +def test_bundle_markitdown_path_uses_existing_base64_extractor(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "report.docx" + doc.write_bytes(b"fake docx") + mock_result = type("Result", (), {"text_content": "![](data:image/png;base64,abc123)"})() + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.render.MarkItDown") as markitdown, + patch("openkb.ingest.render.extract_base64_images", return_value="converted markdown"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + markitdown.return_value.convert.return_value = mock_result + outcome = add_bundle_target(doc, kb_dir) + + assert outcome == "added" + markitdown.return_value.convert.assert_called_once_with(str(doc.resolve()), keep_data_uris=True) + assert (kb_dir / "wiki" / "sources" / "report.md").read_text(encoding="utf-8") == ( + "converted markdown" + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["type"] == "docx" + + +def test_bundle_url_add_writes_staged_artifacts_and_registry(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n importers:\n" + " enabled:\n - url\n", + ) + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch( + "openkb.ingest.importers.url.fetch_url_to_dir", + side_effect=lambda _url, raw_dir: _write_url_markdown( + raw_dir, + "Remote-Article.md", + "# Remote Article\n\nBody.\n", + ), + ) as fetch, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target("https://example.com/article", kb_dir) + + assert outcome == "added" + fetch.assert_called_once() + assert (kb_dir / "raw" / "Remote-Article.md").read_text(encoding="utf-8") == ( + "# Remote Article\n\nBody.\n" + ) + assert (kb_dir / "wiki" / "sources" / "Remote-Article.md").read_text( + encoding="utf-8" + ) == "# Remote Article\n\nBody.\n" + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["path"] == "https://example.com/article" + assert meta["raw_path"] == "raw/Remote-Article.md" + assert meta["source_path"] == "wiki/sources/Remote-Article.md" + + +def test_bundle_feishu_add_uses_lark_cli_and_preserves_source_url(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n importers:\n" + " enabled:\n - feishu\n", + ) + url = "https://acme.feishu.cn/wiki/D3l4w0Y" + stdout = json.dumps({"data": {"title": "Feishu PRD", "content": "# Feishu PRD\n\nBody."}}) + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch( + "openkb.ingest.importers.feishu.subprocess.run", + return_value=_completed_process(stdout), + ) as run, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(url, kb_dir) + + assert outcome == "added" + command = run.call_args.args[0] + assert command[:3] == ["lark-cli", "docs", "+fetch"] + assert command[command.index("--doc") + 1] == url + assert (kb_dir / "raw" / "Feishu-PRD.md").read_text(encoding="utf-8") == ( + "# Feishu PRD\n\nBody.\n" + ) + assert (kb_dir / "wiki" / "sources" / "Feishu-PRD.md").read_text( + encoding="utf-8" + ) == "# Feishu PRD\n\nBody.\n" + bundle_json = json.loads( + (kb_dir / ".openkb" / "bundles" / "Feishu-PRD.json").read_text(encoding="utf-8") + ) + assert bundle_json["title"] == "Feishu PRD" + assert bundle_json["source_uri"] == url + assert bundle_json["metadata"]["source_metadata"]["source_system"] == "feishu" + assert ( + bundle_json["metadata"]["source_metadata"]["permission_boundary"] + == "current lark-cli identity" + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["path"] == url + assert meta["raw_path"] == "raw/Feishu-PRD.md" + assert meta["source_path"] == "wiki/sources/Feishu-PRD.md" + + +def test_bundle_image_add_writes_image_asset_markdown_and_registry(tmp_path): + kb_dir = _setup_kb(tmp_path) + image = tmp_path / "diagram.png" + image_bytes = b"\x89PNG\r\n\x1a\nfake image" + image.write_bytes(image_bytes) + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(image, kb_dir) + + assert outcome == "added" + assert (kb_dir / "raw" / "diagram.png").read_bytes() == image_bytes + rendered = (kb_dir / "wiki" / "sources" / "diagram.md").read_text(encoding="utf-8") + assert rendered == "# diagram\n\n![source image](sources/images/diagram/diagram.png)\n" + assert (kb_dir / "wiki" / "sources" / "images" / "diagram" / "diagram.png").read_bytes() == ( + image_bytes + ) + ((_, meta),) = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries().items() + assert meta["type"] == "png" + assert meta["raw_path"] == "raw/diagram.png" + assert meta["source_path"] == "wiki/sources/diagram.md" + + +def test_bundle_image_vision_enricher_writes_derived_sections(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n" + " image_model: gpt-4o-vision\n enrichers:\n enabled:\n" + " - image_vision\n", + ) + image = tmp_path / "diagram.png" + image.write_bytes(b"\x89PNG\r\n\x1a\nfake image") + vision_payload = json.dumps( + { + "visual_description": "A system architecture diagram with three boxes.", + "ocr_text": "API Gateway", + "keywords": ["architecture", "gateway"], + "uncertainty": "Small labels may be incomplete.", + } + ) + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch( + "litellm.completion", + return_value={"choices": [{"message": {"content": vision_payload}}]}, + ) as completion, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(image, kb_dir) + + assert outcome == "added" + assert completion.call_args.kwargs["model"] == "gpt-4o-vision" + message_content = completion.call_args.kwargs["messages"][0]["content"] + assert message_content[1]["image_url"]["url"].startswith("data:image/png;base64,") + rendered = (kb_dir / "wiki" / "sources" / "diagram.md").read_text(encoding="utf-8") + assert "## 视觉描述" in rendered + assert "模型派生信息，可能不完整或不准确。" in rendered + assert "A system architecture diagram with three boxes." in rendered + assert "## 可见文字" in rendered + assert "API Gateway" in rendered + + +def test_cli_bundle_long_pdf_falls_back_to_legacy(tmp_path): + kb_dir = _setup_kb(tmp_path) + doc = tmp_path / "long.pdf" + doc.write_bytes(b"%PDF-1.4 fake") + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.ingest.normalizers.pdf.get_pdf_page_count", return_value=200), + patch("openkb.cli.add_single_file", return_value="added") as legacy_add, + ): + result = runner.invoke(cli, ["add", str(doc), "--ingest-pipeline", "bundle"]) + + assert result.exit_code == 0, result.output + assert "Falling back to legacy long-document ingest" in result.output + legacy_add.assert_called_once_with(doc.resolve(), kb_dir) + + +def test_link_discovery_enricher_finds_markdown_and_bare_links(tmp_path): + kb_dir = _setup_kb(tmp_path) + context = IngestContext(kb_dir=kb_dir, config={}, staging_dir=kb_dir / ".openkb" / "staging") + bundle = DocumentBundle( + id="root", + title="Root", + source_uri="root.md", + blocks=[ + TextBlock( + text=("See [Child](child.md), ![figure](image.png), and https://example.com/page.") + ) + ], + ) + + enriched = LinkDiscoveryEnricher().enrich(bundle, context) + + assert [(link.uri, link.text) for link in enriched.links] == [ + ("child.md", "Child"), + ("https://example.com/page", None), + ] + + +def test_bundle_recursive_imports_local_markdown_links(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n", + ) + root = tmp_path / "root.md" + child = tmp_path / "child.md" + root.write_text("# Root\n\nSee [Child](child.md).\n", encoding="utf-8") + child.write_text("# Child\n\nBody.\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + assert (kb_dir / "wiki" / "sources" / "root.md").exists() + assert (kb_dir / "wiki" / "sources" / "child.md").exists() + child_bundle = json.loads( + (kb_dir / ".openkb" / "bundles" / "child.json").read_text(encoding="utf-8") + ) + assert child_bundle["provenance"][0]["relationship"] == "child" + assert child_bundle["provenance"][0]["parent_uri"] == _registry_path(root, kb_dir) + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert {meta["doc_name"] for meta in entries.values()} == {"root", "child"} + + +def test_bundle_recursion_respects_max_depth_zero(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 0\n max_documents: 5\n", + ) + root = tmp_path / "root.md" + child = tmp_path / "child.md" + root.write_text("# Root\n\nSee [Child](child.md).\n", encoding="utf-8") + child.write_text("# Child\n\nBody.\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + assert (kb_dir / "wiki" / "sources" / "root.md").exists() + assert not (kb_dir / "wiki" / "sources" / "child.md").exists() + + +def test_bundle_recursion_respects_max_documents(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 1\n", + ) + root = tmp_path / "root.md" + child = tmp_path / "child.md" + root.write_text("# Root\n\nSee [Child](child.md).\n", encoding="utf-8") + child.write_text("# Child\n\nBody.\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + assert (kb_dir / "wiki" / "sources" / "root.md").exists() + assert not (kb_dir / "wiki" / "sources" / "child.md").exists() + + +def test_bundle_recursion_deduplicates_cycles(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 5\n max_documents: 5\n", + ) + root = tmp_path / "root.md" + child = tmp_path / "child.md" + root.write_text("# Root\n\nSee [Child](child.md).\n", encoding="utf-8") + child.write_text("# Child\n\nBack to [Root](root.md).\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert {meta["doc_name"] for meta in entries.values()} == {"root", "child"} + + +def test_bundle_recursion_skips_url_when_url_importer_disabled(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n allow_domains:\n - example.com\n", + ) + root = tmp_path / "root.md" + root.write_text("# Root\n\nSee [Remote](https://example.com/child.md).\n", encoding="utf-8") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert [meta["doc_name"] for meta in entries.values()] == ["root"] + + +def test_bundle_recursion_imports_allowed_url_links(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n allow_domains:\n - example.com\n" + " importers:\n enabled:\n - file\n - url\n", + ) + root = tmp_path / "root.md" + root.write_text("# Root\n\nSee [Remote](https://example.com/child?q=1#section).\n") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch( + "openkb.ingest.importers.url.fetch_url_to_dir", + side_effect=lambda url, raw_dir: _write_url_markdown( + raw_dir, + "child.md", + f"# Child\n\nFetched from {url}.\n", + ), + ) as fetch, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + fetch.assert_called_once() + assert fetch.call_args.args[0] == "https://example.com/child?q=1" + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert {meta["doc_name"] for meta in entries.values()} == {"root", "child"} + child_meta = next(meta for meta in entries.values() if meta["doc_name"] == "child") + assert child_meta["path"] == "https://example.com/child?q=1" + + +def test_bundle_recursion_imports_allowed_feishu_links_without_url_importer(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n allow_domains:\n - feishu.cn\n" + " importers:\n enabled:\n - file\n - feishu\n", + ) + root = tmp_path / "root.md" + feishu_url = "https://acme.feishu.cn/docx/AbCdEf#heading" + root.write_text(f"# Root\n\nSee [Feishu]({feishu_url}).\n") + stdout = json.dumps({"data": {"title": "Child Doc", "content": "# Child Doc\n\nBody."}}) + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch( + "openkb.ingest.importers.feishu.subprocess.run", + return_value=_completed_process(stdout), + ) as run, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + command = run.call_args.args[0] + assert command[command.index("--doc") + 1] == "https://acme.feishu.cn/docx/AbCdEf" + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert {meta["doc_name"] for meta in entries.values()} == {"root", "Child-Doc"} + child_meta = next(meta for meta in entries.values() if meta["doc_name"] == "Child-Doc") + assert child_meta["path"] == "https://acme.feishu.cn/docx/AbCdEf" + + +def test_bundle_recursion_blocks_feishu_link_outside_allow_domains(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n allow_domains:\n - example.com\n" + " importers:\n enabled:\n - file\n - feishu\n", + ) + root = tmp_path / "root.md" + root.write_text("# Root\n\nSee [Feishu](https://acme.feishu.cn/wiki/D3l4w0Y).\n") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.importers.feishu.subprocess.run") as run, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + run.assert_not_called() + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert [meta["doc_name"] for meta in entries.values()] == ["root"] + + +def test_bundle_recursion_blocks_url_outside_allow_domains(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n link_discovery: true\n" + " max_depth: 1\n max_documents: 5\n allow_domains:\n - example.com\n" + " importers:\n enabled:\n - file\n - url\n", + ) + root = tmp_path / "root.md" + root.write_text("# Root\n\nSee [Remote](https://blocked.example.net/child).\n") + + async def compile_noop(*args, **kwargs): + return None + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.importers.url.fetch_url_to_dir") as fetch, + patch("openkb.agent.compiler.compile_short_doc", new=compile_noop), + ): + outcome = add_bundle_target(root, kb_dir) + + assert outcome == "added" + fetch.assert_not_called() + entries = HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() + assert [meta["doc_name"] for meta in entries.values()] == ["root"] + + +def test_bundle_url_long_pdf_fallback_keeps_staged_file_during_fallback(tmp_path): + kb_dir = _setup_kb(tmp_path) + _write_config( + kb_dir, + "model: gpt-4o-mini\ningest:\n pipeline: bundle\n importers:\n" + " enabled:\n - url\n", + ) + fallback_seen = [] + + def fake_fetch(_url, raw_dir): + raw_dir.mkdir(parents=True, exist_ok=True) + path = raw_dir / "long.pdf" + path.write_bytes(b"%PDF-1.4 fake") + return path + + def legacy_fallback(path, kb_dir_arg): + fallback_seen.append((path.exists(), path, kb_dir_arg)) + return "added" + + with ( + patch("openkb.ingest.add._setup_ingest_llm"), + patch("openkb.ingest.importers.url.fetch_url_to_dir", side_effect=fake_fetch), + patch("openkb.ingest.normalizers.pdf.get_pdf_page_count", return_value=200), + ): + outcome = add_bundle_target( + "https://example.com/long.pdf", + kb_dir, + legacy_fallback=legacy_fallback, + ) + + assert outcome == "added" + assert fallback_seen + exists_during_fallback, fallback_path, fallback_kb_dir = fallback_seen[0] + assert exists_during_fallback is True + assert fallback_path.name == "long.pdf" + assert fallback_kb_dir == kb_dir diff --git a/tests/test_remove.py b/tests/test_remove.py index c5318ad5..eb7c2166 100644 --- a/tests/test_remove.py +++ b/tests/test_remove.py @@ -971,6 +971,26 @@ def test_cli_remove_dry_run_does_not_touch_images(kb_dir): assert (images_dir / "p1_img1.png").exists() +def test_cli_remove_deletes_bundle_sidecar(kb_dir): + _seed_two_doc_kb(kb_dir) + registry_path = kb_dir / ".openkb" / "hashes.json" + entries = json.loads(registry_path.read_text(encoding="utf-8")) + entries["h_a"]["bundle_path"] = ".openkb/bundles/attention-h_a.json" + registry_path.write_text(json.dumps(entries), encoding="utf-8") + bundle_path = kb_dir / ".openkb" / "bundles" / "attention-h_a.json" + bundle_path.parent.mkdir(parents=True) + bundle_path.write_text('{"schema_version": 1}', encoding="utf-8") + + dry_run = _invoke(kb_dir, ["remove", "attention.pdf", "--dry-run"]) + assert dry_run.exit_code == 0, dry_run.output + assert ".openkb/bundles/attention-h_a.json" in dry_run.output + assert bundle_path.exists() + + result = _invoke(kb_dir, ["remove", "attention.pdf", "--yes"]) + assert result.exit_code == 0, result.output + assert not bundle_path.exists() + + # --------------------------------------------------------------------------- # Functional-completeness fix: `doc_id` is persisted for long PDFs so a # later `openkb remove` can call PageIndex's delete_document API. diff --git a/tests/test_url_ingest.py b/tests/test_url_ingest.py index 6e39cfcc..45aee84e 100644 --- a/tests/test_url_ingest.py +++ b/tests/test_url_ingest.py @@ -11,6 +11,7 @@ _sanitize_filename, _sniff_content_type, _unique_path, + fetch_url_to_dir, fetch_url_to_raw, looks_like_url, ) @@ -209,6 +210,22 @@ def test_fetch_pdf_writes_chunked_to_raw_dir(tmp_path): assert result.read_bytes() == body +def test_fetch_url_to_dir_writes_to_requested_directory(tmp_path): + body = b"%PDF-1.4\n" + b"x" * 1000 + resp = _fake_response( + body=body, + headers={"Content-Type": "application/pdf"}, + ) + staged_raw = tmp_path / ".openkb" / "staging" / "bundle-1" / "raw" + + with patch("urllib.request.urlopen", return_value=resp): + result = fetch_url_to_dir("https://example.com/paper.pdf", staged_raw) + + assert result == staged_raw / "paper.pdf" + assert result.read_bytes() == body + assert not (tmp_path / "raw" / "paper.pdf").exists() + + def test_fetch_pdf_with_lying_octet_stream_header(tmp_path): """Server says octet-stream but the body starts with %PDF — magic bytes must win and the file gets the .pdf extension."""