From 859860318acb6a2a705bbffa5c5a8325bc4b8328 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 17 Jul 2025 13:17:16 +0300 Subject: [PATCH 01/10] Cohere reranker --- pyproject.toml | 1 + src/haiku/rag/config.py | 6 +++ src/haiku/rag/reranking/__init__.py | 26 +++++++++ src/haiku/rag/reranking/base.py | 13 +++++ src/haiku/rag/reranking/cohere.py | 34 ++++++++++++ tests/test_reranker.py | 83 +++++++++++++++++++++++++++++ uv.lock | 75 +++++++++++++++++++++++++- 7 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 src/haiku/rag/reranking/__init__.py create mode 100644 src/haiku/rag/reranking/base.py create mode 100644 src/haiku/rag/reranking/cohere.py create mode 100644 tests/test_reranker.py diff --git a/pyproject.toml b/pyproject.toml index 3bfd9784..16be0002 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ voyageai = ["voyageai>=0.3.2"] openai = ["openai>=1.0.0"] anthropic = ["anthropic>=0.56.0"] +cohere = ["cohere>=5.16.1"] [project.scripts] haiku-rag = "haiku.rag.cli:cli" diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index e1552873..e32799e1 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,6 +19,9 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 + RERANK_PROVIDER: str = "cohere" + RERANK_MODEL: str = "rerank-v3.5" + QA_PROVIDER: str = "ollama" QA_MODEL: str = "qwen3" @@ -31,6 +34,7 @@ class AppConfig(BaseModel): VOYAGE_API_KEY: str = "" OPENAI_API_KEY: str = "" ANTHROPIC_API_KEY: str = "" + COHERE_API_KEY: str = "" @field_validator("MONITOR_DIRECTORIES", mode="before") @classmethod @@ -52,3 +56,5 @@ if Config.VOYAGE_API_KEY: os.environ["VOYAGE_API_KEY"] = Config.VOYAGE_API_KEY if Config.ANTHROPIC_API_KEY: os.environ["ANTHROPIC_API_KEY"] = Config.ANTHROPIC_API_KEY +if Config.COHERE_API_KEY: + os.environ["CO_API_KEY"] = Config.COHERE_API_KEY diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py new file mode 100644 index 00000000..8fbdae31 --- /dev/null +++ b/src/haiku/rag/reranking/__init__.py @@ -0,0 +1,26 @@ +from haiku.rag.config import Config +from haiku.rag.reranking.base import RerankerBase + +try: + from haiku.rag.reranking.cohere import CohereReranker +except ImportError: + pass + + +def get_reranker() -> RerankerBase: + """ + Factory function to get the appropriate reranker based on the configuration. + """ + + if Config.RERANK_PROVIDER == "cohere": + try: + from haiku.rag.reranking.cohere import CohereReranker + except ImportError: + raise ImportError( + "Cohere reranker requires the 'cohere' package. " + "Please install haiku.rag with the 'cohere' extra:" + "uv pip install haiku.rag --extra cohere" + ) + return CohereReranker() + + raise ValueError(f"Unsupported reranker provider: {Config.RERANK_PROVIDER}") diff --git a/src/haiku/rag/reranking/base.py b/src/haiku/rag/reranking/base.py new file mode 100644 index 00000000..72b3a3df --- /dev/null +++ b/src/haiku/rag/reranking/base.py @@ -0,0 +1,13 @@ +from haiku.rag.config import Config +from haiku.rag.store.models.chunk import Chunk + + +class RerankerBase: + _model: str = Config.RERANK_MODEL + + async def rerank( + self, query: str, chunks: list[Chunk], top_n: int = 10 + ) -> list[Chunk]: + raise NotImplementedError( + "Reranker is an abstract class. Please implement the rerank method in a subclass." + ) diff --git a/src/haiku/rag/reranking/cohere.py b/src/haiku/rag/reranking/cohere.py new file mode 100644 index 00000000..19e91577 --- /dev/null +++ b/src/haiku/rag/reranking/cohere.py @@ -0,0 +1,34 @@ +from haiku.rag.config import Config +from haiku.rag.reranking.base import RerankerBase +from haiku.rag.store.models.chunk import Chunk + +try: + import cohere +except ImportError as e: + raise ImportError( + "cohere is not installed. Please install it with `pip install cohere` or use the cohere optional dependency." + ) from e + + +class CohereReranker(RerankerBase): + def __init__(self): + self._client = cohere.ClientV2(api_key=Config.COHERE_API_KEY) + + async def rerank( + self, query: str, chunks: list[Chunk], top_n: int = 10 + ) -> list[Chunk]: + if not chunks: + return [] + + documents = [chunk.content for chunk in chunks] + + response = self._client.rerank( + model=self._model, query=query, documents=documents, top_n=top_n + ) + + reranked_chunks = [] + for result in response.results: + original_chunk = chunks[result.index] + reranked_chunks.append(original_chunk) + + return reranked_chunks diff --git a/tests/test_reranker.py b/tests/test_reranker.py new file mode 100644 index 00000000..89e2996e --- /dev/null +++ b/tests/test_reranker.py @@ -0,0 +1,83 @@ +import pytest + +from haiku.rag.reranking import get_reranker +from haiku.rag.reranking.base import RerankerBase +from haiku.rag.store.models.chunk import Chunk + + +@pytest.mark.asyncio +async def test_reranker_base(): + reranker = RerankerBase() + assert reranker._model == "rerank-v3.5" + + with pytest.raises(NotImplementedError): + await reranker.rerank("query", []) + + +@pytest.mark.asyncio +async def test_cohere_reranker(): + try: + # Mock the client + class MockResult: + def __init__(self, index): + self.index = index + + class MockResponse: + def __init__(self, results): + self.results = results + + class MockClient: + def __init__(self, api_key=None): + pass + + def rerank(self, model, query, documents, top_n): + return MockResponse([MockResult(1), MockResult(0)]) + + import haiku.rag.reranking.cohere + + original_client = haiku.rag.reranking.cohere.cohere.ClientV2 + haiku.rag.reranking.cohere.cohere.ClientV2 = MockClient + + try: + from haiku.rag.reranking.cohere import CohereReranker + + reranker = CohereReranker() + assert reranker._model == "rerank-v3.5" + + chunks = [ + Chunk(id=1, content="First chunk", document_id=1), + Chunk(id=2, content="Second chunk", document_id=1), + ] + + result = await reranker.rerank("test query", chunks) + assert len(result) == 2 + assert result[0] == chunks[1] # Should return chunk at index 1 first + assert result[1] == chunks[0] # Should return chunk at index 0 second + finally: + haiku.rag.reranking.cohere.cohere.ClientV2 = original_client + + except ImportError: + pytest.skip("Cohere package not installed") + + +@pytest.mark.asyncio +async def test_get_reranker(): + try: + + class MockClient: + def __init__(self, api_key=None): + pass + + import haiku.rag.reranking.cohere + + original_client = haiku.rag.reranking.cohere.cohere.ClientV2 + haiku.rag.reranking.cohere.cohere.ClientV2 = MockClient + + try: + reranker = get_reranker() + assert reranker._model == "rerank-v3.5" + assert hasattr(reranker, "rerank") + finally: + haiku.rag.reranking.cohere.cohere.ClientV2 = original_client + except ImportError: + pytest.skip("Cohere package not installed") diff --git a/uv.lock b/uv.lock index d0ac66e2..e43acd5a 100644 --- a/uv.lock +++ b/uv.lock @@ -429,6 +429,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/e1/3714a2f371985215c219c2a70953d38e3eed81ef165aed061d21de0e998b/cobble-0.1.4-py3-none-any.whl", hash = "sha256:36c91b1655e599fd428e2b95fdd5f0da1ca2e9f1abb0bc871dec21a0e78a2b44", size = 3984, upload-time = "2024-06-01T18:11:07.911Z" }, ] +[[package]] +name = "cohere" +version = "5.16.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fastavro" }, + { name = "httpx" }, + { name = "httpx-sse" }, + { name = "pydantic" }, + { name = "pydantic-core" }, + { name = "requests" }, + { name = "tokenizers" }, + { name = "types-requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ed/c7/fd1e4c61cf3f0aac9d9d73fce63a766c9778e1270f7a26812eb289b4851d/cohere-5.16.1.tar.gz", hash = "sha256:02aa87668689ad0fbac2cda979c190310afdb99fb132552e8848fdd0aff7cd40", size = 162300, upload-time = "2025-07-09T20:47:36.348Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/c6/72309ac75f3567425ca31a601ad394bfee8d0f4a1569dfbc80cbb2890d07/cohere-5.16.1-py3-none-any.whl", hash = "sha256:37e2c1d69b1804071b5e5f5cb44f8b74127e318376e234572d021a1a729c6baa", size = 291894, upload-time = "2025-07-09T20:47:34.919Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -648,6 +668,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674, upload-time = "2025-05-10T17:42:49.33Z" }, ] +[[package]] +name = "fastavro" +version = "1.11.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/48/8f/32664a3245247b13702d13d2657ea534daf64e58a3f72a3a2d10598d6916/fastavro-1.11.1.tar.gz", hash = "sha256:bf6acde5ee633a29fb8dfd6dfea13b164722bc3adc05a0e055df080549c1c2f8", size = 1016250, upload-time = "2025-05-18T04:54:31.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/be/53df3fec7fdabc1848896a76afb0f01ab96b58abb29611aa68a994290167/fastavro-1.11.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:603aa1c1d1be21fb4bcb63e1efb0711a9ddb337de81391c32dac95c6e0dacfcc", size = 944225, upload-time = "2025-05-18T04:54:34.586Z" }, + { url = "https://files.pythonhosted.org/packages/d0/cc/c7c76a082fbf5aaaf82ab7da7b9ede6fc99eb8f008c084c67d230b29c446/fastavro-1.11.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:45653b312d4ce297e2bd802ea3ffd17ecbe718e5e8b6e2ae04cd72cb50bb99d5", size = 3105189, upload-time = "2025-05-18T04:54:36.855Z" }, + { url = "https://files.pythonhosted.org/packages/48/ff/5f1f0b5e3835e788ba8121d6dd6426cd4c6e58ce1bff02cb7810278648b0/fastavro-1.11.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:998a53fc552e6bee9acda32af258f02557313c85fb5b48becba5b71ec82f421e", size = 3113124, upload-time = "2025-05-18T04:54:40.013Z" }, + { url = "https://files.pythonhosted.org/packages/e5/b8/1ac01433b55460dabeb6d3fbb05ba1c971d57137041e8f53b2e9f46cd033/fastavro-1.11.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:9f878c9ad819467120cb066f1c73496c42eb24ecdd7c992ec996f465ef4cedad", size = 3155196, upload-time = "2025-05-18T04:54:42.307Z" }, + { url = "https://files.pythonhosted.org/packages/5e/a8/66e599b946ead031a5caba12772e614a7802d95476e8732e2e9481369973/fastavro-1.11.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:da9e4c231ac4951092c2230ca423d8a3f2966718f072ac1e2c5d2d44c70b2a50", size = 3229028, upload-time = "2025-05-18T04:54:44.503Z" }, + { url = "https://files.pythonhosted.org/packages/0e/e7/17c35e2dfe8a9e4f3735eabdeec366b0edc4041bb1a84fcd528c8efd12af/fastavro-1.11.1-cp310-cp310-win_amd64.whl", hash = "sha256:7423bfad3199567eeee7ad6816402c7c0ee1658b959e8c10540cfbc60ce96c2a", size = 449177, upload-time = "2025-05-18T04:54:46.127Z" }, + { url = "https://files.pythonhosted.org/packages/8e/63/f33d6fd50d8711f305f07ad8c7b4a25f2092288f376f484c979dcf277b07/fastavro-1.11.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:3573340e4564e8962e22f814ac937ffe0d4be5eabbd2250f77738dc47e3c8fe9", size = 957526, upload-time = "2025-05-18T04:54:47.701Z" }, + { url = "https://files.pythonhosted.org/packages/f4/09/a57ad9d8cb9b8affb2e43c29d8fb8cbdc0f1156f8496067a0712c944bacc/fastavro-1.11.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7291cf47735b8bd6ff5d9b33120e6e0974f52fd5dff90cd24151b22018e7fd29", size = 3322808, upload-time = "2025-05-18T04:54:50.419Z" }, + { url = "https://files.pythonhosted.org/packages/86/70/d6df59309d3754d6d4b0c7beca45b9b1a957d6725aed8da3aca247db3475/fastavro-1.11.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bf3bb065d657d5bac8b2cb39945194aa086a9b3354f2da7f89c30e4dc20e08e2", size = 3330870, upload-time = "2025-05-18T04:54:52.406Z" }, + { url = "https://files.pythonhosted.org/packages/ad/ea/122315154d2a799a2787058435ef0d4d289c0e8e575245419436e9b702ca/fastavro-1.11.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8758317c85296b848698132efb13bc44a4fbd6017431cc0f26eaeb0d6fa13d35", size = 3343369, upload-time = "2025-05-18T04:54:54.652Z" }, + { url = "https://files.pythonhosted.org/packages/62/12/7800de5fec36d55a818adf3db3b085b1a033c4edd60323cf6ca0754cf8cb/fastavro-1.11.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:ad99d57228f83bf3e2214d183fbf6e2fda97fd649b2bdaf8e9110c36cbb02624", size = 3430629, upload-time = "2025-05-18T04:54:56.513Z" }, + { url = "https://files.pythonhosted.org/packages/48/65/2b74ccfeba9dcc3f7dbe64907307386b4a0af3f71d2846f63254df0f1e1d/fastavro-1.11.1-cp311-cp311-win_amd64.whl", hash = "sha256:9134090178bdbf9eefd467717ced3dc151e27a7e7bfc728260ce512697efe5a4", size = 451621, upload-time = "2025-05-18T04:54:58.156Z" }, + { url = "https://files.pythonhosted.org/packages/99/58/8e789b0a2f532b22e2d090c20d27c88f26a5faadcba4c445c6958ae566cf/fastavro-1.11.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e8bc238f2637cd5d15238adbe8fb8c58d2e6f1870e0fb28d89508584670bae4b", size = 939583, upload-time = "2025-05-18T04:54:59.853Z" }, + { url = "https://files.pythonhosted.org/packages/34/3f/02ed44742b1224fe23c9fc9b9b037fc61769df716c083cf80b59a02b9785/fastavro-1.11.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b403933081c83fc4d8a012ee64b86e560a024b1280e3711ee74f2abc904886e8", size = 3257734, upload-time = "2025-05-18T04:55:02.366Z" }, + { url = "https://files.pythonhosted.org/packages/cc/bc/9cc8b19eeee9039dd49719f8b4020771e805def262435f823fa8f27ddeea/fastavro-1.11.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3f6ecb4b5f77aa756d973b7dd1c2fb4e4c95b4832a3c98b059aa96c61870c709", size = 3318218, upload-time = "2025-05-18T04:55:04.352Z" }, + { url = "https://files.pythonhosted.org/packages/39/77/3b73a986606494596b6d3032eadf813a05b59d1623f54384a23de4217d5f/fastavro-1.11.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:059893df63ef823b0231b485c9d43016c7e32850cae7bf69f4e9d46dd41c28f2", size = 3297296, upload-time = "2025-05-18T04:55:06.175Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1c/b69ceef6494bd0df14752b5d8648b159ad52566127bfd575e9f5ecc0c092/fastavro-1.11.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5120ffc9a200699218e01777e695a2f08afb3547ba818184198c757dc39417bd", size = 3438056, upload-time = "2025-05-18T04:55:08.276Z" }, + { url = "https://files.pythonhosted.org/packages/ef/11/5c2d0db3bd0e6407546fabae9e267bb0824eacfeba79e7dd81ad88afa27d/fastavro-1.11.1-cp312-cp312-win_amd64.whl", hash = "sha256:7bb9d0d2233f33a52908b6ea9b376fe0baf1144bdfdfb3c6ad326e200a8b56b0", size = 442824, upload-time = "2025-05-18T04:55:10.385Z" }, + { url = "https://files.pythonhosted.org/packages/ec/08/8e25b9e87a98f8c96b25e64565fa1a1208c0095bb6a84a5c8a4b925688a5/fastavro-1.11.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f963b8ddaf179660e814ab420850c1b4ea33e2ad2de8011549d958b21f77f20a", size = 931520, upload-time = "2025-05-18T04:55:11.614Z" }, + { url = "https://files.pythonhosted.org/packages/02/ee/7cf5561ef94781ed6942cee6b394a5e698080f4247f00f158ee396ec244d/fastavro-1.11.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0253e5b6a3c9b62fae9fc3abd8184c5b64a833322b6af7d666d3db266ad879b5", size = 3195989, upload-time = "2025-05-18T04:55:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/b3/31/f02f097d79f090e5c5aca8a743010c4e833a257c0efdeb289c68294f7928/fastavro-1.11.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca637b150e1f4c0e8e564fad40a16bd922bcb7ffd1a6e4836e6084f2c4f4e8db", size = 3239755, upload-time = "2025-05-18T04:55:16.463Z" }, + { url = "https://files.pythonhosted.org/packages/09/4c/46626b4ee4eb8eb5aa7835973c6ba8890cf082ef2daface6071e788d2992/fastavro-1.11.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:76af1709031621828ca6ce7f027f7711fa33ac23e8269e7a5733996ff8d318da", size = 3243788, upload-time = "2025-05-18T04:55:18.544Z" }, + { url = "https://files.pythonhosted.org/packages/a7/6f/8ed42524e9e8dc0554f0f211dd1c6c7a9dde83b95388ddcf7c137e70796f/fastavro-1.11.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8224e6d8d9864d4e55dafbe88920d6a1b8c19cc3006acfac6aa4f494a6af3450", size = 3378330, upload-time = "2025-05-18T04:55:20.887Z" }, + { url = "https://files.pythonhosted.org/packages/b8/51/38cbe243d5facccab40fc43a4c17db264c261be955ce003803d25f0da2c3/fastavro-1.11.1-cp313-cp313-win_amd64.whl", hash = "sha256:cde7ed91b52ff21f0f9f157329760ba7251508ca3e9618af3ffdac986d9faaa2", size = 443115, upload-time = "2025-05-18T04:55:22.107Z" }, + { url = "https://files.pythonhosted.org/packages/d0/57/0d31ed1a49c65ad9f0f0128d9a928972878017781f9d4336f5f60982334c/fastavro-1.11.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e5ed1325c1c414dd954e7a2c5074daefe1eceb672b8c727aa030ba327aa00693", size = 1021401, upload-time = "2025-05-18T04:55:23.431Z" }, + { url = "https://files.pythonhosted.org/packages/56/7a/a3f1a75fbfc16b3eff65dc0efcdb92364967923194312b3f8c8fc2cb95be/fastavro-1.11.1-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8cd3c95baeec37188899824faf44a5ee94dfc4d8667b05b2f867070c7eb174c4", size = 3384349, upload-time = "2025-05-18T04:55:25.575Z" }, + { url = "https://files.pythonhosted.org/packages/be/84/02bceb7518867df84027232a75225db758b9b45f12017c9743f45b73101e/fastavro-1.11.1-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2e0babcd81acceb4c60110af9efa25d890dbb68f7de880f806dadeb1e70fe413", size = 3240658, upload-time = "2025-05-18T04:55:27.633Z" }, + { url = "https://files.pythonhosted.org/packages/f2/17/508c846c644d39bc432b027112068b8e96e7560468304d4c0757539dd73a/fastavro-1.11.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:b2c0cb8063c7208b53b6867983dc6ae7cc80b91116b51d435d2610a5db2fc52f", size = 3372809, upload-time = "2025-05-18T04:55:30.063Z" }, + { url = "https://files.pythonhosted.org/packages/fe/84/9c2917a70ed570ddbfd1d32ac23200c1d011e36c332e59950d2f6d204941/fastavro-1.11.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:1bc2824e9969c04ab6263d269a1e0e5d40b9bd16ade6b70c29d6ffbc4f3cc102", size = 3387171, upload-time = "2025-05-18T04:55:32.531Z" }, +] + [[package]] name = "fastmcp" version = "2.8.1" @@ -836,6 +893,9 @@ dependencies = [ anthropic = [ { name = "anthropic" }, ] +cohere = [ + { name = "cohere" }, +] openai = [ { name = "openai" }, ] @@ -859,6 +919,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.56.0" }, + { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.16.1" }, { name = "fastmcp", specifier = ">=2.8.1" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" }, @@ -873,7 +934,7 @@ requires-dist = [ { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "watchfiles", specifier = ">=1.1.0" }, ] -provides-extras = ["voyageai", "openai", "anthropic"] +provides-extras = ["voyageai", "openai", "anthropic", "cohere"] [package.metadata.requires-dev] dev = [ @@ -2849,6 +2910,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/76/42/3efaf858001d2c2913de7f354563e3a3a2f0decae3efe98427125a8f441e/typer-0.16.0-py3-none-any.whl", hash = "sha256:1f79bed11d4d02d4310e3c1b7ba594183bcedb0ac73b27a9e5f28f6fb5b98855", size = 46317, upload-time = "2025-05-26T14:30:30.523Z" }, ] +[[package]] +name = "types-requests" +version = "2.32.4.20250611" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/7f/73b3a04a53b0fd2a911d4ec517940ecd6600630b559e4505cc7b68beb5a0/types_requests-2.32.4.20250611.tar.gz", hash = "sha256:741c8777ed6425830bf51e54d6abe245f79b4dcb9019f1622b773463946bf826", size = 23118, upload-time = "2025-06-11T03:11:41.272Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3d/ea/0be9258c5a4fa1ba2300111aa5a0767ee6d18eb3fd20e91616c12082284d/types_requests-2.32.4.20250611-py3-none-any.whl", hash = "sha256:ad2fe5d3b0cb3c2c902c8815a70e7fb2302c4b8c1f77bdcd738192cdb3878072", size = 20643, upload-time = "2025-06-11T03:11:40.186Z" }, +] + [[package]] name = "typing-extensions" version = "4.14.0" From 0f8124fd98b457354b86e31585f5a8118747cf5e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 17 Jul 2025 13:17:47 +0300 Subject: [PATCH 02/10] Simply embedders --- src/haiku/rag/embeddings/base.py | 7 +++++-- src/haiku/rag/embeddings/ollama.py | 3 --- src/haiku/rag/embeddings/openai.py | 4 ---- src/haiku/rag/embeddings/voyageai.py | 4 ---- 4 files changed, 5 insertions(+), 13 deletions(-) diff --git a/src/haiku/rag/embeddings/base.py b/src/haiku/rag/embeddings/base.py index 16e19d9c..369a53a6 100644 --- a/src/haiku/rag/embeddings/base.py +++ b/src/haiku/rag/embeddings/base.py @@ -1,6 +1,9 @@ +from haiku.rag.config import Config + + class EmbedderBase: - _model: str = "" - _vector_dim: int = 0 + _model: str = Config.EMBEDDINGS_MODEL + _vector_dim: int = Config.EMBEDDINGS_VECTOR_DIM def __init__(self, model: str, vector_dim: int): self._model = model diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index d7aa97a7..600afe65 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/src/haiku/rag/embeddings/ollama.py @@ -5,9 +5,6 @@ from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1024 - async def embed(self, text: str) -> list[float]: client = AsyncClient(host=Config.OLLAMA_BASE_URL) res = await client.embeddings(model=self._model, prompt=text) diff --git a/src/haiku/rag/embeddings/openai.py b/src/haiku/rag/embeddings/openai.py index 024705cd..818f0e5b 100644 --- a/src/haiku/rag/embeddings/openai.py +++ b/src/haiku/rag/embeddings/openai.py @@ -1,13 +1,9 @@ try: from openai import AsyncOpenAI - from haiku.rag.config import Config from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1536 - async def embed(self, text: str) -> list[float]: client = AsyncOpenAI() response = await client.embeddings.create( diff --git a/src/haiku/rag/embeddings/voyageai.py b/src/haiku/rag/embeddings/voyageai.py index d37378c7..ac7aa1b6 100644 --- a/src/haiku/rag/embeddings/voyageai.py +++ b/src/haiku/rag/embeddings/voyageai.py @@ -1,13 +1,9 @@ try: from voyageai.client import Client # type: ignore - from haiku.rag.config import Config from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1024 - async def embed(self, text: str) -> list[float]: client = Client() res = client.embed([text], model=self._model, output_dtype="float") From 044e1b0a06aab020ce2bd92a2eb93b6cdbe72d7a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 18 Jul 2025 20:29:24 +0300 Subject: [PATCH 03/10] mixedbread reranker as default --- pyproject.toml | 1 + src/haiku/rag/config.py | 6 +- src/haiku/rag/reranking/mxbai.py | 28 +++ tests/test_reranker.py | 26 ++- uv.lock | 342 +++++++++++++++++++++++++++++++ 5 files changed, 400 insertions(+), 3 deletions(-) create mode 100644 src/haiku/rag/reranking/mxbai.py diff --git a/pyproject.toml b/pyproject.toml index 16be0002..1ce95919 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "fastmcp>=2.8.1", "httpx>=0.28.1", "markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2", + "mxbai-rerank>=0.1.6", "ollama>=0.5.1", "pydantic>=2.11.7", "python-dotenv>=1.1.0", diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index e32799e1..6151a80c 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,8 +19,10 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 - RERANK_PROVIDER: str = "cohere" - RERANK_MODEL: str = "rerank-v3.5" + # RERANK_PROVIDER: str = "cohere" + # RERANK_MODEL: str = "rerank-v3.5" + RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" + RERANK_PROVIDER: str = "mxbai" QA_PROVIDER: str = "ollama" QA_MODEL: str = "qwen3" diff --git a/src/haiku/rag/reranking/mxbai.py b/src/haiku/rag/reranking/mxbai.py new file mode 100644 index 00000000..135727ed --- /dev/null +++ b/src/haiku/rag/reranking/mxbai.py @@ -0,0 +1,28 @@ +from mxbai_rerank import MxbaiRerankV2 + +from haiku.rag.config import Config +from haiku.rag.reranking.base import RerankerBase +from haiku.rag.store.models.chunk import Chunk + + +class MxBAIReranker(RerankerBase): + def __init__(self): + self._client = MxbaiRerankV2( + Config.RERANK_MODEL, disable_transformers_warnings=True + ) + + async def rerank( + self, query: str, chunks: list[Chunk], top_n: int = 10 + ) -> list[Chunk]: + if not chunks: + return [] + + documents = [chunk.content for chunk in chunks] + + results = self._client.rank(query=query, documents=documents, top_k=top_n) + reranked_chunks = [] + for result in results: + original_chunk = chunks[result.index] + reranked_chunks.append(original_chunk) + + return reranked_chunks diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 89e2996e..39a515c1 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -2,18 +2,42 @@ import pytest from haiku.rag.reranking import get_reranker from haiku.rag.reranking.base import RerankerBase +from haiku.rag.reranking.mxbai import MxBAIReranker from haiku.rag.store.models.chunk import Chunk @pytest.mark.asyncio async def test_reranker_base(): reranker = RerankerBase() - assert reranker._model == "rerank-v3.5" + assert reranker._model == "mixedbread-ai/mxbai-rerank-base-v2" with pytest.raises(NotImplementedError): await reranker.rerank("query", []) +@pytest.mark.asyncio +async def test_mxbai_reranker(): + reranker = MxBAIReranker() + chunks = [ + Chunk(content=content, document_id=i) + for i, content in enumerate( + [ + "To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.", + "The novel Moby-Dick was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.", + "Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961.", + "Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment upon the British landed gentry at the end of the 18th century.", + "The Harry Potter series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era.", + "The Great Gatsby, a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.", + ] + ) + ] + + reranked = await reranker.rerank( + "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 + ) + assert [r.document_id for r in reranked] == [0, 2] + + @pytest.mark.asyncio async def test_cohere_reranker(): try: diff --git a/uv.lock b/uv.lock index e43acd5a..d01bdaad 100644 --- a/uv.lock +++ b/uv.lock @@ -8,6 +8,25 @@ resolution-markers = [ "python_full_version < '3.11'", ] +[[package]] +name = "accelerate" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/25/969456a95a90ed38f73f68d0f0915bdf1d76145d05054c59ad587b171150/accelerate-1.9.0.tar.gz", hash = "sha256:0e8c61f81af7bf37195b6175a545ed292617dd90563c88f49020aea5b6a0b47f", size = 383234, upload-time = "2025-07-16T16:24:54.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/1c/a17fb513aeb684fb83bef5f395910f53103ab30308bbdd77fd66d6698c46/accelerate-1.9.0-py3-none-any.whl", hash = "sha256:c24739a97ade1d54af4549a65f8b6b046adc87e2b3e4d6c66516e32c53d5a8f1", size = 367073, upload-time = "2025-07-16T16:24:52.957Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -259,6 +278,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/ff/392bff89415399a979be4a65357a41d92729ae8580a66073d8ec8d810f98/backrefs-5.9-py39-none-any.whl", hash = "sha256:f48ee18f6252b8f5777a22a00a09a85de0ca931658f1dd96d4406a34f3748c60", size = 380265, upload-time = "2025-06-22T19:34:12.405Z" }, ] +[[package]] +name = "batched" +version = "0.1.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c9/40/8d9a8ed9b95cb95acf599698557b7074b462df652823a61e7e43899aa519/batched-0.1.5.tar.gz", hash = "sha256:58b8a41d3f8d4d39a0edba79c6238ed204938cfc2c8908224919d70af07c610d", size = 23940, upload-time = "2025-07-14T09:58:31.862Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/c8/16a977fd90cdc974ef7781e237b8a0e0008a6204768ededbef2b2ff1bb43/batched-0.1.5-py3-none-any.whl", hash = "sha256:356dae99f15c906629992e4bd3481a857114790b5316268fa38fe8ad0d0b9480", size = 29367, upload-time = "2025-07-14T09:58:30.968Z" }, +] + [[package]] name = "beautifulsoup4" version = "4.13.4" @@ -879,6 +907,7 @@ dependencies = [ { name = "fastmcp" }, { name = "httpx" }, { name = "markitdown", extra = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"] }, + { name = "mxbai-rerank" }, { name = "ollama" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -923,6 +952,7 @@ requires-dist = [ { name = "fastmcp", specifier = ">=2.8.1" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" }, + { name = "mxbai-rerank", specifier = ">=0.1.6" }, { name = "ollama", specifier = ">=0.5.1" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.0.0" }, { name = "pydantic", specifier = ">=2.11.7" }, @@ -1616,6 +1646,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/d9/f7f9379981e39b8c2511c9e0326d212accacb82f12fbfdc1aa2ce2a7b2b6/multiprocess-0.70.16-py39-none-any.whl", hash = "sha256:a0bafd3ae1b732eac64be2e72038231c1ba97724b60b09400d68f229fcc2fbf3", size = 133351, upload-time = "2024-01-28T18:52:31.981Z" }, ] +[[package]] +name = "mxbai-rerank" +version = "0.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "accelerate" }, + { name = "batched" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "torch" }, + { name = "tqdm" }, + { name = "transformers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/76/a19c864a1025222d3304a888ed4ed9217bfdf55dbaf4ed37500ee03935e0/mxbai_rerank-0.1.6.tar.gz", hash = "sha256:8d08e8464796429a7415314ce6de682bf9b538eb4ee5a7ddcd1a07839ee02879", size = 21449, upload-time = "2025-06-02T14:59:42.45Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/2a/503622b3a80272c662dabef421c9635168e5cbf6d51f0aa1883998561292/mxbai_rerank-0.1.6-py3-none-any.whl", hash = "sha256:aee94e7a14d5fba6520052ff2098f0f03db6cd9cc39553b7d2e82389deec9e05", size = 18458, upload-time = "2025-06-02T14:59:41.003Z" }, +] + +[[package]] +name = "networkx" +version = "3.4.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" }, +] + +[[package]] +name = "networkx" +version = "3.5" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/4f/ccdb8ad3a38e583f214547fd2f7ff1fc160c43a75af88e6aec213404b96a/networkx-3.5.tar.gz", hash = "sha256:d4c6f9cf81f52d69230866796b82afbccdec3db7ae4fbd1b65ea750feed50037", size = 2471065, upload-time = "2025-05-29T11:35:07.804Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" }, +] + [[package]] name = "nodeenv" version = "1.9.1" @@ -1753,6 +1827,139 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/de/bcad52ce972dc26232629ca3a99721fd4b22c1d2bda84d5db6541913ef9c/numpy-2.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:e017a8a251ff4d18d71f139e28bdc7c31edba7a507f72b1414ed902cbe48c74d", size = 12924237, upload-time = "2025-06-07T14:52:44.713Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.6.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/eb/ff4b8c503fa1f1796679dce648854d58751982426e4e4b37d6fce49d259c/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb", size = 393138322, upload-time = "2024-11-20T17:40:25.65Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.6.80" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/60/7b6497946d74bcf1de852a21824d63baad12cd417db4195fc1bfe59db953/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132", size = 8917980, upload-time = "2024-11-20T17:36:04.019Z" }, + { url = "https://files.pythonhosted.org/packages/a5/24/120ee57b218d9952c379d1e026c4479c9ece9997a4fb46303611ee48f038/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73", size = 8917972, upload-time = "2024-10-01T16:58:06.036Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/2e/46030320b5a80661e88039f59060d1790298b4718944a65a7f2aeda3d9e9/nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:35b0cc6ee3a9636d5409133e79273ce1f3fd087abb0532d2d2e8fff1fe9efc53", size = 23650380, upload-time = "2024-10-01T17:00:14.643Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/23/e717c5ac26d26cf39a27fbc076240fad2e3b817e5889d671b67f4f9f49c5/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7", size = 897690, upload-time = "2024-11-20T17:35:30.697Z" }, + { url = "https://files.pythonhosted.org/packages/f0/62/65c05e161eeddbafeca24dc461f47de550d9fa8a7e04eb213e32b55cfd99/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8", size = 897678, upload-time = "2024-10-01T16:57:33.821Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.5.1.17" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/78/4535c9c7f859a64781e43c969a3a7e84c54634e319a996d43ef32ce46f83/nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2", size = 570988386, upload-time = "2024-10-25T19:54:26.39Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/16/73727675941ab8e6ffd86ca3a4b7b47065edcca7a997920b831f8147c99d/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5", size = 200221632, upload-time = "2024-11-20T17:41:32.357Z" }, + { url = "https://files.pythonhosted.org/packages/60/de/99ec247a07ea40c969d904fc14f3a356b3e2a704121675b75c366b694ee1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca", size = 200221622, upload-time = "2024-10-01T17:03:58.79Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.11.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/66/cc9876340ac68ae71b15c743ddb13f8b30d5244af344ec8322b449e35426/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159", size = 1142103, upload-time = "2024-11-20T17:42:11.83Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.7.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/1b/44a01c4e70933637c93e6e1a8063d1e998b50213a6b65ac5a9169c47e98e/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf", size = 56279010, upload-time = "2024-11-20T17:42:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/4a/aa/2c7ff0b5ee02eaef890c0ce7d4f74bc30901871c5e45dee1ae6d0083cd80/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117", size = 56279000, upload-time = "2024-10-01T17:04:45.274Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/6e/c2cf12c9ff8b872e92b4a5740701e51ff17689c4d726fca91875b07f655d/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c", size = 158229790, upload-time = "2024-11-20T17:43:43.211Z" }, + { url = "https://files.pythonhosted.org/packages/9f/81/baba53585da791d043c10084cf9553e074548408e04ae884cfe9193bd484/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6", size = 158229780, upload-time = "2024-10-01T17:05:39.875Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/1e/b8b7c2f4099a37b96af5c9bb158632ea9e5d9d27d7391d7eb8fc45236674/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73", size = 216561367, upload-time = "2024-11-20T17:44:54.824Z" }, + { url = "https://files.pythonhosted.org/packages/43/ac/64c4316ba163e8217a99680c7605f779accffc6a4bcd0c778c12948d3707/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f", size = 216561357, upload-time = "2024-10-01T17:06:29.861Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.6.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/9a/72ef35b399b0e183bc2e8f6f558036922d453c4d8237dab26c666a04244b/nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46", size = 156785796, upload-time = "2024-10-15T21:29:17.709Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.26.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/ca/f42388aed0fddd64ade7493dbba36e1f534d4e6fdbdd355c6a90030ae028/nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:694cf3879a206553cc9d7dbda76b13efaf610fdb70a50cba303de1b0d1530ac6", size = 201319755, upload-time = "2025-03-13T00:29:55.296Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.6.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/d7/c5383e47c7e9bf1c99d5bd2a8c935af2b6d705ad831a7ec5c97db4d82f4f/nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:eedc36df9e88b682efe4309aa16b5b4e78c2407eac59e8c10a6a47535164369a", size = 19744971, upload-time = "2024-11-20T17:46:53.366Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/9a/fff8376f8e3d084cd1530e1ef7b879bb7d6d265620c95c1b322725c694f4/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2", size = 89276, upload-time = "2024-11-20T17:38:27.621Z" }, + { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265, upload-time = "2024-10-01T17:00:38.172Z" }, +] + [[package]] name = "ollama" version = "0.5.1" @@ -2146,6 +2353,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/af/ab3c51ab7507a7325e98ffe691d9495ee3d3aa5f589afad65ec920d39821/protobuf-6.31.1-py3-none-any.whl", hash = "sha256:720a6c7e6b77288b85063569baae8536671b39f15cc22037ec7045658d80489e", size = 168724, upload-time = "2025-05-28T19:25:53.926Z" }, ] +[[package]] +name = "psutil" +version = "7.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2a/80/336820c1ad9286a4ded7e845b2eccfcb27851ab8ac6abece774a6ff4d3de/psutil-7.0.0.tar.gz", hash = "sha256:7be9c3eba38beccb6495ea33afd982a44074b78f28c434a1f51cc07fd315c456", size = 497003, upload-time = "2025-02-13T21:54:07.946Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/e6/2d26234410f8b8abdbf891c9da62bee396583f713fb9f3325a4760875d22/psutil-7.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:101d71dc322e3cffd7cea0650b09b3d08b8e7c4109dd6809fe452dfd00e58b25", size = 238051, upload-time = "2025-02-13T21:54:12.36Z" }, + { url = "https://files.pythonhosted.org/packages/04/8b/30f930733afe425e3cbfc0e1468a30a18942350c1a8816acfade80c005c4/psutil-7.0.0-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:39db632f6bb862eeccf56660871433e111b6ea58f2caea825571951d4b6aa3da", size = 239535, upload-time = "2025-02-13T21:54:16.07Z" }, + { url = "https://files.pythonhosted.org/packages/2a/ed/d362e84620dd22876b55389248e522338ed1bf134a5edd3b8231d7207f6d/psutil-7.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1fcee592b4c6f146991ca55919ea3d1f8926497a713ed7faaf8225e174581e91", size = 275004, upload-time = "2025-02-13T21:54:18.662Z" }, + { url = "https://files.pythonhosted.org/packages/bf/b9/b0eb3f3cbcb734d930fdf839431606844a825b23eaf9a6ab371edac8162c/psutil-7.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b1388a4f6875d7e2aff5c4ca1cc16c545ed41dd8bb596cefea80111db353a34", size = 277986, upload-time = "2025-02-13T21:54:21.811Z" }, + { url = "https://files.pythonhosted.org/packages/eb/a2/709e0fe2f093556c17fbafda93ac032257242cabcc7ff3369e2cb76a97aa/psutil-7.0.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5f098451abc2828f7dc6b58d44b532b22f2088f4999a937557b603ce72b1993", size = 279544, upload-time = "2025-02-13T21:54:24.68Z" }, + { url = "https://files.pythonhosted.org/packages/50/e6/eecf58810b9d12e6427369784efe814a1eec0f492084ce8eb8f4d89d6d61/psutil-7.0.0-cp37-abi3-win32.whl", hash = "sha256:ba3fcef7523064a6c9da440fc4d6bd07da93ac726b5733c29027d7dc95b39d99", size = 241053, upload-time = "2025-02-13T21:54:34.31Z" }, + { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" }, +] + [[package]] name = "pyarrow" version = "20.0.0" @@ -2654,6 +2876,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/bf/b273dd11673fed8a6bd46032c0ea2a04b2ac9bfa9c628756a5856ba113b0/ruff-0.11.13-py3-none-win_arm64.whl", hash = "sha256:b4385285e9179d608ff1d2fb9922062663c658605819a6876d8beef0c30b7f3b", size = 10683928, upload-time = "2025-06-05T21:00:13.758Z" }, ] +[[package]] +name = "safetensors" +version = "0.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/71/7e/2d5d6ee7b40c0682315367ec7475693d110f512922d582fef1bd4a63adc3/safetensors-0.5.3.tar.gz", hash = "sha256:b6b0d6ecacec39a4fdd99cc19f4576f5219ce858e6fd8dbe7609df0b8dc56965", size = 67210, upload-time = "2025-02-26T09:15:13.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/ae/88f6c49dbd0cc4da0e08610019a3c78a7d390879a919411a410a1876d03a/safetensors-0.5.3-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:bd20eb133db8ed15b40110b7c00c6df51655a2998132193de2f75f72d99c7073", size = 436917, upload-time = "2025-02-26T09:15:03.702Z" }, + { url = "https://files.pythonhosted.org/packages/b8/3b/11f1b4a2f5d2ab7da34ecc062b0bc301f2be024d110a6466726bec8c055c/safetensors-0.5.3-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:21d01c14ff6c415c485616b8b0bf961c46b3b343ca59110d38d744e577f9cce7", size = 418419, upload-time = "2025-02-26T09:15:01.765Z" }, + { url = "https://files.pythonhosted.org/packages/5d/9a/add3e6fef267658075c5a41573c26d42d80c935cdc992384dfae435feaef/safetensors-0.5.3-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:11bce6164887cd491ca75c2326a113ba934be596e22b28b1742ce27b1d076467", size = 459493, upload-time = "2025-02-26T09:14:51.812Z" }, + { url = "https://files.pythonhosted.org/packages/df/5c/bf2cae92222513cc23b3ff85c4a1bb2811a2c3583ac0f8e8d502751de934/safetensors-0.5.3-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4a243be3590bc3301c821da7a18d87224ef35cbd3e5f5727e4e0728b8172411e", size = 472400, upload-time = "2025-02-26T09:14:53.549Z" }, + { url = "https://files.pythonhosted.org/packages/58/11/7456afb740bd45782d0f4c8e8e1bb9e572f1bf82899fb6ace58af47b4282/safetensors-0.5.3-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8bd84b12b1670a6f8e50f01e28156422a2bc07fb16fc4e98bded13039d688a0d", size = 522891, upload-time = "2025-02-26T09:14:55.717Z" }, + { url = "https://files.pythonhosted.org/packages/57/3d/fe73a9d2ace487e7285f6e157afee2383bd1ddb911b7cb44a55cf812eae3/safetensors-0.5.3-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:391ac8cab7c829452175f871fcaf414aa1e292b5448bd02620f675a7f3e7abb9", size = 537694, upload-time = "2025-02-26T09:14:57.036Z" }, + { url = "https://files.pythonhosted.org/packages/a6/f8/dae3421624fcc87a89d42e1898a798bc7ff72c61f38973a65d60df8f124c/safetensors-0.5.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cead1fa41fc54b1e61089fa57452e8834f798cb1dc7a09ba3524f1eb08e0317a", size = 471642, upload-time = "2025-02-26T09:15:00.544Z" }, + { url = "https://files.pythonhosted.org/packages/ce/20/1fbe16f9b815f6c5a672f5b760951e20e17e43f67f231428f871909a37f6/safetensors-0.5.3-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1077f3e94182d72618357b04b5ced540ceb71c8a813d3319f1aba448e68a770d", size = 502241, upload-time = "2025-02-26T09:14:58.303Z" }, + { url = "https://files.pythonhosted.org/packages/5f/18/8e108846b506487aa4629fe4116b27db65c3dde922de2c8e0cc1133f3f29/safetensors-0.5.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:799021e78287bac619c7b3f3606730a22da4cda27759ddf55d37c8db7511c74b", size = 638001, upload-time = "2025-02-26T09:15:05.79Z" }, + { url = "https://files.pythonhosted.org/packages/82/5a/c116111d8291af6c8c8a8b40628fe833b9db97d8141c2a82359d14d9e078/safetensors-0.5.3-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:df26da01aaac504334644e1b7642fa000bfec820e7cef83aeac4e355e03195ff", size = 734013, upload-time = "2025-02-26T09:15:07.892Z" }, + { url = "https://files.pythonhosted.org/packages/7d/ff/41fcc4d3b7de837963622e8610d998710705bbde9a8a17221d85e5d0baad/safetensors-0.5.3-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:32c3ef2d7af8b9f52ff685ed0bc43913cdcde135089ae322ee576de93eae5135", size = 670687, upload-time = "2025-02-26T09:15:09.979Z" }, + { url = "https://files.pythonhosted.org/packages/40/ad/2b113098e69c985a3d8fbda4b902778eae4a35b7d5188859b4a63d30c161/safetensors-0.5.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:37f1521be045e56fc2b54c606d4455573e717b2d887c579ee1dbba5f868ece04", size = 643147, upload-time = "2025-02-26T09:15:11.185Z" }, + { url = "https://files.pythonhosted.org/packages/0a/0c/95aeb51d4246bd9a3242d3d8349c1112b4ee7611a4b40f0c5c93b05f001d/safetensors-0.5.3-cp38-abi3-win32.whl", hash = "sha256:cfc0ec0846dcf6763b0ed3d1846ff36008c6e7290683b61616c4b040f6a54ace", size = 296677, upload-time = "2025-02-26T09:15:16.554Z" }, + { url = "https://files.pythonhosted.org/packages/69/e2/b011c38e5394c4c18fb5500778a55ec43ad6106126e74723ffaee246f56e/safetensors-0.5.3-cp38-abi3-win_amd64.whl", hash = "sha256:836cbbc320b47e80acd40e44c8682db0e8ad7123209f69b093def21ec7cafd11", size = 308878, upload-time = "2025-02-26T09:15:14.99Z" }, +] + +[[package]] +name = "setuptools" +version = "80.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/18/5d/3bf57dcd21979b887f014ea83c24ae194cfcd12b9e0fda66b957c69d1fca/setuptools-80.9.0.tar.gz", hash = "sha256:f36b47402ecde768dbfafc46e8e4207b4360c654f1f3bb84475f0a28628fb19c", size = 1319958, upload-time = "2025-05-27T00:56:51.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/dc/17031897dae0efacfea57dfd3a82fdd2a2aeb58e0ff71b77b87e44edc772/setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922", size = 1201486, upload-time = "2025-05-27T00:56:49.664Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" @@ -2883,6 +3136,58 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6e/c2/61d3e0f47e2b74ef40a68b9e6ad5984f6241a942f7cd3bbfbdbd03861ea9/tomli-2.2.1-py3-none-any.whl", hash = "sha256:cb55c73c5f4408779d0cf3eef9f762b9c9f147a77de7b258bef0a5628adc85cc", size = 14257, upload-time = "2024-11-27T22:38:35.385Z" }, ] +[[package]] +name = "torch" +version = "2.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.12'" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/27/2e06cb52adf89fe6e020963529d17ed51532fc73c1e6d1b18420ef03338c/torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a103b5d782af5bd119b81dbcc7ffc6fa09904c423ff8db397a1e6ea8fd71508f", size = 99089441, upload-time = "2025-06-04T17:38:48.268Z" }, + { url = "https://files.pythonhosted.org/packages/0a/7c/0a5b3aee977596459ec45be2220370fde8e017f651fecc40522fd478cb1e/torch-2.7.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:fe955951bdf32d182ee8ead6c3186ad54781492bf03d547d31771a01b3d6fb7d", size = 821154516, upload-time = "2025-06-04T17:36:28.556Z" }, + { url = "https://files.pythonhosted.org/packages/f9/91/3d709cfc5e15995fb3fe7a6b564ce42280d3a55676dad672205e94f34ac9/torch-2.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:885453d6fba67d9991132143bf7fa06b79b24352f4506fd4d10b309f53454162", size = 216093147, upload-time = "2025-06-04T17:39:38.132Z" }, + { url = "https://files.pythonhosted.org/packages/92/f6/5da3918414e07da9866ecb9330fe6ffdebe15cb9a4c5ada7d4b6e0a6654d/torch-2.7.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d72acfdb86cee2a32c0ce0101606f3758f0d8bb5f8f31e7920dc2809e963aa7c", size = 68630914, upload-time = "2025-06-04T17:39:31.162Z" }, + { url = "https://files.pythonhosted.org/packages/11/56/2eae3494e3d375533034a8e8cf0ba163363e996d85f0629441fa9d9843fe/torch-2.7.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:236f501f2e383f1cb861337bdf057712182f910f10aeaf509065d54d339e49b2", size = 99093039, upload-time = "2025-06-04T17:39:06.963Z" }, + { url = "https://files.pythonhosted.org/packages/e5/94/34b80bd172d0072c9979708ccd279c2da2f55c3ef318eceec276ab9544a4/torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:06eea61f859436622e78dd0cdd51dbc8f8c6d76917a9cf0555a333f9eac31ec1", size = 821174704, upload-time = "2025-06-04T17:37:03.799Z" }, + { url = "https://files.pythonhosted.org/packages/50/9e/acf04ff375b0b49a45511c55d188bcea5c942da2aaf293096676110086d1/torch-2.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:8273145a2e0a3c6f9fd2ac36762d6ee89c26d430e612b95a99885df083b04e52", size = 216095937, upload-time = "2025-06-04T17:39:24.83Z" }, + { url = "https://files.pythonhosted.org/packages/5b/2b/d36d57c66ff031f93b4fa432e86802f84991477e522adcdffd314454326b/torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730", size = 68640034, upload-time = "2025-06-04T17:39:17.989Z" }, + { url = "https://files.pythonhosted.org/packages/87/93/fb505a5022a2e908d81fe9a5e0aa84c86c0d5f408173be71c6018836f34e/torch-2.7.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ea1e518df4c9de73af7e8a720770f3628e7f667280bce2be7a16292697e3fa", size = 98948276, upload-time = "2025-06-04T17:39:12.852Z" }, + { url = "https://files.pythonhosted.org/packages/56/7e/67c3fe2b8c33f40af06326a3d6ae7776b3e3a01daa8f71d125d78594d874/torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c33360cfc2edd976c2633b3b66c769bdcbbf0e0b6550606d188431c81e7dd1fc", size = 821025792, upload-time = "2025-06-04T17:34:58.747Z" }, + { url = "https://files.pythonhosted.org/packages/a1/37/a37495502bc7a23bf34f89584fa5a78e25bae7b8da513bc1b8f97afb7009/torch-2.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:d8bf6e1856ddd1807e79dc57e54d3335f2b62e6f316ed13ed3ecfe1fc1df3d8b", size = 216050349, upload-time = "2025-06-04T17:38:59.709Z" }, + { url = "https://files.pythonhosted.org/packages/3a/60/04b77281c730bb13460628e518c52721257814ac6c298acd25757f6a175c/torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb", size = 68645146, upload-time = "2025-06-04T17:38:52.97Z" }, + { url = "https://files.pythonhosted.org/packages/66/81/e48c9edb655ee8eb8c2a6026abdb6f8d2146abd1f150979ede807bb75dcb/torch-2.7.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:03563603d931e70722dce0e11999d53aa80a375a3d78e6b39b9f6805ea0a8d28", size = 98946649, upload-time = "2025-06-04T17:38:43.031Z" }, + { url = "https://files.pythonhosted.org/packages/3a/24/efe2f520d75274fc06b695c616415a1e8a1021d87a13c68ff9dce733d088/torch-2.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d632f5417b6980f61404a125b999ca6ebd0b8b4bbdbb5fbbba44374ab619a412", size = 821033192, upload-time = "2025-06-04T17:38:09.146Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d9/9c24d230333ff4e9b6807274f6f8d52a864210b52ec794c5def7925f4495/torch-2.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:23660443e13995ee93e3d844786701ea4ca69f337027b05182f5ba053ce43b38", size = 216055668, upload-time = "2025-06-04T17:38:36.253Z" }, + { url = "https://files.pythonhosted.org/packages/95/bf/e086ee36ddcef9299f6e708d3b6c8487c1651787bb9ee2939eb2a7f74911/torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585", size = 68925988, upload-time = "2025-06-04T17:38:29.273Z" }, + { url = "https://files.pythonhosted.org/packages/69/6a/67090dcfe1cf9048448b31555af6efb149f7afa0a310a366adbdada32105/torch-2.7.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e08d7e6f21a617fe38eeb46dd2213ded43f27c072e9165dc27300c9ef9570934", size = 99028857, upload-time = "2025-06-04T17:37:50.956Z" }, + { url = "https://files.pythonhosted.org/packages/90/1c/48b988870823d1cc381f15ec4e70ed3d65e043f43f919329b0045ae83529/torch-2.7.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:30207f672328a42df4f2174b8f426f354b2baa0b7cca3a0adb3d6ab5daf00dc8", size = 821098066, upload-time = "2025-06-04T17:37:33.939Z" }, + { url = "https://files.pythonhosted.org/packages/7b/eb/10050d61c9d5140c5dc04a89ed3257ef1a6b93e49dd91b95363d757071e0/torch-2.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:79042feca1c634aaf6603fe6feea8c6b30dfa140a6bbc0b973e2260c7e79a22e", size = 216336310, upload-time = "2025-06-04T17:36:09.862Z" }, + { url = "https://files.pythonhosted.org/packages/b1/29/beb45cdf5c4fc3ebe282bf5eafc8dfd925ead7299b3c97491900fe5ed844/torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946", size = 68645708, upload-time = "2025-06-04T17:34:39.852Z" }, +] + [[package]] name = "tqdm" version = "4.67.1" @@ -2895,6 +3200,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, ] +[[package]] +name = "transformers" +version = "4.53.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "requests" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/67/80f51466ec447028fd84469b208eb742533ce06cc8fad2e3181380199e5c/transformers-4.53.2.tar.gz", hash = "sha256:6c3ed95edfb1cba71c4245758f1b4878c93bf8cde77d076307dacb2cbbd72be2", size = 9201233, upload-time = "2025-07-11T12:39:08.742Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/88/beb33a79a382fcd2aed0be5222bdc47f41e4bfe7aaa90ae1374f1d8ea2af/transformers-4.53.2-py3-none-any.whl", hash = "sha256:db8f4819bb34f000029c73c3c557e7d06fc1b8e612ec142eecdae3947a9c78bf", size = 10826609, upload-time = "2025-07-11T12:39:05.461Z" }, +] + +[[package]] +name = "triton" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "setuptools" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/a9/549e51e9b1b2c9b854fd761a1d23df0ba2fbc60bd0c13b489ffa518cfcb7/triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e", size = 155600257, upload-time = "2025-05-29T23:39:36.085Z" }, + { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937, upload-time = "2025-05-29T23:39:44.182Z" }, + { url = "https://files.pythonhosted.org/packages/24/5f/950fb373bf9c01ad4eb5a8cd5eaf32cdf9e238c02f9293557a2129b9c4ac/triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43", size = 155669138, upload-time = "2025-05-29T23:39:51.771Z" }, + { url = "https://files.pythonhosted.org/packages/74/1f/dfb531f90a2d367d914adfee771babbd3f1a5b26c3f5fbc458dee21daa78/triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240", size = 155673035, upload-time = "2025-05-29T23:40:02.468Z" }, + { url = "https://files.pythonhosted.org/packages/28/71/bd20ffcb7a64c753dc2463489a61bf69d531f308e390ad06390268c4ea04/triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42", size = 155735832, upload-time = "2025-05-29T23:40:10.522Z" }, +] + [[package]] name = "typer" version = "0.16.0" From f25416b0ffdb04cef5fbace73646989182cae4e7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 18 Jul 2025 20:55:28 +0300 Subject: [PATCH 04/10] Properly test cohere reranker --- src/haiku/rag/config.py | 2 - tests/test_reranker.py | 95 +++++++++-------------------------------- 2 files changed, 21 insertions(+), 76 deletions(-) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 6151a80c..8fb2bef1 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,8 +19,6 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 - # RERANK_PROVIDER: str = "cohere" - # RERANK_MODEL: str = "rerank-v3.5" RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" RERANK_PROVIDER: str = "mxbai" diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 39a515c1..c83c2d3d 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -1,10 +1,23 @@ import pytest -from haiku.rag.reranking import get_reranker from haiku.rag.reranking.base import RerankerBase from haiku.rag.reranking.mxbai import MxBAIReranker from haiku.rag.store.models.chunk import Chunk +chunks = [ + Chunk(content=content, document_id=i) + for i, content in enumerate( + [ + "To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.", + "The novel Moby-Dick was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.", + "Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961.", + "Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment upon the British landed gentry at the end of the 18th century.", + "The Harry Potter series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era.", + "The Great Gatsby, a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.", + ] + ) +] + @pytest.mark.asyncio async def test_reranker_base(): @@ -18,20 +31,6 @@ async def test_reranker_base(): @pytest.mark.asyncio async def test_mxbai_reranker(): reranker = MxBAIReranker() - chunks = [ - Chunk(content=content, document_id=i) - for i, content in enumerate( - [ - "To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.", - "The novel Moby-Dick was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.", - "Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961.", - "Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment upon the British landed gentry at the end of the 18th century.", - "The Harry Potter series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era.", - "The Great Gatsby, a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.", - ] - ) - ] - reranked = await reranker.rerank( "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 ) @@ -41,67 +40,15 @@ async def test_mxbai_reranker(): @pytest.mark.asyncio async def test_cohere_reranker(): try: - # Mock the client - class MockResult: - def __init__(self, index): - self.index = index + from haiku.rag.reranking.cohere import CohereReranker - class MockResponse: - def __init__(self, results): - self.results = results + reranker = CohereReranker() + assert reranker._model == "rerank-v3.5" - class MockClient: - def __init__(self, api_key=None): - pass - - def rerank(self, model, query, documents, top_n): - return MockResponse([MockResult(1), MockResult(0)]) - - import haiku.rag.reranking.cohere - - original_client = haiku.rag.reranking.cohere.cohere.ClientV2 - haiku.rag.reranking.cohere.cohere.ClientV2 = MockClient - - try: - from haiku.rag.reranking.cohere import CohereReranker - - reranker = CohereReranker() - assert reranker._model == "rerank-v3.5" - - chunks = [ - Chunk(id=1, content="First chunk", document_id=1), - Chunk(id=2, content="Second chunk", document_id=1), - ] - - result = await reranker.rerank("test query", chunks) - assert len(result) == 2 - assert result[0] == chunks[1] # Should return chunk at index 1 first - assert result[1] == chunks[0] # Should return chunk at index 0 second - finally: - haiku.rag.reranking.cohere.cohere.ClientV2 = original_client + reranked = await reranker.rerank( + "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 + ) + assert [r.document_id for r in reranked] == [0, 2] except ImportError: pytest.skip("Cohere package not installed") - - -@pytest.mark.asyncio -async def test_get_reranker(): - try: - - class MockClient: - def __init__(self, api_key=None): - pass - - import haiku.rag.reranking.cohere - - original_client = haiku.rag.reranking.cohere.cohere.ClientV2 - haiku.rag.reranking.cohere.cohere.ClientV2 = MockClient - - try: - reranker = get_reranker() - assert reranker._model == "rerank-v3.5" - assert hasattr(reranker, "rerank") - finally: - haiku.rag.reranking.cohere.cohere.ClientV2 = original_client - except ImportError: - pytest.skip("Cohere package not installed") From 2343ab7751e9898d599c404fc8c56ff415c73e6c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 12:16:17 +0300 Subject: [PATCH 05/10] Return tuple (Chunk, score,) in rerank() --- src/haiku/rag/reranking/base.py | 2 +- src/haiku/rag/reranking/cohere.py | 4 ++-- src/haiku/rag/reranking/mxbai.py | 4 ++-- tests/test_reranker.py | 8 +++++--- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/haiku/rag/reranking/base.py b/src/haiku/rag/reranking/base.py index 72b3a3df..0e95e26b 100644 --- a/src/haiku/rag/reranking/base.py +++ b/src/haiku/rag/reranking/base.py @@ -7,7 +7,7 @@ class RerankerBase: async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 - ) -> list[Chunk]: + ) -> list[tuple[Chunk, float]]: raise NotImplementedError( "Reranker is an abstract class. Please implement the rerank method in a subclass." ) diff --git a/src/haiku/rag/reranking/cohere.py b/src/haiku/rag/reranking/cohere.py index 19e91577..6d30952d 100644 --- a/src/haiku/rag/reranking/cohere.py +++ b/src/haiku/rag/reranking/cohere.py @@ -16,7 +16,7 @@ class CohereReranker(RerankerBase): async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 - ) -> list[Chunk]: + ) -> list[tuple[Chunk, float]]: if not chunks: return [] @@ -29,6 +29,6 @@ class CohereReranker(RerankerBase): reranked_chunks = [] for result in response.results: original_chunk = chunks[result.index] - reranked_chunks.append(original_chunk) + reranked_chunks.append((original_chunk, result.relevance_score)) return reranked_chunks diff --git a/src/haiku/rag/reranking/mxbai.py b/src/haiku/rag/reranking/mxbai.py index 135727ed..032edac5 100644 --- a/src/haiku/rag/reranking/mxbai.py +++ b/src/haiku/rag/reranking/mxbai.py @@ -13,7 +13,7 @@ class MxBAIReranker(RerankerBase): async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 - ) -> list[Chunk]: + ) -> list[tuple[Chunk, float]]: if not chunks: return [] @@ -23,6 +23,6 @@ class MxBAIReranker(RerankerBase): reranked_chunks = [] for result in results: original_chunk = chunks[result.index] - reranked_chunks.append(original_chunk) + reranked_chunks.append((original_chunk, result.score)) return reranked_chunks diff --git a/tests/test_reranker.py b/tests/test_reranker.py index c83c2d3d..b3de96cf 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -34,7 +34,8 @@ async def test_mxbai_reranker(): reranked = await reranker.rerank( "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 ) - assert [r.document_id for r in reranked] == [0, 2] + assert [chunk.document_id for chunk, score in reranked] == [0, 2] + assert all(isinstance(score, float) for chunk, score in reranked) @pytest.mark.asyncio @@ -43,12 +44,13 @@ async def test_cohere_reranker(): from haiku.rag.reranking.cohere import CohereReranker reranker = CohereReranker() - assert reranker._model == "rerank-v3.5" + reranker._model = "rerank-v3.5" reranked = await reranker.rerank( "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 ) - assert [r.document_id for r in reranked] == [0, 2] + assert [chunk.document_id for chunk, score in reranked] == [0, 2] + assert all(isinstance(score, float) for chunk, score in reranked) except ImportError: pytest.skip("Cohere package not installed") From 1437942edcccc15709d6f7355c2ebb781ba3249a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 12:51:22 +0300 Subject: [PATCH 06/10] Make reranker a global module object to avoid re-initialization --- src/haiku/rag/config.py | 2 +- src/haiku/rag/reranking/__init__.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 8fb2bef1..d4fd7421 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,8 +19,8 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 - RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" RERANK_PROVIDER: str = "mxbai" + RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" QA_PROVIDER: str = "ollama" QA_MODEL: str = "qwen3" diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index 8fbdae31..ccef8b9f 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -6,11 +6,21 @@ try: except ImportError: pass +_reranker: RerankerBase | None = None + def get_reranker() -> RerankerBase: """ Factory function to get the appropriate reranker based on the configuration. """ + global _reranker + if _reranker is not None: + return _reranker + if Config.RERANK_PROVIDER == "mxbai": + from haiku.rag.reranking.mxbai import MxBAIReranker + + _reranker = MxBAIReranker() + return _reranker if Config.RERANK_PROVIDER == "cohere": try: @@ -21,6 +31,7 @@ def get_reranker() -> RerankerBase: "Please install haiku.rag with the 'cohere' extra:" "uv pip install haiku.rag --extra cohere" ) - return CohereReranker() + _reranker = CohereReranker() + return _reranker raise ValueError(f"Unsupported reranker provider: {Config.RERANK_PROVIDER}") From 2347612fb6af34b8938ed2af6a33cad5780dc698 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 13:12:34 +0300 Subject: [PATCH 07/10] Optionally perform reranking when searching --- src/haiku/rag/client.py | 22 +++++++++++++++++++--- src/haiku/rag/config.py | 1 + 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 7f375155..11a60089 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -10,6 +10,7 @@ import httpx from haiku.rag.config import Config from haiku.rag.reader import FileReader +from haiku.rag.reranking import get_reranker from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -277,9 +278,9 @@ class HaikuRAG: return await self.document_repository.list_all(limit=limit, offset=offset) async def search( - self, query: str, limit: int = 5, k: int = 60 + self, query: str, limit: int = 3, k: int = 60, rerank=Config.RERANK ) -> list[tuple[Chunk, float]]: - """Search for relevant chunks using hybrid search (vector similarity + full-text search). + """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking. Args: query: The search query string. @@ -289,7 +290,22 @@ class HaikuRAG: Returns: List of (chunk, score) tuples ordered by relevance. """ - return await self.chunk_repository.search_chunks_hybrid(query, limit, k) + + if not rerank: + return await self.chunk_repository.search_chunks_hybrid(query, limit, k) + + # Get more initial results (3X) for reranking + search_results = await self.chunk_repository.search_chunks_hybrid( + query, limit * 3, k + ) + + # Apply reranking + reranker = get_reranker() + chunks = [chunk for chunk, _ in search_results] + reranked_results = await reranker.rerank(query, chunks, top_n=limit) + + # Return reranked results with scores from reranker + return reranked_results async def ask(self, question: str) -> str: """Ask a question using the configured QA agent. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index d4fd7421..f174a81c 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,6 +19,7 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 + RERANK: bool = False RERANK_PROVIDER: str = "mxbai" RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" From a6595b15974f59e1cd63a2e472bd1e3c0c61e3ec Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 13:16:09 +0300 Subject: [PATCH 08/10] Update docs --- README.md | 5 +++-- docs/configuration.md | 29 +++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 968c6fa1..2d70755b 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite. - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI - **Multiple QA providers**: Ollama, OpenAI, Anthropic - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion +- **Reranking**: Optional result reranking with MixedBread AI or Cohere - **Question answering**: Built-in QA agents on your documents - **File monitoring**: Auto-index files when run as server - **40+ file formats**: PDF, DOCX, HTML, Markdown, audio, URLs @@ -49,8 +50,8 @@ async with HaikuRAG("database.db") as client: # Add document doc = await client.create_document("Your content") - # Search - results = await client.search("query") + # Search (with optional reranking) + results = await client.search("query", rerank=True) for chunk, score in results: print(f"{score:.3f}: {chunk.content}") diff --git a/docs/configuration.md b/docs/configuration.md index 5dbba71b..ef5fb00a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -103,6 +103,35 @@ QA_MODEL="claude-3-5-haiku-20241022" # or claude-3-5-sonnet-20241022, etc. ANTHROPIC_API_KEY="your-api-key" ``` +## Reranking + +Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. + +### MixedBread AI (Default) + +```bash +RERANK=true +RERANK_PROVIDER="mxbai" +RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2" +``` + +### Cohere + +For Cohere reranking, install with Cohere extras: + +```bash +uv pip install haiku.rag --extra cohere +``` + +Then configure: + +```bash +RERANK=true +RERANK_PROVIDER="cohere" +RERANK_MODEL="rerank-v3.5" +COHERE_API_KEY="your-api-key" +``` + ## Other Settings ### Database and Storage From c1d53348041f22c7c6ad6264e0cb89e4856dd2c7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 19:02:47 +0300 Subject: [PATCH 09/10] Update benchmarks with reranking --- docs/benchmarks.md | 20 +++++++++++--------- src/haiku/rag/qa/prompts.py | 3 ++- tests/llm_judge.py | 34 +++++++++++++++++++++++----------- 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index d5a85a04..6b043fcf 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -12,17 +12,19 @@ In order to calculate recall, we load the `News Stories` from `repliqa_3` which The recall obtained is ~0.73 for matching in the top result, raising to ~0.75 for the top 3 results. -| Model | Document in top 1 | Document in top 3 | -|---------------------------------------|-------------------|-------------------| -| Ollama / `mxbai-embed-large` | 0.77 | 0.89 | -| Ollama / `nomic-embed-text` | 0.74 | 0.88 | -| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | +| Model | Document in top 1 | Document in top 3 | Reranker | +|---------------------------------------|-------------------|-------------------|----------------------| +| Ollama / `mxbai-embed-large` | 0.77 | 0.89 | None | +| Ollama / `mxbai-embed-large` | 0.81 | 0.91 | mxbai-rerank-base-v2 | +| Ollama / `nomic-embed-text` | 0.74 | 0.88 | None | +| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None | ## Question/Answer evaluation Again using the same dataset, we use a QA agent to answer the question. In addition we use an LLM judge (using the Ollama `qwen3`) to evaluate whether the answer is correct or not. The obtained accuracy is as follows: -| Embedding Model | QA Model | Accuracy | -|------------------------------|-----------------------------------|-----------| -| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.64 | -| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | +| Embedding Model | QA Model | Accuracy | Reranker | +|------------------------------|-----------------------------------|-----------|----------------------| +| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.64 | None | +| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.72 | mxbai-rerank-base-v2 | +| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None | diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index 283c40e2..68d42cc2 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -6,7 +6,7 @@ Your process: 2. Search with specific keywords and phrases from the user's question 3. Review the search results and their relevance scores 4. If you need additional context, perform follow-up searches with different keywords -5. Provide a comprehensive answer based only on the retrieved documents +5. Provide a short and to the point comprehensive answer based only on the retrieved documents Guidelines: - Base your answers strictly on the provided document content @@ -15,6 +15,7 @@ Guidelines: - Indicate when information is incomplete or when you need to search for additional context - If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question." - For complex questions, consider breaking them down and performing multiple searches +- Stick to the answer, do not ellaborate or provde context unless asked for it. Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents. """ diff --git a/tests/llm_judge.py b/tests/llm_judge.py index 5af4cf0e..f7e3f6e2 100644 --- a/tests/llm_judge.py +++ b/tests/llm_judge.py @@ -35,23 +35,35 @@ class LLMJudge: - score: str rating from 1-5 """ - prompt = f""" - You are an expert judge evaluating the equivalence of two answers to the same question. + prompt = f"""You are an expert evaluator determining whether two answers to the same question are semantically equivalent. - Question: {question} +QUESTION: {question} - Generated Answer: {answer} +GENERATED ANSWER: {answer} - Expected Answer: {expected_answer} +EXPECTED ANSWER: {expected_answer} - Your task is to determine if these two answers are equivalent in meaning and both correctly answer the question. Consider: +EVALUATION CRITERIA: +Rate as EQUIVALENT (true) if: +✓ Both answers contain the same core factual information +✓ Both directly address the question asked +✓ The key claims and conclusions are consistent +✓ Any additional detail in one answer doesn't contradict the other - 1. Do both answers provide the same answer? - 2. Do both answers directly address the question asked? - 3. Minor differences in wording or style are acceptable if the meaning of the answer is the same. - 4. If one answer is more detailed but the other is correct, they can still be considered equivalent. +Rate as NOT EQUIVALENT (false) if: +✗ Factual contradictions exist between the answers +✗ One answer fails to address the core question +✗ Key information is missing from one answer that changes the meaning +✗ The answers lead to different conclusions or implications - Be strict but fair in your evaluation. Focus on factual correctness and whether both answers would satisfy someone asking the question.""" +GUIDELINES: +- Ignore minor differences in phrasing, style, or formatting +- Focus on semantic meaning rather than exact wording +- Consider both answers correct if they convey the same essential information +- Be tolerant of different levels of detail if the core answer is preserved +- Evaluate based on what a person asking this question would need to know + +Respond with JSON containing only: {{"equivalent": true}} or {{"equivalent": false}}""" response = await self.client.chat( model=self.model, From 942be2271cca9671fdbacf205c5df04c17b4e845 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 19 Jul 2025 19:57:37 +0300 Subject: [PATCH 10/10] Update docs --- README.md | 6 +++--- docs/configuration.md | 10 +++++++--- docs/index.md | 6 +++--- docs/python.md | 7 +++++-- src/haiku/rag/config.py | 2 +- 5 files changed, 19 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 2d70755b..2220b6d7 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite. - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI - **Multiple QA providers**: Ollama, OpenAI, Anthropic - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion -- **Reranking**: Optional result reranking with MixedBread AI or Cohere +- **Reranking**: Default search result reranking with MixedBread AI or Cohere - **Question answering**: Built-in QA agents on your documents - **File monitoring**: Auto-index files when run as server - **40+ file formats**: PDF, DOCX, HTML, Markdown, audio, URLs @@ -50,8 +50,8 @@ async with HaikuRAG("database.db") as client: # Add document doc = await client.create_document("Your content") - # Search (with optional reranking) - results = await client.search("query", rerank=True) + # Search (reranking enabled by default) + results = await client.search("query") for chunk, score in results: print(f"{score:.3f}: {chunk.content}") diff --git a/docs/configuration.md b/docs/configuration.md index ef5fb00a..4dea64e6 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -105,12 +105,17 @@ ANTHROPIC_API_KEY="your-api-key" ## Reranking -Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. +Reranking is **enabled by default** and improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. + +If you use the default reranked (running locally), it can slow down searching significantly. To disable reranking for faster searches: + +```bash +RERANK=false +``` ### MixedBread AI (Default) ```bash -RERANK=true RERANK_PROVIDER="mxbai" RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2" ``` @@ -126,7 +131,6 @@ uv pip install haiku.rag --extra cohere Then configure: ```bash -RERANK=true RERANK_PROVIDER="cohere" RERANK_MODEL="rerank-v3.5" COHERE_API_KEY="your-api-key" diff --git a/docs/index.md b/docs/index.md index 19da8d6d..aa969d6a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,13 +1,13 @@ # haiku.rag -`haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work on SQLite alone without the need for external vector databases. It uses [sqlite-vec](https://github.com/asg017/sqlite-vec) for storing the embeddings and performs semantic (vector) search as well as full-text search combined through Reciprocal Rank Fusion. Both open-source (Ollama) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. - +`haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work on SQLite alone without the need for external vector databases. It uses [sqlite-vec](https://github.com/asg017/sqlite-vec) for storing the embeddings and performs semantic (vector) search as well as full-text search combined through Reciprocal Rank Fusion. Both open-source (Ollama, MixedBread AI) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. ## Features - **Local SQLite**: No need to run additional servers - **Support for various embedding providers**: Ollama, VoyageAI, OpenAI or add your own - **Hybrid Search**: Vector search using `sqlite-vec` combined with full-text search `FTS5`, using Reciprocal Rank Fusion +- **Reranking**: Optional result reranking with MixedBread AI or Cohere - **Question Answering**: Built-in QA agents using Ollama, OpenAI, or Anthropic. - **File monitoring**: Automatically index files when run as a server - **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more. Or add a URL! @@ -34,7 +34,7 @@ async with HaikuRAG("database.db") as client: results = await client.search("query") # Ask questions - answer = await client.ask("Who is the author of haiku.rag?") + answer = await client.ask("Who is the author of haiku.rag?", rerank=False) ``` Or use the CLI: diff --git a/docs/python.md b/docs/python.md index 8ad47f4e..dc8121b5 100644 --- a/docs/python.md +++ b/docs/python.md @@ -76,7 +76,9 @@ async for doc_id in client.rebuild_database(): ## Searching Documents -Basic search: +The search method performs hybrid search (vector + full-text) with **reranking enabled by default** for improved relevance: + +Basic search (with reranking): ```python results = await client.search("machine learning algorithms", limit=5) for chunk, score in results: @@ -90,7 +92,8 @@ With options: results = await client.search( query="machine learning", limit=5, # Maximum results to return - k=60 # RRF parameter for reciprocal rank fusion + k=60, # RRF parameter for reciprocal rank fusion + rerank=False # Disable reranking for faster search ) # Process results diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index f174a81c..898f3856 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,7 +19,7 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 - RERANK: bool = False + RERANK: bool = True RERANK_PROVIDER: str = "mxbai" RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2"