Reranking support for vLLM, documentation update

README.md

@@ -9,10 +9,10 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 ## Features
 
 - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
-- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI
+- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
 - **Multiple QA providers**: Any provider/model supported by Pydantic AI
 - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
-- **Reranking**: Default search result reranking with MixedBread AI or Cohere
+- **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
 - **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, code files, URLs

docs/configuration.md

@@ -53,6 +53,18 @@ EMBEDDINGS_VECTOR_DIM=1536
 OPENAI_API_KEY="your-api-key"
 ```
 
+### vLLM
+For high-performance local inference, you can use vLLM to serve embedding models with OpenAI-compatible APIs:
+
+```bash
+EMBEDDINGS_PROVIDER="vllm"
+EMBEDDINGS_MODEL="mixedbread-ai/mxbai-embed-large-v1"  # Any embedding model supported by vLLM
+EMBEDDINGS_VECTOR_DIM=512  # Dimension depends on the model
+VLLM_EMBEDDINGS_BASE_URL="http://localhost:8000"  # vLLM server URL
+```
+
+**Note:** You need to run a vLLM server separately with an embedding model loaded.
+
 ## Question Answering Providers
 
 Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used.
@@ -136,6 +148,18 @@ RERANK_MODEL="rerank-v3.5"
 COHERE_API_KEY="your-api-key"
 ```
 
+### vLLM
+
+For high-performance local reranking using dedicated reranking models:
+
+```bash
+RERANK_PROVIDER="vllm"
+RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2"  # Any reranking model supported by vLLM
+VLLM_RERANK_BASE_URL="http://localhost:8001"  # vLLM server URL
+```
+
+**Note:** vLLM reranking uses the `/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded. Consult the specific model's documentation for proper vLLM serving configuration.
+
 ## Other Settings
 
 ### Database and Storage

docs/installation.md

@@ -11,6 +11,7 @@ This includes support for:
 - **OpenAI** (GPT models for QA and embeddings)
 - **Anthropic** (Claude models for QA)
 - **Cohere** (reranking models)
+- **vLLM** (high-performance local inference for embeddings, QA, and reranking)
 
 ## Provider-Specific Installation
 
@@ -28,7 +29,38 @@ uv pip install haiku.rag[voyageai]
 uv pip install haiku.rag[mxbai]
 ```
 
+### vLLM Setup
+
+vLLM requires no additional installation - it works with the base haiku.rag package. However, you need to run vLLM servers separately:
+
+```bash
+# Install vLLM
+pip install vllm
+
+# Serve an embedding model
+vllm serve mixedbread-ai/mxbai-embed-large-v1 --port 8000
+
+# Serve a model for reranking
+vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --port 8001
+```
+
+Then configure haiku.rag to use the vLLM servers:
+
+```bash
+# Embeddings
+EMBEDDINGS_PROVIDER="vllm"
+EMBEDDINGS_MODEL="mixedbread-ai/mxbai-embed-large-v1"
+EMBEDDINGS_VECTOR_DIM=512
+VLLM_EMBEDDINGS_BASE_URL="http://localhost:8000"
+
+# Reranking (optional)
+RERANK_PROVIDER="vllm"
+RERANK_MODEL="microsoft/DialoGPT-medium"
+VLLM_RERANK_BASE_URL="http://localhost:8001"
+```
+
 ## Requirements
 
 - Python 3.10+
 - Ollama (for default embeddings)
+- vLLM server (for vLLM provider)

src/haiku/rag/reranking/vllm.py

@@ -0,0 +1,44 @@
+import httpx
+
+from haiku.rag.config import Config
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class VLLMReranker(RerankerBase):
+    def __init__(self, model: str):
+        self._model = model
+        self._base_url = Config.VLLM_RERANK_BASE_URL
+
+    async def rerank(
+        self, query: str, chunks: list[Chunk], top_n: int = 10
+    ) -> list[tuple[Chunk, float]]:
+        if not chunks:
+            return []
+
+        # Prepare documents for reranking
+        documents = [chunk.content for chunk in chunks]
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self._base_url}/v1/rerank",
+                json={"model": self._model, "query": query, "documents": documents},
+                headers={
+                    "accept": "application/json",
+                    "Content-Type": "application/json",
+                },
+            )
+            response.raise_for_status()
+
+            result = response.json()
+
+            # Extract scores and pair with chunks
+            scored_chunks = []
+            for item in result.get("results", []):
+                index = item["index"]
+                score = item["relevance_score"]
+                scored_chunks.append((chunks[index], score))
+
+            # Sort by score (descending) and return top_n
+            scored_chunks.sort(key=lambda x: x[1], reverse=True)
+            return scored_chunks[:top_n]

tests/test_reranker.py

@@ -2,9 +2,11 @@ import pytest
 
 from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.reranking.vllm import VLLMReranker
 from haiku.rag.store.models.chunk import Chunk
 
 COHERE_AVAILABLE = bool(Config.COHERE_API_KEY)
+VLLM_RERANK_AVAILABLE = bool(Config.VLLM_RERANK_BASE_URL)
 
 chunks = [
     Chunk(content=content, document_id=str(i))
@@ -66,3 +68,22 @@ async def test_cohere_reranker():
 
     except ImportError:
         pytest.skip("Cohere package not installed")
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not VLLM_RERANK_AVAILABLE, reason="vLLM rerank server not configured"
+)
+async def test_vllm_reranker():
+    try:
+        reranker = VLLMReranker("mixedbread-ai/mxbai-rerank-base-v2")
+
+        reranked = await reranker.rerank(
+            "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
+        )
+        assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
+        assert all(isinstance(score, float) for chunk, score in reranked)
+
+    except Exception:
+        # Skip test if vLLM rerank server is not available
+        pytest.skip("vLLM rerank server not available")