diff --git a/src/haiku/rag/embeddings/base.py b/src/haiku/rag/embeddings/base.py index 16e19d9c..369a53a6 100644 --- a/src/haiku/rag/embeddings/base.py +++ b/src/haiku/rag/embeddings/base.py @@ -1,6 +1,9 @@ +from haiku.rag.config import Config + + class EmbedderBase: - _model: str = "" - _vector_dim: int = 0 + _model: str = Config.EMBEDDINGS_MODEL + _vector_dim: int = Config.EMBEDDINGS_VECTOR_DIM def __init__(self, model: str, vector_dim: int): self._model = model diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index d7aa97a7..600afe65 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/src/haiku/rag/embeddings/ollama.py @@ -5,9 +5,6 @@ from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1024 - async def embed(self, text: str) -> list[float]: client = AsyncClient(host=Config.OLLAMA_BASE_URL) res = await client.embeddings(model=self._model, prompt=text) diff --git a/src/haiku/rag/embeddings/openai.py b/src/haiku/rag/embeddings/openai.py index 024705cd..818f0e5b 100644 --- a/src/haiku/rag/embeddings/openai.py +++ b/src/haiku/rag/embeddings/openai.py @@ -1,13 +1,9 @@ try: from openai import AsyncOpenAI - from haiku.rag.config import Config from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1536 - async def embed(self, text: str) -> list[float]: client = AsyncOpenAI() response = await client.embeddings.create( diff --git a/src/haiku/rag/embeddings/voyageai.py b/src/haiku/rag/embeddings/voyageai.py index d37378c7..ac7aa1b6 100644 --- a/src/haiku/rag/embeddings/voyageai.py +++ b/src/haiku/rag/embeddings/voyageai.py @@ -1,13 +1,9 @@ try: from voyageai.client import Client # type: ignore - from haiku.rag.config import Config from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): - _model: str = Config.EMBEDDINGS_MODEL - _vector_dim: int = 1024 - async def embed(self, text: str) -> list[float]: client = Client() res = client.embed([text], model=self._model, output_dtype="float")