e8f00fcff4
haiku.rag.config exported two configuration instances: the lazy _config behind get_config/set_config, and Config, loaded at import time. Nothing linked them, and eleven signatures captured Config as a default argument, so set_config could not reach the factories, the client, the store or the MCP server. reranking/base.py went further and snapshotted the configured reranker name into a class attribute at import. Config is removed. Internal defaults are config: AppConfig | None = None, resolved through get_config() per call. RerankerBase._model is None and CohereReranker takes its model name as an argument, like every other reranker. The suite patched attributes on Config while production read the instance get_config() returns, a different object, so those patches were no-ops waiting to happen. They now go through get_config().
33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
try:
|
|
import cohere
|
|
except ImportError as e: # pragma: no cover
|
|
raise ImportError(
|
|
"cohere is not installed. Please install it with `pip install cohere` or use the cohere optional dependency."
|
|
) from e
|
|
|
|
|
|
class CohereReranker(RerankerBase): # pragma: no cover
|
|
def __init__(self, model: str | None = None):
|
|
self._model = model
|
|
# Cohere SDK reads CO_API_KEY from environment by default
|
|
self._client = cohere.AsyncClientV2()
|
|
|
|
async def _rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
documents = [chunk.content for chunk in chunks]
|
|
|
|
model_name = self._model or "rerank-v3.5"
|
|
response = await self._client.rerank(
|
|
model=model_name, 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, result.relevance_score))
|
|
|
|
return reranked_chunks
|