Spaces:
Running
Running
| """ | |
| 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:])) | |