Skip to content

Commit b8a3a14

Browse files
committed
Add semantic query timing and FastEmbed parallel guardrails
Signed-off-by: phernandez <paul@basicmachines.co>
1 parent 9b199c6 commit b8a3a14

3 files changed

Lines changed: 343 additions & 9 deletions

File tree

src/basic_memory/repository/fastembed_provider.py

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import asyncio
66
from typing import TYPE_CHECKING
77

8+
from loguru import logger
9+
810
from basic_memory.repository.embedding_provider import EmbeddingProvider
911
from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
1012

@@ -19,6 +21,9 @@ class FastEmbedEmbeddingProvider(EmbeddingProvider):
1921
"bge-small-en-v1.5": "BAAI/bge-small-en-v1.5",
2022
}
2123

24+
def _effective_parallel(self) -> int | None:
25+
return self.parallel if self.parallel is not None and self.parallel > 1 else None
26+
2227
def __init__(
2328
self,
2429
model_name: str = "bge-small-en-v1.5",
@@ -71,18 +76,39 @@ def _create_model() -> "TextEmbedding":
7176
return TextEmbedding(model_name=resolved_model_name)
7277

7378
self._model = await asyncio.to_thread(_create_model)
79+
logger.info(
80+
"FastEmbed model loaded: model_name={model_name} batch_size={batch_size} "
81+
"threads={threads} configured_parallel={configured_parallel} "
82+
"effective_parallel={effective_parallel}",
83+
model_name=self._MODEL_ALIASES.get(self.model_name, self.model_name),
84+
batch_size=self.batch_size,
85+
threads=self.threads,
86+
configured_parallel=self.parallel,
87+
effective_parallel=self._effective_parallel(),
88+
)
7489
return self._model
7590

7691
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
7792
if not texts:
7893
return []
7994

8095
model = await self._load_model()
96+
effective_parallel = self._effective_parallel()
97+
logger.debug(
98+
"FastEmbed embed_documents call: text_count={text_count} batch_size={batch_size} "
99+
"threads={threads} configured_parallel={configured_parallel} "
100+
"effective_parallel={effective_parallel}",
101+
text_count=len(texts),
102+
batch_size=self.batch_size,
103+
threads=self.threads,
104+
configured_parallel=self.parallel,
105+
effective_parallel=effective_parallel,
106+
)
81107

82108
def _embed_batch() -> list[list[float]]:
83109
embed_kwargs: dict[str, int] = {"batch_size": self.batch_size}
84-
if self.parallel is not None:
85-
embed_kwargs["parallel"] = self.parallel
110+
if effective_parallel is not None:
111+
embed_kwargs["parallel"] = effective_parallel
86112
vectors = list(model.embed(texts, **embed_kwargs))
87113
normalized: list[list[float]] = []
88114
for vector in vectors:

0 commit comments

Comments
 (0)