1"""PyPI — `https://pypi.org/pypi/<pkg>/json` (public JSON document) for the AI Python ecosystem. Tier 2.23Packages = the `pypi` names of registry/repositories.yaml (identifiers `pypi` + `github_repo` → merges with the GitHub entity) plus a4curated list of standalone libraries (identifiers `pypi` only, entity type `library`).56Claims: `pypi_version`, `pypi_release_at` (earliest upload time of the latest release's files), `license`, `description` (summary),7`homepage`, `requires_python`, `metric.releases`, `repository_url` (when the project links GitHub). For packages that are NOT mapped to a8GitHub repository the connector also sets `latest_version` / `latest_release_at`; for mapped packages those two properties are owned by9the GitHub release feed so two same-tier sources never flip-flop the same property when a tag lands hours before the wheel.10"""11from __future__ import annotations1213import re14from datetime import UTC15from typing import Any1617from aiatlas.registry import org_by_github, org_ref18from aiatlas.sdk.connector import BaseConnector, Parsed, RunContext19from aiatlas.sdk.extract.dates import parse_datetime20from aiatlas.sdk.facts import EntityRef, Facts, Target21from aiatlas.sdk.fetch import FetchResult2223from .github import repo_entity, repo_entries2425EXTRA_PACKAGES = [26 "torch", "transformers", "vllm", "mlx", "mlx-lm", "llama-cpp-python", "langchain", "langgraph", "llama-index", "sglang", "jax", "onnxruntime",27 "diffusers", "peft", "trl", "unsloth", "openai", "anthropic", "google-genai", "mistralai", "cohere", "litellm", "faiss-cpu", "chromadb",28 "qdrant-client", "pymilvus", "weaviate-client", "lancedb", "crewai", "pydantic-ai", "browser-use", "ragas", "langfuse", "ray", "bentoml",29 "triton", "flash-attn", "lm-eval", "sentence-transformers", "accelerate", "bitsandbytes", "autoawq", "gguf", "safetensors", "tokenizers",30 "datasets", "evaluate", "gradio", "streamlit", "fastapi", "huggingface-hub", "optimum", "xformers", "einops", "tiktoken", "instructor",31 "dspy", "outlines", "guidance", "semantic-kernel", "haystack-ai", "autogen-agentchat", "smolagents", "mcp", "openai-agents", "letta",32 "vector-quantize-pytorch", "timm", "torchvision", "torchaudio", "scikit-learn", "xgboost", "lightgbm", "keras", "tensorflow", "deepspeed",33 "megatron-core", "flax", "optax", "numpyro", "pyro-ppl", "spacy", "nltk", "whisperx", "faster-whisper", "ctranslate2", "onnx", "tensorrt",34 "nvidia-modelopt", "auto-gptq", "exllamav2", "llmcompressor", "sparseml", "lmdeploy", "text-generation", "openllm", "modal", "wandb", "mlflow",35 "comet-ml", "arize-phoenix", "opentelemetry-instrumentation-openai", "promptfoo", "deepeval", "inspect-ai", "trulens", "guardrails-ai",36]37GITHUB_URL = re.compile(r"https?://github\.com/([\w.-]+/[\w.-]+)")383940class PyPIConnector(BaseConnector):41 name = "pypi"42 label = "PyPI — versions and metadata of the AI Python ecosystem"43 description = "Public JSON documents of the packages in registry/repositories.yaml plus a curated list of standalone AI libraries."44 source_key = "pypi.org"45 version = "1"46 parser_version = "1"47 interval_seconds = 12 * 360048 min_interval_seconds = 6 * 360049 max_interval_seconds = 3 * 8640050 rate_per_min = 3051 tier = 252 priority = 253 expected_min_records = 6054 concurrency = 35556 def packages(self) -> dict[str, dict[str, Any] | None]:57 out: dict[str, dict[str, Any] | None] = {}58 for entry in repo_entries():59 if entry.get("pypi"):60 out[_norm(entry["pypi"])] = entry61 for pkg in self.config.get("packages") or EXTRA_PACKAGES:62 out.setdefault(_norm(pkg), None)63 return out6465 async def discover(self, ctx: RunContext) -> list[Target]:66 return [Target(url=f"https://pypi.org/pypi/{pkg}/json", doc_type="package", key=f"pkg:{pkg}", min_bytes=200,67 meta={"pypi": pkg, "repo": (entry or {}).get("repo"), "content_type": "application/json"}, priority=2)68 for pkg, entry in self.packages().items()]6970 async def extract(self, ctx: RunContext, target: Target, res: FetchResult, parsed: Parsed) -> Facts:71 facts = Facts()72 data = parsed.json if parsed.kind == "json" else None73 if not isinstance(data, dict) or not isinstance(data.get("info"), dict):74 return facts75 info = data["info"]76 pkg = _norm(info.get("name") or target.meta.get("pypi") or "")77 if not pkg:78 return facts79 entry = next((e for e in repo_entries() if e.get("pypi") and _norm(e["pypi"]) == pkg), None)80 urls = {k.lower(): v for k, v in (info.get("project_urls") or {}).items() if isinstance(v, str)}81 gh = _github(urls, info.get("home_page"))82 if entry:83 ref = repo_entity(facts, entry)84 else:85 org = None86 if gh:87 known = org_by_github(gh.split("/")[0])88 org = org_ref(known["key"]) if known else None89 ids = {"pypi": pkg}90 if gh and _norm(gh.split("/")[1]) == pkg: # `gguf` links to ggml-org/llama.cpp: a dependency, not the same project91 ids["github_repo"] = gh92 ref = facts.entity("library", info.get("name") or pkg, identifiers=ids, organization=org, aliases=[pkg])93 if org:94 facts.entities.append(org)95 facts.relate(org, "develops", ref)96 if gh:97 facts.claim(ref, "repository_url", f"https://github.com/{gh}")98 facts.claim(ref, "pypi", pkg)99 facts.claim(ref, "pypi_url", info.get("package_url") or f"https://pypi.org/project/{pkg}/")100 version = info.get("version")101 released = _release_time(data, version)102 facts.claim(ref, "pypi_version", version)103 facts.claim(ref, "pypi_release_at", released)104 if not entry:105 facts.claim(ref, "latest_version", version)106 facts.claim(ref, "latest_release_at", released)107 lic = _license(info)108 facts.claim(ref, "pypi_license", lic)109 if not entry:110 facts.claim(ref, "license", lic) # mapped repos: GitHub's SPDX id owns `license` (avoids same-tier flip-flops)111 summary = (info.get("summary") or "").strip()112 facts.claim(ref, "description", summary[:1000] or None)113 homepage = urls.get("homepage") or urls.get("home") or info.get("home_page")114 if not entry and isinstance(homepage, str) and homepage.startswith("http"):115 facts.claim(ref, "homepage", homepage)116 docs = urls.get("documentation") or urls.get("docs")117 facts.claim(ref, "docs_url", docs if isinstance(docs, str) and docs.startswith("http") else None)118 facts.claim(ref, "requires_python", info.get("requires_python") or None)119 facts.claim(ref, "author", (info.get("author") or "").strip() or None)120 facts.claim(ref, "metric.releases", len(data.get("releases") or {}))121 facts.claim(ref, "python_classifiers", [c for c in info.get("classifiers") or [] if c.startswith("Programming Language :: Python :: 3.")][:12])122 facts.document_entity = ref123 facts.document_title = f"{info.get('name') or pkg} {version or ''} on PyPI".strip()124 return facts125126127def _norm(name: str) -> str:128 return re.sub(r"[-_.]+", "-", name.strip()).lower()129130131def _github(urls: dict[str, str], home: Any) -> str | None:132 for key in ("source", "source code", "repository", "code", "github", "homepage", "home"):133 v = urls.get(key)134 if v:135 m = GITHUB_URL.search(v)136 if m:137 return m.group(1).removesuffix(".git")138 if isinstance(home, str):139 m = GITHUB_URL.search(home)140 if m:141 return m.group(1).removesuffix(".git")142 return None143144145def _license(info: dict[str, Any]) -> str | None:146 expr = info.get("license_expression")147 if isinstance(expr, str) and expr.strip():148 return expr.strip()149 lic = info.get("license")150 if isinstance(lic, str) and 0 < len(lic.strip()) <= 80:151 return lic.strip()152 for c in info.get("classifiers") or []:153 if c.startswith("License :: OSI Approved :: "):154 return c.rsplit(" :: ", 1)[1]155 if c.startswith("License :: ") and c != "License :: OSI Approved":156 return c.rsplit(" :: ", 1)[1]157 return None158159160def _release_time(data: dict[str, Any], version: str | None) -> str | None:161 files = (data.get("releases") or {}).get(version or "") or data.get("urls") or []162 times = [parse_datetime(f.get("upload_time_iso_8601") or f.get("upload_time")) for f in files if isinstance(f, dict)]163 times = [t for t in times if t]164 return min(times).astimezone(UTC).isoformat(timespec="seconds") if times else None165166167__all__ = ["EXTRA_PACKAGES", "PyPIConnector"]168_ = EntityRef # re-exported type for tests169CONNECTORS = [PyPIConnector]170