"""PyPI — `https://pypi.org/pypi//json` (public JSON document) for the AI Python ecosystem. Tier 2. Packages = the `pypi` names of registry/repositories.yaml (identifiers `pypi` + `github_repo` → merges with the GitHub entity) plus a curated list of standalone libraries (identifiers `pypi` only, entity type `library`). Claims: `pypi_version`, `pypi_release_at` (earliest upload time of the latest release's files), `license`, `description` (summary), `homepage`, `requires_python`, `metric.releases`, `repository_url` (when the project links GitHub). For packages that are NOT mapped to a GitHub repository the connector also sets `latest_version` / `latest_release_at`; for mapped packages those two properties are owned by the GitHub release feed so two same-tier sources never flip-flop the same property when a tag lands hours before the wheel. """ from __future__ import annotations import re from datetime import UTC from typing import Any from aiatlas.registry import org_by_github, org_ref from aiatlas.sdk.connector import BaseConnector, Parsed, RunContext from aiatlas.sdk.extract.dates import parse_datetime from aiatlas.sdk.facts import EntityRef, Facts, Target from aiatlas.sdk.fetch import FetchResult from .github import claim_license, repo_entity, repo_entries EXTRA_PACKAGES = [ "torch", "transformers", "vllm", "mlx", "mlx-lm", "llama-cpp-python", "langchain", "langgraph", "llama-index", "sglang", "jax", "onnxruntime", "diffusers", "peft", "trl", "unsloth", "openai", "anthropic", "google-genai", "mistralai", "cohere", "litellm", "faiss-cpu", "chromadb", "qdrant-client", "pymilvus", "weaviate-client", "lancedb", "crewai", "pydantic-ai", "browser-use", "ragas", "langfuse", "ray", "bentoml", "triton", "flash-attn", "lm-eval", "sentence-transformers", "accelerate", "bitsandbytes", "autoawq", "gguf", "safetensors", "tokenizers", "datasets", "evaluate", "gradio", "streamlit", "fastapi", "huggingface-hub", "optimum", "xformers", "einops", "tiktoken", "instructor", "dspy", "outlines", "guidance", "semantic-kernel", "haystack-ai", "autogen-agentchat", "smolagents", "mcp", "openai-agents", "letta", "vector-quantize-pytorch", "timm", "torchvision", "torchaudio", "scikit-learn", "xgboost", "lightgbm", "keras", "tensorflow", "deepspeed", "megatron-core", "flax", "optax", "numpyro", "pyro-ppl", "spacy", "nltk", "whisperx", "faster-whisper", "ctranslate2", "onnx", "tensorrt", "nvidia-modelopt", "auto-gptq", "exllamav2", "llmcompressor", "sparseml", "lmdeploy", "text-generation", "openllm", "modal", "wandb", "mlflow", "comet-ml", "arize-phoenix", "opentelemetry-instrumentation-openai", "promptfoo", "deepeval", "inspect-ai", "trulens", "guardrails-ai", ] GITHUB_URL = re.compile(r"https?://github\.com/([\w.-]+/[\w.-]+)") class PyPIConnector(BaseConnector): name = "pypi" label = "PyPI — versions and metadata of the AI Python ecosystem" description = "Public JSON documents of the packages in registry/repositories.yaml plus a curated list of standalone AI libraries." source_key = "pypi.org" version = "2" parser_version = "2" interval_seconds = 12 * 3600 min_interval_seconds = 6 * 3600 max_interval_seconds = 3 * 86400 rate_per_min = 30 tier = 2 priority = 2 expected_min_records = 60 concurrency = 3 def packages(self) -> dict[str, dict[str, Any] | None]: out: dict[str, dict[str, Any] | None] = {} for entry in repo_entries(): if entry.get("pypi"): out[_norm(entry["pypi"])] = entry for pkg in self.config.get("packages") or EXTRA_PACKAGES: out.setdefault(_norm(pkg), None) return out async def discover(self, ctx: RunContext) -> list[Target]: return [Target(url=f"https://pypi.org/pypi/{pkg}/json", doc_type="package", key=f"pkg:{pkg}", min_bytes=200, meta={"pypi": pkg, "repo": (entry or {}).get("repo"), "content_type": "application/json"}, priority=2) for pkg, entry in self.packages().items()] async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts: facts = Facts() data = parsed.json if parsed.kind == "json" else None if not isinstance(data, dict) or not isinstance(data.get("info"), dict): return facts info = data["info"] pkg = _norm(info.get("name") or target.meta.get("pypi") or "") if not pkg: return facts entry = next((e for e in repo_entries() if e.get("pypi") and _norm(e["pypi"]) == pkg), None) urls = {k.lower(): v for k, v in (info.get("project_urls") or {}).items() if isinstance(v, str)} gh = _github(urls, info.get("home_page")) if entry: ref = repo_entity(facts, entry) else: org = None if gh: known = org_by_github(gh.split("/")[0]) org = org_ref(known["key"]) if known else None ids = {"pypi": pkg} if gh and _norm(gh.split("/")[1]) == pkg: # `gguf` links to ggml-org/llama.cpp: a dependency, not the same project ids["github_repo"] = gh ref = facts.entity("library", info.get("name") or pkg, identifiers=ids, organization=org, aliases=[pkg]) if org: facts.entities.append(org) facts.relate(org, "develops", ref) if gh: facts.claim(ref, "repository_url", f"https://github.com/{gh}") facts.claim(ref, "pypi", pkg) facts.claim(ref, "pypi_url", info.get("package_url") or f"https://pypi.org/project/{pkg}/") version = info.get("version") released = _release_time(data, version) facts.claim(ref, "pypi_version", version) facts.claim(ref, "pypi_release_at", released) if not entry: facts.claim(ref, "latest_version", version) facts.claim(ref, "latest_release_at", released) lic = _license(info) claim_license(facts, ref, lic, prop="pypi_license") if not entry: claim_license(facts, ref, lic) # mapped repos: GitHub's SPDX id owns `license` (avoids same-tier flip-flops) facts.claim(ref, "kind", "library") summary = (info.get("summary") or "").strip() facts.claim(ref, "description", summary[:1000] or None) homepage = urls.get("homepage") or urls.get("home") or info.get("home_page") if not entry and isinstance(homepage, str) and homepage.startswith("http"): facts.claim(ref, "homepage", homepage) docs = urls.get("documentation") or urls.get("docs") facts.claim(ref, "docs_url", docs if isinstance(docs, str) and docs.startswith("http") else None) facts.claim(ref, "requires_python", info.get("requires_python") or None) facts.claim(ref, "author", (info.get("author") or "").strip() or None) facts.claim(ref, "metric.releases", len(data.get("releases") or {})) facts.claim(ref, "python_classifiers", [c for c in info.get("classifiers") or [] if c.startswith("Programming Language :: Python :: 3.")][:12]) facts.document_entity = ref facts.document_title = f"{info.get('name') or pkg} {version or ''} on PyPI".strip() return facts def _norm(name: str) -> str: return re.sub(r"[-_.]+", "-", name.strip()).lower() def _github(urls: dict[str, str], home: Any) -> str | None: for key in ("source", "source code", "repository", "code", "github", "homepage", "home"): v = urls.get(key) if v: m = GITHUB_URL.search(v) if m: return m.group(1).removesuffix(".git") if isinstance(home, str): m = GITHUB_URL.search(home) if m: return m.group(1).removesuffix(".git") return None def _license(info: dict[str, Any]) -> str | None: expr = info.get("license_expression") if isinstance(expr, str) and expr.strip(): return expr.strip() lic = info.get("license") if isinstance(lic, str) and 0 < len(lic.strip()) <= 80: return lic.strip() for c in info.get("classifiers") or []: if c.startswith("License :: OSI Approved :: "): return c.rsplit(" :: ", 1)[1] if c.startswith("License :: ") and c != "License :: OSI Approved": return c.rsplit(" :: ", 1)[1] return None def _release_time(data: dict[str, Any], version: str | None) -> str | None: files = (data.get("releases") or {}).get(version or "") or data.get("urls") or [] times = [parse_datetime(f.get("upload_time_iso_8601") or f.get("upload_time")) for f in files if isinstance(f, dict)] times = [t for t in times if t] return min(times).astimezone(UTC).isoformat(timespec="seconds") if times else None __all__ = ["EXTRA_PACKAGES", "PyPIConnector"] _ = EntityRef # re-exported type for tests CONNECTORS = [PyPIConnector]