""" Install only the pyproject.toml optional-dependency groups that the active .env requires. Usage (no TUI required): python scripts/install_extras.py # auto-detect from .env python scripts/install_extras.py docling-ocr # force a specific group python scripts/install_extras.py --all # install every optional group python scripts/install_extras.py --dry-run # print pip command, don't run python scripts/install_extras.py --check # exit non-zero if missing Reads ./.env (or $BRAINAPI_DOTENV) for OCR_MODE and any future toggles, then calls `python -m pip install` against the active interpreter with the matching specs from pyproject.toml's [project.optional-dependencies]. Safe to re-run: pip is a no-op when the specs are already satisfied. """ from __future__ import annotations import argparse import importlib.util import os import re import subprocess import sys from pathlib import Path try: import tomllib except ModuleNotFoundError: import tomli as tomllib # type: ignore[no-redef] REPO_ROOT = Path(__file__).resolve().parent.parent PYPROJECT_PATH = REPO_ROOT / "pyproject.toml" DEFAULT_ENV_PATH = Path(os.environ.get("BRAINAPI_DOTENV", REPO_ROOT / ".env")) def parse_dotenv(path: Path) -> dict[str, str]: if not path.exists(): return {} result: dict[str, str] = {} for raw in path.read_text(encoding="utf-8").splitlines(): line = raw.strip() if not line or line.startswith("#") or "=" not in line: continue key, _, value = line.partition("=") result[key.strip()] = value.strip().strip('"').strip("'") return result def _load_pyproject() -> dict: if not PYPROJECT_PATH.exists(): return {} with PYPROJECT_PATH.open("rb") as fh: return tomllib.load(fh) def _poetry_dep_to_spec(name: str, value) -> str | None: if name == "python": return None if isinstance(value, str): return name if value == "*" else f"{name}{value}" if isinstance(value, dict): version = value.get("version", "") extras = value.get("extras", []) or [] head = name + (f"[{','.join(extras)}]" if extras else "") if not version or version == "*": return head return f"{head}{version}" return None def load_base_deps() -> list[str]: data = _load_pyproject() project_deps = data.get("project", {}).get("dependencies", []) or [] if project_deps: return list(project_deps) poetry_deps = ( data.get("tool", {}).get("poetry", {}).get("dependencies", {}) or {} ) out: list[str] = [] for name, value in poetry_deps.items(): spec = _poetry_dep_to_spec(name, value) if spec: out.append(spec) return out def load_optional_groups() -> dict[str, list[str]]: data = _load_pyproject() project_optional = ( data.get("project", {}).get("optional-dependencies", {}) or {} ) if project_optional: return {k: list(v) for k, v in project_optional.items()} poetry_groups = data.get("tool", {}).get("poetry", {}).get("group", {}) or {} out: dict[str, list[str]] = {} for group_name, group_value in poetry_groups.items(): if not isinstance(group_value, dict): continue deps = group_value.get("dependencies", {}) or {} specs: list[str] = [] for name, value in deps.items(): spec = _poetry_dep_to_spec(name, value) if spec: specs.append(spec) if specs: out[group_name] = specs return out POSTGRESQL_BACKEND_SPECS = [ "psycopg2-binary>=2.9.6,<3.0.0", "pgvector>=0.3.0,<1.0.0", "networkx>=3.0,<4.0.0", ] def uses_postgresql_backend(env: dict[str, str]) -> bool: vector = env.get("VECTOR_DB", "milvus").lower() data = env.get("DATA_DB", "mongo").lower() graph = env.get("GRAPH_DB", "neo4j").lower() return vector == "postgresql" or data == "postgresql" or graph == "networkx" def required_groups(env: dict[str, str]) -> list[str]: groups: list[str] = [] if env.get("OCR_MODE", "docparser").lower() == "docling": groups.append("docling-ocr") if uses_postgresql_backend(env): groups.append("postgresql-backend") return groups def required_backend_specs(env: dict[str, str]) -> list[str]: if not uses_postgresql_backend(env): return [] optional = load_optional_groups() if "postgresql-backend" in optional: return [] return list(POSTGRESQL_BACKEND_SPECS) def package_name_from_spec(spec: str) -> str: spec = spec.strip() match = re.match(r"^([A-Za-z0-9][A-Za-z0-9._-]*)", spec) if match: return match.group(1) return spec def package_installed(name: str) -> bool: pkg = package_name_from_spec(name) try: from importlib.metadata import PackageNotFoundError, version version(pkg) return True except PackageNotFoundError: pass except Exception: pass import_names = {pkg.replace("-", "_"), pkg} if pkg == "psycopg2-binary": import_names.add("psycopg2") return any(importlib.util.find_spec(c) is not None for c in import_names) def filter_missing(specs: list[str]) -> list[str]: return [s for s in specs if not package_installed(package_name_from_spec(s))] def install_specs(specs: list[str], dry_run: bool) -> int: if not specs: return 0 cmd = [sys.executable, "-m", "pip", "install", *specs] print(f"$ {' '.join(cmd)}") if dry_run: return 0 return subprocess.call(cmd) def main(argv: list[str]) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "groups", nargs="*", help="Optional groups to install. Omit to auto-detect from .env.", ) parser.add_argument( "--env-file", default=str(DEFAULT_ENV_PATH), help="Path to .env (default: ./.env, or $BRAINAPI_DOTENV).", ) parser.add_argument( "--dry-run", action="store_true", help="Print the pip command without running it.", ) parser.add_argument( "--check", action="store_true", help="Exit non-zero if any required extras are missing; do not install.", ) parser.add_argument( "--all", action="store_true", help="Install every optional group declared in pyproject.toml.", ) parser.add_argument( "--include-base", action="store_true", help="Also install [project].dependencies / [tool.poetry.dependencies] base deps.", ) parser.add_argument( "--upgrade-pip", action="store_true", help="Run `pip install --upgrade pip` before installing anything.", ) args = parser.parse_args(argv) optional_groups = load_optional_groups() if args.all: wanted = list(optional_groups.keys()) elif args.groups: wanted = args.groups else: env = parse_dotenv(Path(args.env_file)) wanted = required_groups(env) if wanted: print(f"Auto-detected required groups from {args.env_file}: {wanted}") elif not args.include_base: print( f"No optional groups required by {args.env_file} " f"(OCR_MODE={env.get('OCR_MODE', 'docparser')}).", ) unknown = [g for g in wanted if g not in optional_groups] known = [g for g in wanted if g in optional_groups] if unknown: print( f"Skipping optional groups not declared in pyproject.toml: {unknown}", file=sys.stderr, ) if optional_groups: print( f"Available optional groups: {list(optional_groups.keys())}", file=sys.stderr, ) else: print( "No optional-dependency groups are declared in pyproject.toml. " "These packages may already be part of base dependencies — " "continuing without them.", file=sys.stderr, ) wanted = known extra_specs: list[str] = [] for g in wanted: extra_specs.extend(optional_groups[g]) env = parse_dotenv(Path(args.env_file)) backend_specs = ( required_backend_specs(env) if not args.groups and not args.all else [] ) if args.check: missing = filter_missing([*extra_specs, *backend_specs]) if missing: print(f"Missing extras: {missing}", file=sys.stderr) hint_groups = wanted or ( ["postgresql-backend"] if backend_specs else [] ) if hint_groups: print( f"Install them with: {sys.executable} " f"scripts/install_extras.py {' '.join(hint_groups)}", file=sys.stderr, ) return 1 print("All required extras are installed.") return 0 base_specs = load_base_deps() if args.include_base else [] all_specs = [*base_specs, *extra_specs, *backend_specs] if not all_specs: return 0 if args.upgrade_pip and not args.dry_run: subprocess.call([sys.executable, "-m", "pip", "install", "--upgrade", "pip"]) return install_specs(all_specs, args.dry_run) if __name__ == "__main__": sys.exit(main(sys.argv[1:]))