Retrieval

Protocol

DocumentRetriever

Bases: Protocol

Protocol for document retrievers used in schema data assessment.

Implementors fetch a list of documents (any type) for a given query. The schema data assessment agent calls this to sample example documents that are used to evaluate whether the generated schema fits real-world data.

Adapters

HuggingFace

HuggingFaceRetriever

HuggingFaceRetriever(
    dataset_name: str,
    text_column: str,
    embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    max_documents: int | None = None,
    split: str = "train",
    batch_size: int = 256,
    device: str | None = None,
    index_path: str | Path | None = None,
    model_kwargs: dict | None = None,
)

Bases: _BaseHuggingFaceRetriever

DocumentRetriever backed by a HuggingFace dataset + FAISS index.

Embeddings are computed with a sentence-transformers model. The index is built lazily on first use. Pass index_path to persist the index to disk — subsequent instantiations will load it instead of re-embedding.

Parameters:

Name	Type	Description	Default
`dataset_name`	`str`	HuggingFace dataset identifier (e.g. `"JuDDGES/pl-court-raw"`).	required
`text_column`	`str`	Dataset column to embed and return as search target.	required
`embedding_model`	`str`	Sentence-transformers model name or path.	`'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'`
`max_documents`	`int \| None`	Number of rows to stream and index (`None` = no limit).	`None`
`split`	`str`	Dataset split to load.	`'train'`
`batch_size`	`int`	Batch size for embedding computation.	`256`
`device`	`str \| None`	Device for sentence-transformers (`None` = auto-detect).	`None`
`index_path`	`str \| Path \| None`	Directory to save/load the FAISS index and dataset arrow files.	`None`

MMLWRobertaV2Retriever

MMLWRobertaV2Retriever(
    dataset_name: str,
    text_column: str,
    max_documents: int | None = None,
    split: str = "train",
    batch_size: int = 256,
    device: str | None = None,
    index_path: str | Path | None = None,
    model_kwargs: dict | None = None,
)

Bases: _BaseHuggingFaceRetriever

DocumentRetriever using sdadas/mmlw-retrieval-roberta-large-v2 for Polish retrieval.

Optimised for information retrieval (NDCG@10 of 60.71 on PIRB). Queries are prefixed with [query]: as required by the model. Pass index_path to persist the index to disk — subsequent instantiations will load it instead of re-embedding.

Parameters:

Name	Type	Description	Default
`dataset_name`	`str`	HuggingFace dataset identifier.	required
`text_column`	`str`	Dataset column to embed and return as search target.	required
`max_documents`	`int \| None`	Number of rows to stream and index (`None` = no limit).	`None`
`split`	`str`	Dataset split to load.	`'train'`
`batch_size`	`int`	Batch size for embedding computation.	`256`
`device`	`str \| None`	Device for sentence-transformers (`None` = auto-detect).	`None`
`index_path`	`str \| Path \| None`	Directory to save/load the FAISS index and dataset arrow files.	`None`
`model_kwargs`	`dict \| None`	Extra kwargs forwarded to `SentenceTransformer`. Defaults to `{"model_kwargs": {"torch_dtype": "bfloat16"}}` as recommended.	`None`

Weaviate

WeaviateRetriever

WeaviateRetriever(
    collection_name: str,
    target_vector: Optional[str] = None,
    filters: Optional[dict[str, str]] = None,
    host: Optional[str] = None,
    port: Optional[int] = None,
    grpc_port: Optional[int] = None,
    api_key: Optional[str] = None,
)

DocumentRetriever backed by a Weaviate hybrid-search index.

Works with any Weaviate collection. Connection parameters are read from environment variables by default but can be overridden.

Parameters:

Name	Type	Description	Default
`collection_name`	`str`	Weaviate collection to query.	required
`target_vector`	`Optional[str]`	Named vector to use for hybrid search. `None` uses the collection default.	`None`
`filters`	`Optional[dict[str, str]]`	Property-equality filters applied to every query, e.g. `{"language": "pl", "document_type": "judgment"}`.	`None`
`host`	`Optional[str]`	Weaviate host (default: `WV_URL` env var or `"localhost"`).	`None`
`port`	`Optional[int]`	Weaviate HTTP port (default: `WV_PORT` env var or `8080`).	`None`
`grpc_port`	`Optional[int]`	Weaviate gRPC port (default: `WV_GRPC_PORT` env var or `50051`).	`None`
`api_key`	`Optional[str]`	Weaviate API key (default: `WV_API_KEY` env var).	`None`