Extending supernova¶

# supernova/sources/common_crawl.py
from typing import Iterator
from supernova.models import Record
from supernova.sources.base import DatasetSource


class CommonCrawlSource(DatasetSource):
    def __init__(self, snapshot: str, path_filter: str | None = None):
        self.snapshot = snapshot
        self.path_filter = path_filter

    @property
    def source_name(self) -> str:
        return f"common-crawl/{self.snapshot}"

    def stream(self) -> Iterator[dict]:
        # Yield raw rows. Pull from WARC files, an HTTP feed, whatever.
        for record in self._iter_warc_records():
            yield {"url": record.url, "text": record.payload}

    def format_record(self, row: dict) -> Record:
        # Decide what becomes `text` and which fields tag along as columns.
        return Record(text=row["text"], columns={"url": row["url"]})

    def get_total_rows(self) -> int:
        # Used for distributed slicing. If unknown, return a best estimate or
        # implement total_rows_override via the config (see HuggingFaceSource).
        return self._count_rows()

SOURCE_REGISTRY = {
    "huggingface": HuggingFaceSource,
    "huggingface_parquet": HuggingFaceSource,   # alias of huggingface
    "common_crawl": CommonCrawlSource,          # new
}

# supernova/destinations.py

@dataclass(frozen=True)
class GsDestination:
    bucket: str
    prefix: str

    @property
    def scheme(self) -> str:
        return "gs"

    @property
    def root_uri(self) -> str:
        return f"gs://{self.bucket}/{self.prefix}".rstrip("/")

    def child_uri(self, sub: str) -> str:
        sub = sub.lstrip("/")
        if not self.prefix:
            return f"gs://{self.bucket}/{sub}"
        return f"gs://{self.bucket}/{self.prefix}/{sub}"

    def eval_uri(self, filename: str) -> str:
        return self.child_uri(f"{EVAL_SUBDIR}/{filename}")

def parse_destination(uri: str) -> Destination:
    if uri.startswith("s3://"):
        ...
    if uri.startswith("hf://buckets/"):
        ...
    if uri.startswith("gs://"):
        rest = uri[len("gs://"):]
        bucket, _, prefix = rest.partition("/")
        if not bucket:
            raise ValueError(f"gs:// URI is missing bucket: {uri!r}")
        return GsDestination(bucket=bucket, prefix=prefix.rstrip("/"))
    raise ValueError(...)


def _discover_gs(dest: GsDestination) -> list[str]:
    from google.cloud import storage
    client = storage.Client()
    eval_segment = f"/{EVAL_SUBDIR}/"
    uris = []
    for blob in client.list_blobs(dest.bucket, prefix=dest.prefix):
        if blob.name.endswith(".parquet") and eval_segment not in blob.name:
            uris.append(f"gs://{dest.bucket}/{blob.name}")
    return sorted(uris)


def filesystem_for_uri(uri: str):
    ...
    if uri.startswith("gs://"):
        from gcsfs import GCSFileSystem
        return GCSFileSystem()


def bare_key_for_uri(uri: str) -> str:
    ...
    if uri.startswith("gs://"):
        rest = uri[len("gs://"):]
        _, _, key = rest.partition("/")
        return key

from supernova.storage.base import StorageBackend


class GcsBackend(StorageBackend):
    def __init__(self, bucket: str, prefix: str):
        self.bucket = bucket
        self.prefix = prefix
        self._client = None
        self._ready = False

    @property
    def destination(self) -> str:
        return f"gs://{self.bucket}/{self.prefix}"

    async def ensure_ready(self) -> None:
        # Create the bucket if it doesn't exist, etc. async-safe via asyncio.to_thread.

    async def upload_file(self, local_path: str, remote_subpath: str | None = None) -> None:
        # Upload to f"{self.prefix}/{remote_subpath or basename}".
        # Honor remote_subpath verbatim (keeps shard_by_rank paths like rank00/batch_*.parquet).

    async def upload_bytes(self, data: bytes, filename: str) -> None:
        # Used for manifest JSON. Per-backend convention: where does it live?
        # S3Backend and HuggingFaceBackend both put it under the prefix alongside
        # the parquets — bucket URIs are flat, no auto-detection subdir needed.

def build_storage(cfg: dict):
    storage_type = cfg.pop("type", "s3")
    if storage_type == "s3":
        return S3Backend(...)
    elif storage_type == "hf":
        return HuggingFaceBackend(...)
    elif storage_type == "gcs":
        return GcsBackend(bucket=cfg["bucket"], prefix=cfg["prefix"])
    elif storage_type == "local":
        return LocalBackend(...)
    raise ValueError(...)

from typing import Iterable
from .base import DataReader


class GcsDataReader(DataReader):
    def __init__(self, gcs_bucket: str, gcs_prefix: str, file_list: list[str] | None = None,
                 id_expression: str = "row_id", vectors=None, payload_fields=None,
                 duckdb_memory_limit="2GB", duckdb_threads=2):
        super().__init__(
            id_expression=id_expression, vectors=vectors,
            payload_fields=payload_fields,
            duckdb_memory_limit=duckdb_memory_limit, duckdb_threads=duckdb_threads,
        )
        self.gcs_bucket = gcs_bucket
        self.gcs_prefix = gcs_prefix.rstrip("/")
        self.file_list = file_list

    @property
    def glob_path(self) -> str:
        return f"gs://{self.gcs_bucket}/{self.gcs_prefix}/**/*.parquet"

    def _root_uri_prefix(self) -> str:
        # The base class uses this to register vf_point_id. The bare key fed
        # to make_point_id is everything after this prefix — must agree with
        # bare_key_for_uri("gs://...") on the brute-force / query-gen side.
        return f"gs://{self.gcs_bucket}/"

    def _iter_sources(self) -> Iterable[str]:
        suffix = self._parquet_kwargs
        if self.file_list:
            for f in self.file_list:
                yield f"read_parquet('{f}'{suffix})"
        elif self._parquet_kwargs:
            yield f"read_parquet('{self.glob_path}'{suffix})"
        else:
            yield f"'{self.glob_path}'"

    def _configure_connection(self) -> None:
        conn = self._conn
        conn.execute("INSTALL httpfs; LOAD httpfs;")
        # GCS auth — service account JSON, application default credentials, etc.
        # See S3DataReader / HuggingFaceDataReader for the auth-injection pattern.

DATASOURCE_REGISTRY = {
    "s3": S3DataReader,
    "huggingface": HuggingFaceDataReader,
    "gcs": GcsDataReader,
}

# supernova/loader/vectorstore/weaviate.py
from supernova.loader.vectorstore.base import VectorStore


class WeaviateVectorStore(VectorStore):
    def __init__(self, url: str, api_key: str | None = None,
                 collection_name: str = "default", vectors=None, params=None):
        self.url = url
        self.collection_name = collection_name
        self.vectors = vectors
        self.params = params or {}
        # construct client...

    @property
    def name(self) -> str:
        return f"weaviate({self.collection_name})"

    async def ensure_collection(self, dimensions: dict[str, int]) -> None:
        # Create collection if it doesn't exist. dimensions is {vector_name: dim}
        # for dense + multivector; sparse vectors are absent.

    async def upsert_batch(self, points: list[dict]) -> None:
        # Each point: {id: str, vectors: {name: value}, payload: dict}.
        # For multivector: value is list[list[float]]. For sparse: {"indices", "values"}.

    async def close(self) -> None:
        # Clean up the client.

    # Optional — implement these if your store has a fast-bulk-load mode.
    # No-op default in the base class is fine if it doesn't.
    # async def defer_indexing(self) -> None: ...
    # async def enable_indexing(self) -> None: ...
    # async def wait_for_indexing(self) -> None: ...

VECTORSTORE_REGISTRY = {
    "qdrant": QdrantVectorStore,
    "weaviate": WeaviateVectorStore,
}