Merge pull request #120 from ggozad/fix/embedder-config
Pass custom config to embedders if provided.
This commit is contained in:
commit
a6eccc79e4
7 changed files with 152 additions and 13 deletions
|
|
@ -129,9 +129,11 @@ This is useful for:
|
||||||
- Testing with different configurations
|
- Testing with different configurations
|
||||||
- Applications that need multiple clients with different configurations
|
- Applications that need multiple clients with different configurations
|
||||||
|
|
||||||
## API Keys
|
## Environment Variables
|
||||||
|
|
||||||
API keys are configured through **environment variables**, not in the YAML file.
|
API keys and some provider settings are configured through **environment variables**, not in the YAML file.
|
||||||
|
|
||||||
|
### API Keys
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# OpenAI
|
# OpenAI
|
||||||
|
|
@ -147,6 +149,13 @@ export VOYAGE_API_KEY=your-key-here
|
||||||
export CO_API_KEY=your-key-here
|
export CO_API_KEY=your-key-here
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Provider Configuration
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Ollama base URL (defaults to http://localhost:11434)
|
||||||
|
export OLLAMA_BASE_URL=http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
## File Monitoring
|
## File Monitoring
|
||||||
|
|
||||||
Set directories to monitor for automatic indexing:
|
Set directories to monitor for automatic indexing:
|
||||||
|
|
@ -171,6 +180,23 @@ embeddings:
|
||||||
vector_dim: 1024
|
vector_dim: 1024
|
||||||
```
|
```
|
||||||
|
|
||||||
|
The Ollama base URL can be configured via environment variable or config file:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Via environment variable (recommended)
|
||||||
|
export OLLAMA_BASE_URL=http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
Or in your config file:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
providers:
|
||||||
|
ollama:
|
||||||
|
base_url: http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
If neither is set, it defaults to `http://localhost:11434`.
|
||||||
|
|
||||||
### VoyageAI
|
### VoyageAI
|
||||||
|
|
||||||
If you want to use VoyageAI embeddings you will need to install `haiku.rag` with the VoyageAI extras:
|
If you want to use VoyageAI embeddings you will need to install `haiku.rag` with the VoyageAI extras:
|
||||||
|
|
@ -236,7 +262,17 @@ Configure which LLM provider to use for question answering. Any provider and mod
|
||||||
qa:
|
qa:
|
||||||
provider: ollama
|
provider: ollama
|
||||||
model: gpt-oss
|
model: gpt-oss
|
||||||
|
```
|
||||||
|
|
||||||
|
The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export OLLAMA_BASE_URL=http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
Or in your config file:
|
||||||
|
|
||||||
|
```yaml
|
||||||
providers:
|
providers:
|
||||||
ollama:
|
ollama:
|
||||||
base_url: http://localhost:11434
|
base_url: http://localhost:11434
|
||||||
|
|
|
||||||
|
|
@ -46,7 +46,11 @@ class ProcessingConfig(BaseModel):
|
||||||
|
|
||||||
|
|
||||||
class OllamaConfig(BaseModel):
|
class OllamaConfig(BaseModel):
|
||||||
base_url: str = "http://localhost:11434"
|
base_url: str = Field(
|
||||||
|
default_factory=lambda: __import__("os").environ.get(
|
||||||
|
"OLLAMA_BASE_URL", "http://localhost:11434"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class VLLMConfig(BaseModel):
|
class VLLMConfig(BaseModel):
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,9 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase:
|
||||||
"""
|
"""
|
||||||
|
|
||||||
if config.embeddings.provider == "ollama":
|
if config.embeddings.provider == "ollama":
|
||||||
return OllamaEmbedder(config.embeddings.model, config.embeddings.vector_dim)
|
return OllamaEmbedder(
|
||||||
|
config.embeddings.model, config.embeddings.vector_dim, config
|
||||||
|
)
|
||||||
|
|
||||||
if config.embeddings.provider == "voyageai":
|
if config.embeddings.provider == "voyageai":
|
||||||
try:
|
try:
|
||||||
|
|
@ -26,16 +28,22 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase:
|
||||||
"Please install haiku.rag with the 'voyageai' extra: "
|
"Please install haiku.rag with the 'voyageai' extra: "
|
||||||
"uv pip install haiku.rag[voyageai]"
|
"uv pip install haiku.rag[voyageai]"
|
||||||
)
|
)
|
||||||
return VoyageAIEmbedder(config.embeddings.model, config.embeddings.vector_dim)
|
return VoyageAIEmbedder(
|
||||||
|
config.embeddings.model, config.embeddings.vector_dim, config
|
||||||
|
)
|
||||||
|
|
||||||
if config.embeddings.provider == "openai":
|
if config.embeddings.provider == "openai":
|
||||||
from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
|
from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
|
||||||
|
|
||||||
return OpenAIEmbedder(config.embeddings.model, config.embeddings.vector_dim)
|
return OpenAIEmbedder(
|
||||||
|
config.embeddings.model, config.embeddings.vector_dim, config
|
||||||
|
)
|
||||||
|
|
||||||
if config.embeddings.provider == "vllm":
|
if config.embeddings.provider == "vllm":
|
||||||
from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
|
from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
|
||||||
|
|
||||||
return VllmEmbedder(config.embeddings.model, config.embeddings.vector_dim)
|
return VllmEmbedder(
|
||||||
|
config.embeddings.model, config.embeddings.vector_dim, config
|
||||||
|
)
|
||||||
|
|
||||||
raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")
|
raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")
|
||||||
|
|
|
||||||
|
|
@ -1,15 +1,17 @@
|
||||||
from typing import overload
|
from typing import overload
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
|
|
||||||
|
|
||||||
class EmbedderBase:
|
class EmbedderBase:
|
||||||
_model: str = Config.embeddings.model
|
_model: str = Config.embeddings.model
|
||||||
_vector_dim: int = Config.embeddings.vector_dim
|
_vector_dim: int = Config.embeddings.vector_dim
|
||||||
|
_config: AppConfig = Config
|
||||||
|
|
||||||
def __init__(self, model: str, vector_dim: int):
|
def __init__(self, model: str, vector_dim: int, config: AppConfig = Config):
|
||||||
self._model = model
|
self._model = model
|
||||||
self._vector_dim = vector_dim
|
self._vector_dim = vector_dim
|
||||||
|
self._config = config
|
||||||
|
|
||||||
@overload
|
@overload
|
||||||
async def embed(self, text: str) -> list[float]: ...
|
async def embed(self, text: str) -> list[float]: ...
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,6 @@ from typing import overload
|
||||||
|
|
||||||
from openai import AsyncOpenAI
|
from openai import AsyncOpenAI
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.embeddings.base import EmbedderBase
|
from haiku.rag.embeddings.base import EmbedderBase
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -15,7 +14,7 @@ class Embedder(EmbedderBase):
|
||||||
|
|
||||||
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
||||||
client = AsyncOpenAI(
|
client = AsyncOpenAI(
|
||||||
base_url=f"{Config.providers.ollama.base_url}/v1", api_key="dummy"
|
base_url=f"{self._config.providers.ollama.base_url}/v1", api_key="dummy"
|
||||||
)
|
)
|
||||||
if not text:
|
if not text:
|
||||||
return []
|
return []
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,6 @@ from typing import overload
|
||||||
|
|
||||||
from openai import AsyncOpenAI
|
from openai import AsyncOpenAI
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.embeddings.base import EmbedderBase
|
from haiku.rag.embeddings.base import EmbedderBase
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -15,7 +14,8 @@ class Embedder(EmbedderBase):
|
||||||
|
|
||||||
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
|
||||||
client = AsyncOpenAI(
|
client = AsyncOpenAI(
|
||||||
base_url=f"{Config.providers.vllm.embeddings_base_url}/v1", api_key="dummy"
|
base_url=f"{self._config.providers.vllm.embeddings_base_url}/v1",
|
||||||
|
api_key="dummy",
|
||||||
)
|
)
|
||||||
if not text:
|
if not text:
|
||||||
return []
|
return []
|
||||||
|
|
|
||||||
90
tests/test_embedder_config.py
Normal file
90
tests/test_embedder_config.py
Normal file
|
|
@ -0,0 +1,90 @@
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.config import (
|
||||||
|
AppConfig,
|
||||||
|
EmbeddingsConfig,
|
||||||
|
OllamaConfig,
|
||||||
|
ProvidersConfig,
|
||||||
|
VLLMConfig,
|
||||||
|
)
|
||||||
|
from haiku.rag.embeddings import get_embedder
|
||||||
|
|
||||||
|
|
||||||
|
def test_embedder_uses_config_from_get_embedder():
|
||||||
|
"""Test that embedders use the config passed to get_embedder."""
|
||||||
|
custom_config = AppConfig(
|
||||||
|
embeddings=EmbeddingsConfig(
|
||||||
|
provider="ollama",
|
||||||
|
model="custom-model",
|
||||||
|
vector_dim=512,
|
||||||
|
),
|
||||||
|
providers=ProvidersConfig(
|
||||||
|
ollama=OllamaConfig(base_url="http://custom-ollama:8080"),
|
||||||
|
vllm=VLLMConfig(embeddings_base_url="http://custom-vllm:9000"),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
embedder = get_embedder(custom_config)
|
||||||
|
|
||||||
|
assert embedder._model == "custom-model"
|
||||||
|
assert embedder._vector_dim == 512
|
||||||
|
assert embedder._config.providers.ollama.base_url == "http://custom-ollama:8080"
|
||||||
|
|
||||||
|
|
||||||
|
def test_vllm_embedder_uses_config():
|
||||||
|
"""Test that vllm embedder uses the config passed to get_embedder."""
|
||||||
|
custom_config = AppConfig(
|
||||||
|
embeddings=EmbeddingsConfig(
|
||||||
|
provider="vllm",
|
||||||
|
model="custom-vllm-model",
|
||||||
|
vector_dim=768,
|
||||||
|
),
|
||||||
|
providers=ProvidersConfig(
|
||||||
|
vllm=VLLMConfig(embeddings_base_url="http://custom-vllm:9001"),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
embedder = get_embedder(custom_config)
|
||||||
|
|
||||||
|
assert embedder._model == "custom-vllm-model"
|
||||||
|
assert embedder._vector_dim == 768
|
||||||
|
assert (
|
||||||
|
embedder._config.providers.vllm.embeddings_base_url == "http://custom-vllm:9001"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_openai_embedder_uses_config():
|
||||||
|
"""Test that openai embedder uses the config passed to get_embedder."""
|
||||||
|
custom_config = AppConfig(
|
||||||
|
embeddings=EmbeddingsConfig(
|
||||||
|
provider="openai",
|
||||||
|
model="text-embedding-3-large",
|
||||||
|
vector_dim=3072,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
embedder = get_embedder(custom_config)
|
||||||
|
|
||||||
|
assert embedder._model == "text-embedding-3-large"
|
||||||
|
assert embedder._vector_dim == 3072
|
||||||
|
assert embedder._config == custom_config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
True, reason="VoyageAI is an optional dependency, may not be installed"
|
||||||
|
)
|
||||||
|
def test_voyageai_embedder_uses_config():
|
||||||
|
"""Test that voyageai embedder uses the config passed to get_embedder."""
|
||||||
|
custom_config = AppConfig(
|
||||||
|
embeddings=EmbeddingsConfig(
|
||||||
|
provider="voyageai",
|
||||||
|
model="voyage-large-2",
|
||||||
|
vector_dim=1536,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
embedder = get_embedder(custom_config)
|
||||||
|
|
||||||
|
assert embedder._model == "voyage-large-2"
|
||||||
|
assert embedder._vector_dim == 1536
|
||||||
|
assert embedder._config == custom_config
|
||||||
Loading…
Reference in a new issue