diff --git a/README.md b/README.md index 850f9df7..65217135 100644 --- a/README.md +++ b/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 diff --git a/docs/configuration.md b/docs/configuration.md index 7e472bd5..ae9dd7e0 100644 --- a/docs/configuration.md +++ b/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. @@ -85,6 +97,18 @@ QA_MODEL="claude-3-5-haiku-20241022" # or claude-3-5-sonnet-20241022, etc. ANTHROPIC_API_KEY="your-api-key" ``` +### vLLM + +For high-performance local inference, you can use vLLM to serve models with OpenAI-compatible APIs: + +```bash +QA_PROVIDER="vllm" +QA_MODEL="Qwen/Qwen3-4B" # Any model with tool support in vLLM +VLLM_QA_BASE_URL="http://localhost:8002" # vLLM server URL +``` + +**Note:** You need to run a vLLM server separately with a model that supports tool calling loaded. Consult the specific model's documentation for proper vLLM serving configuration. + ### Other Providers Any provider supported by Pydantic AI can be used. Examples include: @@ -136,6 +160,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 diff --git a/docs/installation.md b/docs/installation.md index 727fadc4..eb1e0750 100644 --- a/docs/installation.md +++ b/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,46 @@ 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 QA (requires tool calling support) +vllm serve Qwen/Qwen3-4B --port 8002 --enable-auto-tool-choice --tool-call-parser hermes + +# 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" + +# QA (optional) +QA_PROVIDER="vllm" +QA_MODEL="Qwen/Qwen3-4B" +VLLM_QA_BASE_URL="http://localhost:8002" + +# Reranking (optional) +RERANK_PROVIDER="vllm" +RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2" +VLLM_RERANK_BASE_URL="http://localhost:8001" +``` + ## Requirements - Python 3.10+ - Ollama (for default embeddings) +- vLLM server (for vLLM provider) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 285c36f5..38a0b89f 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -33,6 +33,9 @@ class AppConfig(BaseModel): CONTEXT_CHUNK_RADIUS: int = 0 OLLAMA_BASE_URL: str = "http://localhost:11434" + VLLM_EMBEDDINGS_BASE_URL: str = "" + VLLM_RERANK_BASE_URL: str = "" + VLLM_QA_BASE_URL: str = "" # Provider keys VOYAGE_API_KEY: str = "" diff --git a/src/haiku/rag/embeddings/vllm.py b/src/haiku/rag/embeddings/vllm.py new file mode 100644 index 00000000..0f9a1aee --- /dev/null +++ b/src/haiku/rag/embeddings/vllm.py @@ -0,0 +1,16 @@ +from openai import AsyncOpenAI + +from haiku.rag.config import Config +from haiku.rag.embeddings.base import EmbedderBase + + +class Embedder(EmbedderBase): + async def embed(self, text: str) -> list[float]: + client = AsyncOpenAI( + base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy" + ) + response = await client.embeddings.create( + model=self._model, + input=text, + ) + return response.data[0].embedding diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index 711106ba..594fad6e 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -2,6 +2,7 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.ollama import OllamaProvider +from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.client import HaikuRAG from haiku.rag.config import Config @@ -65,6 +66,13 @@ class QuestionAnswerAgent: model_name=model, provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), ) + elif provider == "vllm": + return OpenAIChatModel( + model_name=model, + provider=OpenAIProvider( + base_url=f"{Config.VLLM_QA_BASE_URL}/v1", api_key="none" + ), + ) else: # For all other providers, use the provider:model format return f"{provider}:{model}" diff --git a/src/haiku/rag/reranking/vllm.py b/src/haiku/rag/reranking/vllm.py new file mode 100644 index 00000000..c99ddbb4 --- /dev/null +++ b/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] diff --git a/tests/llm_judge.py b/tests/llm_judge.py index 68ffe8c8..435e517b 100644 --- a/tests/llm_judge.py +++ b/tests/llm_judge.py @@ -1,6 +1,6 @@ from pydantic import BaseModel from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIModel +from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.ollama import OllamaProvider from haiku.rag.config import Config @@ -37,9 +37,9 @@ class LLMJudgeResponseSchema(BaseModel): class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" - def __init__(self, model: str = Config.QA_MODEL): + def __init__(self, model: str = "qwen3"): # Create Ollama model - ollama_model = OpenAIModel( + ollama_model = OpenAIChatModel( model_name=model, provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), ) diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 7d227b1c..984ff889 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -4,9 +4,11 @@ import pytest from haiku.rag.config import Config from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder +from haiku.rag.embeddings.vllm import Embedder as VLLMEmbedder OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY) VOYAGEAI_AVAILABLE = bool(Config.VOYAGE_API_KEY) +VLLM_EMBEDDINGS_AVAILABLE = bool(Config.VLLM_EMBEDDINGS_BASE_URL) # Calculate cosine similarity @@ -111,3 +113,35 @@ async def test_voyageai_embedder(): except ImportError: pytest.skip("VoyageAI package not installed") + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not VLLM_EMBEDDINGS_AVAILABLE, reason="vLLM embeddings server not configured" +) +async def test_vllm_embedder(): + embedder = VLLMEmbedder("mixedbread-ai/mxbai-embed-large-v1", 512) + phrases = [ + "I enjoy eating great food.", + "Python is my favorite programming language.", + "I love to travel and see new places.", + ] + embeddings = [np.array(await embedder.embed(phrase)) for phrase in phrases] + + test_phrase = "I am going for a camping trip." + test_embedding = await embedder.embed(test_phrase) + + sims = similarities(embeddings, test_embedding) + assert max(sims) == sims[2] + + test_phrase = "When is dinner ready?" + test_embedding = await embedder.embed(test_phrase) + + sims = similarities(embeddings, test_embedding) + assert max(sims) == sims[0] + + test_phrase = "I work as a software developer." + test_embedding = await embedder.embed(test_phrase) + + sims = similarities(embeddings, test_embedding) + assert max(sims) == sims[1] diff --git a/tests/test_qa.py b/tests/test_qa.py index 837c8f54..2fa2d417 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -9,6 +9,7 @@ from .llm_judge import LLMJudge OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY) ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY) +VLLM_QA_AVAILABLE = bool(Config.VLLM_QA_BASE_URL) @pytest.mark.asyncio @@ -80,3 +81,26 @@ async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path): assert is_equivalent, ( f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not VLLM_QA_AVAILABLE, reason="vLLM QA server not configured") +async def test_qa_vllm(qa_corpus: Dataset, temp_db_path): + """Test vLLM QA with LLM judge.""" + client = HaikuRAG(temp_db_path) + qa = QuestionAnswerAgent(client, "vllm", "Qwen/Qwen3-4B") + llm_judge = LLMJudge() + + doc = qa_corpus[1] + await client.create_document( + content=doc["document_extracted"], uri=doc["document_id"] + ) + + question = doc["question"] + expected_answer = doc["answer"] + answer = await qa.answer(question) + is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) + + assert is_equivalent, ( + f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" + ) diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 5e309c0d..1d7dc980 100644 --- a/tests/test_reranker.py +++ b/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")