Cohere reranker

This commit is contained in:
Yiorgis Gozadinos 2025-07-17 13:17:16 +03:00
parent 20beb23a38
commit 859860318a
No known key found for this signature in database
7 changed files with 237 additions and 1 deletions

View file

@ -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"

View file

@ -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

View file

@ -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}")

View file

@ -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."
)

View file

@ -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

83
tests/test_reranker.py Normal file
View file

@ -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")

75
uv.lock
View file

@ -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"