diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index fd5f88652..eafdfd9b1 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -195,6 +195,12 @@ def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]: ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL" ENV_EMBEDDINGS_LOCAL_FORCE_CPU = "HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU" ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE" +# Provider-common prefix env vars. Prepended to texts before each backend's +# native encode call, distinguishing query-side vs document-side inputs for +# asymmetric retrieval models (e.g., ruri v3, intfloat/multilingual-e5, bge-zh). +# Empty (default) preserves byte-identical pre-existing behavior. +ENV_EMBEDDINGS_QUERY_PREFIX = "HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX" +ENV_EMBEDDINGS_DOC_PREFIX = "HINDSIGHT_API_EMBEDDINGS_DOC_PREFIX" ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL" ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY" ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL" @@ -955,6 +961,8 @@ class HindsightConfig: embeddings_local_model: str embeddings_local_force_cpu: bool embeddings_local_trust_remote_code: bool + embeddings_query_prefix: str + embeddings_doc_prefix: str embeddings_tei_url: str | None embeddings_openai_base_url: str | None embeddings_cohere_api_key: str | None @@ -1530,6 +1538,8 @@ def from_env(cls) -> "HindsightConfig": ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE) ).lower() in ("true", "1"), + embeddings_query_prefix=os.getenv(ENV_EMBEDDINGS_QUERY_PREFIX, ""), + embeddings_doc_prefix=os.getenv(ENV_EMBEDDINGS_DOC_PREFIX, ""), embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL), embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None, embeddings_openai_batch_size=_parse_positive_int( diff --git a/hindsight-api-slim/hindsight_api/engine/embeddings.py b/hindsight-api-slim/hindsight_api/engine/embeddings.py index f6ce6bcf7..bdbd44ebb 100644 --- a/hindsight-api-slim/hindsight_api/engine/embeddings.py +++ b/hindsight-api-slim/hindsight_api/engine/embeddings.py @@ -13,6 +13,7 @@ import os import warnings from abc import ABC, abstractmethod +from typing import Literal from urllib.parse import parse_qs, urlparse, urlunparse import httpx @@ -42,6 +43,8 @@ ENV_LLM_API_KEY, ) +EmbeddingPurpose = Literal["query", "document"] + logger = logging.getLogger(__name__) @@ -51,8 +54,26 @@ class Embeddings(ABC): The embedding dimension is determined by the model and detected at initialization. The database schema is automatically adjusted to match the model's dimension. + + Subclasses implement `_encode_impl(texts)` for the actual model call. The base + class's concrete `encode(texts, purpose)` applies the configured query/document + prefix (for asymmetric retrieval models) before delegating to `_encode_impl`. """ + def __init__(self, query_prefix: str = "", doc_prefix: str = ""): + """ + Initialize prefix configuration shared across all backends. + + Args: + query_prefix: Prefix prepended to texts when `purpose="query"`. Empty + (default) skips prepending — preserves byte-identical pre-existing + behavior for backends/models that don't need asymmetric prefixes. + doc_prefix: Prefix prepended to texts when `purpose="document"`. Same + default semantics as `query_prefix`. + """ + self.query_prefix = query_prefix + self.doc_prefix = doc_prefix + @property @abstractmethod def provider_name(self) -> str: @@ -75,17 +96,33 @@ async def initialize(self) -> None: """ pass - @abstractmethod - def encode(self, texts: list[str]) -> list[list[float]]: + def _apply_prefix(self, texts: list[str], purpose: EmbeddingPurpose) -> list[str]: + prefix = self.query_prefix if purpose == "query" else self.doc_prefix + if not prefix: + return texts + return [f"{prefix}{text}" for text in texts] + + def encode(self, texts: list[str], purpose: EmbeddingPurpose = "document") -> list[list[float]]: """ - Generate embeddings for a list of texts. + Generate embeddings for a list of texts, applying purpose-specific prefix. Args: texts: List of text strings to encode + purpose: "query" or "document" — controls which configured prefix + (if any) is prepended to each input text. Returns: List of embedding vectors (each is a list of floats) """ + return self._encode_impl(self._apply_prefix(texts, purpose)) + + @abstractmethod + def _encode_impl(self, texts: list[str]) -> list[list[float]]: + """ + Backend-specific embedding generation. Receives texts with any prefix + already applied. Subclasses must implement this; callers should use + the public `encode()` method instead. + """ pass @@ -97,7 +134,15 @@ class LocalSTEmbeddings(Embeddings): The embedding dimension is auto-detected from the model. """ - def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False): + def __init__( + self, + model_name: str | None = None, + force_cpu: bool = False, + trust_remote_code: bool = False, + *, + query_prefix: str = "", + doc_prefix: str = "", + ): """ Initialize local SentenceTransformers embeddings. @@ -109,7 +154,9 @@ def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust trust_remote_code: Allow loading models with custom code (security risk). Required for some models with custom architectures. Default: False (disabled for security) + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL self.force_cpu = force_cpu self.trust_remote_code = trust_remote_code @@ -191,16 +238,7 @@ async def initialize(self) -> None: self._dimension = self._model.get_sentence_embedding_dimension() logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings for a list of texts. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._model is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -225,6 +263,9 @@ def __init__( batch_size: int = 32, max_retries: int = 3, retry_delay: float = 0.5, + *, + query_prefix: str = "", + doc_prefix: str = "", ): """ Initialize remote TEI embeddings client. @@ -235,7 +276,9 @@ def __init__( batch_size: Maximum batch size for embedding requests (default: 32) max_retries: Maximum number of retries for failed requests (default: 3) retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5) + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.base_url = base_url.rstrip("/") self.timeout = timeout self.batch_size = batch_size @@ -326,16 +369,7 @@ async def initialize(self) -> None: except httpx.HTTPError as e: raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the remote TEI server. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -386,6 +420,9 @@ def __init__( base_url: str | None = None, batch_size: int = 100, max_retries: int = 3, + *, + query_prefix: str = "", + doc_prefix: str = "", ): """ Initialize OpenAI embeddings client. @@ -396,7 +433,9 @@ def __init__( base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint) batch_size: Maximum batch size for embedding requests (default: 100) max_retries: Maximum number of retries for failed requests (default: 3) + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.api_key = api_key self.model = model self.base_url = base_url @@ -458,16 +497,7 @@ async def initialize(self) -> None: logger.info(f"Embeddings: OpenAI provider initialized (model: {self.model}, dim: {self._dimension})") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the OpenAI API. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -520,6 +550,9 @@ def __init__( batch_size: int = 96, timeout: float = 60.0, input_type: str = "search_document", + *, + query_prefix: str = "", + doc_prefix: str = "", ): """ Initialize Cohere embeddings client. @@ -533,7 +566,9 @@ def __init__( timeout: Request timeout in seconds (default: 60.0) input_type: Input type for embeddings (default: search_document). Options: search_document, search_query, classification, clustering + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.api_key = api_key self.model = model self.base_url = base_url @@ -590,16 +625,7 @@ async def initialize(self) -> None: logger.info(f"Embeddings: Cohere provider initialized (model: {self.model}, dim: {self._dimension})") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the Cohere API. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -657,6 +683,9 @@ def __init__( model: str = DEFAULT_EMBEDDINGS_LITELLM_MODEL, batch_size: int = 100, timeout: float = 60.0, + *, + query_prefix: str = "", + doc_prefix: str = "", ): """ Initialize LiteLLM embeddings client. @@ -668,7 +697,9 @@ def __init__( Use provider prefix for non-OpenAI models (e.g., cohere/embed-english-v3.0) batch_size: Maximum batch size for embedding requests (default: 100) timeout: Request timeout in seconds (default: 60.0) + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.api_base = api_base.rstrip("/") self.api_key = api_key self.model = model @@ -714,16 +745,7 @@ async def initialize(self) -> None: except httpx.HTTPError as e: raise RuntimeError(f"Failed to connect to LiteLLM proxy at {self.api_base}: {e}") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the LiteLLM proxy. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -773,6 +795,9 @@ def __init__( batch_size: int = 100, timeout: float = 60.0, encoding_format: str | None = "float", + *, + query_prefix: str = "", + doc_prefix: str = "", ): """ Initialize LiteLLM SDK embeddings client. @@ -786,7 +811,9 @@ def __init__( timeout: Request timeout in seconds (default: 60.0) encoding_format: Encoding format for embeddings (default: "float"). Set to None or empty string to omit (needed for Voyage AI, Gemini). + query_prefix / doc_prefix: see base `Embeddings` docstring. """ + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.api_key = api_key self.model = model self.api_base = api_base @@ -853,16 +880,7 @@ async def initialize(self) -> None: logger.info(f"Embeddings: LiteLLM SDK provider initialized (model: {self.model}, dim: {self._dimension})") - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the LiteLLM SDK. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors (one per input text) - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._litellm is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -932,7 +950,11 @@ def __init__( output_dimensionality: int | None = None, batch_size: int = 100, force_ipv4: bool = False, + *, + query_prefix: str = "", + doc_prefix: str = "", ): + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) self.model = model self.api_key = api_key self.vertexai_project_id = vertexai_project_id @@ -1056,16 +1078,7 @@ def _init_vertexai(self, genai) -> None: f"model={self.model}, auth={auth_method})" ) - def encode(self, texts: list[str]) -> list[list[float]]: - """ - Generate embeddings using the Google genai SDK. - - Args: - texts: List of text strings to encode - - Returns: - List of embedding vectors - """ + def _encode_impl(self, texts: list[str]) -> list[list[float]]: if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -1113,17 +1126,25 @@ def create_embeddings_from_env() -> Embeddings: config = get_config() provider = config.embeddings_provider.lower() + # Provider-common prefix kwargs — every backend constructor accepts these, + # so plumb uniformly. Missing one branch creates silent "prefix not applied + # for provider X" bugs that are invisible until benchmark time. + prefix_kwargs = { + "query_prefix": config.embeddings_query_prefix, + "doc_prefix": config.embeddings_doc_prefix, + } if provider == "tei": url = config.embeddings_tei_url if not url: raise ValueError(f"{ENV_EMBEDDINGS_TEI_URL} is required when {ENV_EMBEDDINGS_PROVIDER} is 'tei'") - return RemoteTEIEmbeddings(base_url=url) + return RemoteTEIEmbeddings(base_url=url, **prefix_kwargs) elif provider == "local": return LocalSTEmbeddings( model_name=config.embeddings_local_model, force_cpu=config.embeddings_local_force_cpu, trust_remote_code=config.embeddings_local_trust_remote_code, + **prefix_kwargs, ) elif provider == "openai": # Use dedicated embeddings API key, or fall back to LLM API key @@ -1140,6 +1161,7 @@ def create_embeddings_from_env() -> Embeddings: model=model, base_url=base_url, batch_size=config.embeddings_openai_batch_size, + **prefix_kwargs, ) elif provider == "openrouter": api_key = config.embeddings_openrouter_api_key @@ -1153,6 +1175,7 @@ def create_embeddings_from_env() -> Embeddings: model=config.embeddings_openrouter_model, base_url="https://openrouter.ai/api/v1", batch_size=config.embeddings_openai_batch_size, + **prefix_kwargs, ) elif provider == "cohere": api_key = config.embeddings_cohere_api_key @@ -1163,12 +1186,14 @@ def create_embeddings_from_env() -> Embeddings: model=config.embeddings_cohere_model, base_url=config.embeddings_cohere_base_url, output_dimensions=config.embeddings_cohere_output_dimensions, + **prefix_kwargs, ) elif provider == "litellm": return LiteLLMEmbeddings( api_base=config.embeddings_litellm_api_base, api_key=config.embeddings_litellm_api_key, model=config.embeddings_litellm_model, + **prefix_kwargs, ) elif provider == "litellm-sdk": api_key = config.embeddings_litellm_sdk_api_key @@ -1182,6 +1207,7 @@ def create_embeddings_from_env() -> Embeddings: api_base=config.embeddings_litellm_sdk_api_base, output_dimensions=config.embeddings_litellm_sdk_output_dimensions, encoding_format=config.embeddings_litellm_sdk_encoding_format, + **prefix_kwargs, ) elif provider == "google": vertexai_project_id = config.embeddings_vertexai_project_id @@ -1202,6 +1228,7 @@ def create_embeddings_from_env() -> Embeddings: vertexai_service_account_key=config.embeddings_vertexai_service_account_key, output_dimensionality=config.embeddings_gemini_output_dimensionality, force_ipv4=config.embeddings_gemini_force_ipv4, + **prefix_kwargs, ) else: raise ValueError( diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 4f633e91b..8f2f1b270 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -3002,7 +3002,9 @@ async def _search_with_retries( embedding_span.set_attribute("hindsight.query", query[:100]) try: - query_embeddings = await embedding_utils.generate_embeddings_batch(self.embeddings, [query]) + query_embeddings = await embedding_utils.generate_embeddings_batch( + self.embeddings, [query], purpose="query" + ) query_embedding = query_embeddings[0] step_duration = time.time() - step_start log_buffer.append(f" [1] Generate query embedding: {step_duration:.3f}s") @@ -6075,7 +6077,9 @@ async def reflect_async( async def search_mental_models_fn(q: str, max_results: int = 5) -> dict[str, Any]: # Generate embedding for the query - embeddings = await embedding_utils.generate_embeddings_batch(self.embeddings, [q]) + embeddings = await embedding_utils.generate_embeddings_batch( + self.embeddings, [q], purpose="query" + ) query_embedding = embeddings[0] async with backend.acquire() as conn: return await tool_search_mental_models( diff --git a/hindsight-api-slim/hindsight_api/engine/retain/embedding_processing.py b/hindsight-api-slim/hindsight_api/engine/retain/embedding_processing.py index eccf874df..01c9f2f94 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/embedding_processing.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/embedding_processing.py @@ -6,6 +6,7 @@ import logging +from ..embeddings import EmbeddingPurpose from . import embedding_utils from .types import ExtractedFact @@ -46,13 +47,17 @@ def augment_texts_with_dates(facts: list[ExtractedFact], format_date_fn) -> list return augmented_texts -async def generate_embeddings_batch(embeddings_model, texts: list[str]) -> list[list[float]]: +async def generate_embeddings_batch( + embeddings_model, texts: list[str], purpose: EmbeddingPurpose = "document" +) -> list[list[float]]: """ Generate embeddings for a batch of texts. Args: embeddings_model: Embeddings model instance texts: List of text strings to embed + purpose: "query" or "document" — controls which configured prefix + (if any) the backend prepends. Returns: List of embedding vectors (same length as texts) @@ -60,6 +65,6 @@ async def generate_embeddings_batch(embeddings_model, texts: list[str]) -> list[ if not texts: return [] - embeddings = await embedding_utils.generate_embeddings_batch(embeddings_model, texts) + embeddings = await embedding_utils.generate_embeddings_batch(embeddings_model, texts, purpose=purpose) return embeddings diff --git a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py index 6978eef51..8f4d24270 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py @@ -4,29 +4,36 @@ import asyncio import logging +from functools import partial + +from ..embeddings import EmbeddingPurpose logger = logging.getLogger(__name__) -def generate_embedding(embeddings_backend, text: str) -> list[float]: +def generate_embedding(embeddings_backend, text: str, purpose: EmbeddingPurpose = "document") -> list[float]: """ Generate embedding for text using the provided embeddings backend. Args: embeddings_backend: Embeddings instance to use for encoding text: Text to embed + purpose: "query" or "document" — controls which configured prefix + (if any) the backend prepends. Returns: Embedding vector (dimension depends on embeddings backend) """ try: - embeddings = embeddings_backend.encode([text]) + embeddings = embeddings_backend.encode([text], purpose=purpose) return embeddings[0] except Exception as e: raise Exception(f"Failed to generate embedding: {str(e)}") -async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]: +async def generate_embeddings_batch( + embeddings_backend, texts: list[str], purpose: EmbeddingPurpose = "document" +) -> list[list[float]]: """ Generate embeddings for multiple texts using the provided embeddings backend. @@ -36,16 +43,19 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis Args: embeddings_backend: Embeddings instance to use for encoding texts: List of texts to embed + purpose: "query" or "document" — controls which configured prefix + (if any) the backend prepends. Returns: List of embeddings in same order as input texts """ try: loop = asyncio.get_event_loop() + # `run_in_executor` cannot pass kwargs directly to the callable; bind + # `purpose` via partial so the executor sees a no-kwargs callable. embeddings = await loop.run_in_executor( None, - embeddings_backend.encode, - texts, + partial(embeddings_backend.encode, texts, purpose=purpose), ) except Exception as e: raise Exception(f"Failed to generate batch embeddings: {str(e)}") diff --git a/hindsight-api-slim/tests/test_embeddings_prefix.py b/hindsight-api-slim/tests/test_embeddings_prefix.py new file mode 100644 index 000000000..181cdb2a3 --- /dev/null +++ b/hindsight-api-slim/tests/test_embeddings_prefix.py @@ -0,0 +1,158 @@ +""" +Tests for HINDSIGHT_API_EMBEDDINGS_{QUERY,DOC}_PREFIX env vars and the +template-method prefix-aware encoding interface. + +The base `Embeddings.encode(texts, purpose)` applies the configured prefix +before delegating to subclass `_encode_impl(texts)`. This is the single point +of prefix application — provider-common, so every backend honors the same env +vars without each implementing prefix logic independently. +""" + +import asyncio +import os +from typing import Any + +import pytest + +from hindsight_api.engine.embeddings import EmbeddingPurpose, Embeddings + + +class FakeEmbeddings(Embeddings): + """Test fake that records the texts passed to `_encode_impl`.""" + + def __init__(self, *, query_prefix: str = "", doc_prefix: str = ""): + super().__init__(query_prefix=query_prefix, doc_prefix=doc_prefix) + self.last_texts: list[str] | None = None + + @property + def provider_name(self) -> str: + return "fake" + + @property + def dimension(self) -> int: + return 4 + + async def initialize(self) -> None: + return None + + def _encode_impl(self, texts: list[str]) -> list[list[float]]: + self.last_texts = list(texts) + return [[0.0] * self.dimension for _ in texts] + + +def test_query_purpose_prepends_query_prefix(): + fake = FakeEmbeddings(query_prefix="検索クエリ: ", doc_prefix="検索文書: ") + fake.encode(["Hermes 起動ループ", "vchord BM25"], purpose="query") + assert fake.last_texts == ["検索クエリ: Hermes 起動ループ", "検索クエリ: vchord BM25"] + + +def test_document_purpose_prepends_doc_prefix(): + fake = FakeEmbeddings(query_prefix="検索クエリ: ", doc_prefix="検索文書: ") + fake.encode(["Hermes 起動ループ"], purpose="document") + assert fake.last_texts == ["検索文書: Hermes 起動ループ"] + + +def test_default_purpose_is_document(): + fake = FakeEmbeddings(query_prefix="Q: ", doc_prefix="D: ") + fake.encode(["x"]) # no purpose kwarg + assert fake.last_texts == ["D: x"] + + +def test_empty_prefixes_pass_through_unchanged(): + """Byte-identical to pre-patch behavior: empty defaults add no prefix.""" + fake = FakeEmbeddings() + fake.encode(["raw1", "raw2"], purpose="query") + assert fake.last_texts == ["raw1", "raw2"] + fake.encode(["doc1"], purpose="document") + assert fake.last_texts == ["doc1"] + + +def test_only_query_prefix_set_leaves_documents_unchanged(): + fake = FakeEmbeddings(query_prefix="Q: ") + fake.encode(["x"], purpose="document") + assert fake.last_texts == ["x"] + fake.encode(["x"], purpose="query") + assert fake.last_texts == ["Q: x"] + + +def test_async_executor_path_threads_purpose_correctly(): + """`generate_embeddings_batch` runs encode inside `run_in_executor` via + functools.partial — the executor must receive `purpose` as a kwarg.""" + from hindsight_api.engine.retain import embedding_utils + + fake = FakeEmbeddings(query_prefix="Q: ", doc_prefix="D: ") + results = asyncio.run( + embedding_utils.generate_embeddings_batch(fake, ["hello", "world"], purpose="query") + ) + assert len(results) == 2 + assert fake.last_texts == ["Q: hello", "Q: world"] + + results = asyncio.run( + embedding_utils.generate_embeddings_batch(fake, ["a doc"], purpose="document") + ) + assert fake.last_texts == ["D: a doc"] + + +def test_async_default_purpose_is_document(): + """Backwards-compatible default for callers that don't specify purpose.""" + from hindsight_api.engine.retain import embedding_utils + + fake = FakeEmbeddings(doc_prefix="D: ") + asyncio.run(embedding_utils.generate_embeddings_batch(fake, ["x"])) + assert fake.last_texts == ["D: x"] + + +def test_purpose_type_is_literal(): + """Smoke check that EmbeddingPurpose is the documented Literal.""" + assert EmbeddingPurpose is not None # importable + # Runtime: any string is accepted; static type checking enforces the literal. + # We don't enforce runtime narrowing — Python's behavior matches the contract. + + +@pytest.fixture +def env_isolation(): + """Save/restore env vars touched by these tests.""" + from hindsight_api.config import clear_config_cache + + keys = [ + "HINDSIGHT_API_LLM_PROVIDER", + "HINDSIGHT_API_EMBEDDINGS_PROVIDER", + "HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX", + "HINDSIGHT_API_EMBEDDINGS_DOC_PREFIX", + "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL", + ] + original: dict[str, Any] = {k: os.environ.get(k) for k in keys} + clear_config_cache() + yield + for k, v in original.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + clear_config_cache() + + +def test_config_reads_prefix_env_vars(env_isolation): + """HINDSIGHT_API_EMBEDDINGS_{QUERY,DOC}_PREFIX populate the config dataclass.""" + from hindsight_api.config import HindsightConfig + + os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock" + os.environ["HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX"] = "検索クエリ: " + os.environ["HINDSIGHT_API_EMBEDDINGS_DOC_PREFIX"] = "検索文書: " + + config = HindsightConfig.from_env() + assert config.embeddings_query_prefix == "検索クエリ: " + assert config.embeddings_doc_prefix == "検索文書: " + + +def test_config_default_prefixes_are_empty(env_isolation): + """No env var set → empty strings → byte-identical pre-existing behavior.""" + from hindsight_api.config import HindsightConfig + + os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock" + os.environ.pop("HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX", None) + os.environ.pop("HINDSIGHT_API_EMBEDDINGS_DOC_PREFIX", None) + + config = HindsightConfig.from_env() + assert config.embeddings_query_prefix == "" + assert config.embeddings_doc_prefix == ""