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/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]