feat: add LEANN connector support #183
Conversation
|
|
zxqfd555
left a comment
zxqfd555
left a comment
There was a problem hiding this comment.
Thank you for your interest in Pathway!
I have left several comments: could you please address them?
|
I think the change also deserves a mention in the CHANGELOG :) |
|
issues resolved bc6044b 👍🏾 |
zxqfd555
left a comment
zxqfd555
left a comment
There was a problem hiding this comment.
Thank you for updating the PR!
I have several minor concerns to be addressed before the merge:
- In Python parts of the connectors we use
pw.ColumnReferenceto specify columns. It's easier this way to avoid typos that may happen if one specifies the name of the column as a string. It also allows for passing a transformed column, which may be useful in metadata formation. - I have taken a look at LEANN repo, and there is a way to pass metadata. I propose to support it. If the metadata isn't supported elsewhere, let's document it in more detail.
- The black formatter currently fails in the precommit check. Could you please fix it?
There are several documentation comments as well, but nothing large.
| def write( | ||
| table: Table, | ||
| index_path: str | os.PathLike, | ||
| text_column: str = "text", |
There was a problem hiding this comment.
Let's use pw.ColumnReference here; it will be a more idiomatic way to pass the column. This way, it will be impossible to pass a column that doesn't belong to a table.
When a column reference is passed, you can use text_column.name to get the name of a column, and use something like this to check that the reference indeed belongs to a table.
| - Existing index files are overwritten on each build. | ||
| - LEANN requires the ``leann`` package. Please visit | ||
| https://github.com/yichuan-w/LEANN for installation instructions. | ||
| - metadata_columns are collected but not yet indexed (pending LEANN API support). |
There was a problem hiding this comment.
At a glance, LEANN builder has a metadata parameter in add_text. Could you clarify more about the extent of the limitation? If it isn't supported yet, we can probably remove it so as not to mislead the users, and add a TODO.
Besides, if we have metadata columns, I would propose to use ColumnReference to pass them, same as I've proposed with text_column.
| index_path: Path where the LEANN index will be saved (.leann file). | ||
| text_column: Name of the column containing text to index. | ||
| metadata_columns: List of additional columns to include as metadata. | ||
| backend_name: LEANN backend - "hnsw" or "diskann". |
There was a problem hiding this comment.
I would suggest quoting everything that can be passed as a parameter to the method, or more generally speaking, used as a code.
| backend_name: LEANN backend - "hnsw" or "diskann". | |
| backend_name: LEANN backend - ``"hnsw"`` or ``"diskann"``. |
| embedding_mode: Embedding provider - "sentence-transformers", "openai", | ||
| "mlx", or "ollama". | ||
| embedding_model: Specific embedding model name (e.g., "facebook/contriever"). |
There was a problem hiding this comment.
| embedding_mode: Embedding provider - "sentence-transformers", "openai", | |
| "mlx", or "ollama". | |
| embedding_model: Specific embedding model name (e.g., "facebook/contriever"). | |
| embedding_mode: Embedding provider - ``"sentence-transformers"``, ``"openai"``, | |
| ``"mlx"``, or ``"ollama"``. | |
| embedding_model: Specific embedding model name (e.g., ``"facebook/contriever"``). |
| embedding_mode: Embedding provider - "sentence-transformers", "openai", | ||
| "mlx", or "ollama". | ||
| embedding_model: Specific embedding model name (e.g., "facebook/contriever"). | ||
| embedding_options: Additional embedding config (api_key, base_url, etc.). |
There was a problem hiding this comment.
Is there any list for what these parameters may be?
If there is no such list, we can just leave it as it is, otherwise I'd propose linking it.
| for doc in self.documents.values(): | ||
| # TODO: Pass metadata to LEANN when the API supports it. | ||
| # Currently metadata is collected but not indexed. | ||
| builder.add_text(doc["text"]) |
There was a problem hiding this comment.
If we can pass metadata, should it probably be as follows?
| builder.add_text(doc["text"]) | |
| builder.add_text(**doc) |
3 participants
Introduction
To contribute code to the Pathway project, start by discussing your proposed changes on Discord or by filing an issue.
Once approved, follow the fork + pull request model against the main branch, ensuring you've signed the contributor license agreement.
Context
This PR adds a LEANN output connector (
pw.io.leann.write) to support LEANN vector indices in RAG pipelines. LEANN is a storage-efficient vector database that uses graph-based selective recomputation to achieve 97% storage reduction compared to traditional solutions.As discussed in issue #173, LEANN is implemented as a Python output connector following the established pattern from issue #167. The connector accumulates documents and builds the LEANN index on pipeline completion, with support for both HNSW and DiskANN backends.
How has this been tested?
python/pathway/tests/test_io_leann.py): Test observer behavior including additions, deletions, empty text handling, metadata extraction, and index build logic with mocked LEANN dependencies.integration_tests/leann/test_leann.py): End-to-end tests verifying index creation, search functionality, metadata columns, DiskANN backend, custom text columns, directory creation, and edge cases like empty tables.Types of changes
Checklist: