diff --git a/README.md b/README.md index 14df8b4a..fd62392b 100644 --- a/README.md +++ b/README.md @@ -27,24 +27,24 @@ If you want to use an alternative embeddings provider (Ollama being the default) By default: ```bash -EMBEDDING_PROVIDER="ollama" -EMBEDDING_MODEL="mxbai-embed-large" # or any other model -EMBEDDING_VECTOR_DIM=1024 +EMBEDDINGS_PROVIDER="ollama" +EMBEDDINGS_MODEL="mxbai-embed-large" # or any other model +EMBEDDINGS_VECTOR_DIM=1024 ``` For VoyageAI: ```bash -EMBEDDING_PROVIDER="voyageai" -EMBEDDING_MODEL="voyage-3.5" # or any other model -EMBEDDING_VECTOR_DIM=1024 +EMBEDDINGS_PROVIDER="voyageai" +EMBEDDINGS_MODEL="voyage-3.5" # or any other model +EMBEDDINGS_VECTOR_DIM=1024 VOYAGE_API_KEY="your-api-key" ``` For OpenAI: ```bash -EMBEDDING_PROVIDER="openai" -EMBEDDING_MODEL="text-embedding-3-small" # or text-embedding-3-large -EMBEDDING_VECTOR_DIM=1536 +EMBEDDINGS_PROVIDER="openai" +EMBEDDINGS_MODEL="text-embedding-3-small" # or text-embedding-3-large +EMBEDDINGS_VECTOR_DIM=1536 OPENAI_API_KEY="your-api-key" ``` diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index ffbe3054..dd8546d9 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -14,9 +14,9 @@ class AppConfig(BaseModel): DEFAULT_DATA_DIR: Path = get_default_data_dir() - EMBEDDING_PROVIDER: str = "ollama" - EMBEDDING_MODEL: str = "mxbai-embed-large" - EMBEDDING_VECTOR_DIM: int = 1024 + EMBEDDINGS_PROVIDER: str = "ollama" + EMBEDDINGS_MODEL: str = "mxbai-embed-large" + EMBEDDINGS_VECTOR_DIM: int = 1024 CHUNK_SIZE: int = 256 CHUNK_OVERLAP: int = 32 diff --git a/src/haiku/rag/embeddings/__init__.py b/src/haiku/rag/embeddings/__init__.py index add55026..eed18a75 100644 --- a/src/haiku/rag/embeddings/__init__.py +++ b/src/haiku/rag/embeddings/__init__.py @@ -8,10 +8,10 @@ def get_embedder() -> EmbedderBase: Factory function to get the appropriate embedder based on the configuration. """ - if Config.EMBEDDING_PROVIDER == "ollama": - return OllamaEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM) + if Config.EMBEDDINGS_PROVIDER == "ollama": + return OllamaEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM) - if Config.EMBEDDING_PROVIDER == "voyageai": + if Config.EMBEDDINGS_PROVIDER == "voyageai": try: from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder except ImportError: @@ -20,9 +20,9 @@ def get_embedder() -> EmbedderBase: "Please install haiku.rag with the 'voyageai' extra:" "uv pip install haiku.rag --extra voyageai" ) - return VoyageAIEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM) + return VoyageAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM) - if Config.EMBEDDING_PROVIDER == "openai": + if Config.EMBEDDINGS_PROVIDER == "openai": try: from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder except ImportError: @@ -31,6 +31,6 @@ def get_embedder() -> EmbedderBase: "Please install haiku.rag with the 'openai' extra:" "uv pip install haiku.rag --extra openai" ) - return OpenAIEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM) + return OpenAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM) - raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDING_PROVIDER}") + raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDINGS_PROVIDER}") diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index 6de861a9..d7aa97a7 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/src/haiku/rag/embeddings/ollama.py @@ -5,7 +5,7 @@ from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDING_MODEL + _model: str = Config.EMBEDDINGS_MODEL _vector_dim: int = 1024 async def embed(self, text: str) -> list[float]: diff --git a/src/haiku/rag/embeddings/openai.py b/src/haiku/rag/embeddings/openai.py index 15161937..024705cd 100644 --- a/src/haiku/rag/embeddings/openai.py +++ b/src/haiku/rag/embeddings/openai.py @@ -5,7 +5,7 @@ try: from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDING_MODEL + _model: str = Config.EMBEDDINGS_MODEL _vector_dim: int = 1536 async def embed(self, text: str) -> list[float]: diff --git a/src/haiku/rag/embeddings/voyageai.py b/src/haiku/rag/embeddings/voyageai.py index 63b7da30..d37378c7 100644 --- a/src/haiku/rag/embeddings/voyageai.py +++ b/src/haiku/rag/embeddings/voyageai.py @@ -5,7 +5,7 @@ try: from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDING_MODEL + _model: str = Config.EMBEDDINGS_MODEL _vector_dim: int = 1024 async def embed(self, text: str) -> list[float]: diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 46ebde61..f6421631 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -50,8 +50,8 @@ async def test_similarity(): @pytest.mark.asyncio async def test_openai_embedder(monkeypatch): - monkeypatch.setenv("EMBEDDING_PROVIDER", "openai") - monkeypatch.setenv("EMBEDDING_MODEL", "text-embedding-3-small") + monkeypatch.setenv("EMBEDDINGS_PROVIDER", "openai") + monkeypatch.setenv("EMBEDDINGS_MODEL", "text-embedding-3-small") try: from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder @@ -94,8 +94,8 @@ async def test_openai_embedder(monkeypatch): @pytest.mark.asyncio async def test_voyageai_embedder(monkeypatch): - monkeypatch.setenv("EMBEDDING_PROVIDER", "voyageai") - monkeypatch.setenv("EMBEDDING_MODEL", "voyage-3.5") + monkeypatch.setenv("EMBEDDINGS_PROVIDER", "voyageai") + monkeypatch.setenv("EMBEDDINGS_MODEL", "voyage-3.5") try: from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder