diff --git a/custom_components/openrag/opensearch_multimodal.py b/custom_components/openrag/opensearch_multimodal.py index 4a1b1c6eb4..1ac2c096fb 100644 --- a/custom_components/openrag/opensearch_multimodal.py +++ b/custom_components/openrag/opensearch_multimodal.py @@ -1777,6 +1777,21 @@ def _is_placeholder_term(self, term_obj: dict) -> bool: # term_obj like {"filename": "__IMPOSSIBLE_VALUE__"} return any(v == "__IMPOSSIBLE_VALUE__" for v in term_obj.values()) + @staticmethod + def _is_valid_range_clause(clause: Any) -> bool: + """Return True for a single-field range clause with at least one bound.""" + if not isinstance(clause, dict) or not isinstance(clause.get("range"), dict): + return False + range_obj = clause["range"] + if len(range_obj) != 1: + return False + bounds = next(iter(range_obj.values())) + if not isinstance(bounds, dict): + return False + return any( + key in bounds and bounds[key] not in (None, "") for key in ("gte", "gt", "lte", "lt") + ) + def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]: """Convert filter expressions into OpenSearch-compatible filter clauses. @@ -1794,7 +1809,7 @@ def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]: filter_obj: Filter configuration dictionary or None Returns: - List of OpenSearch filter clauses (term/terms objects) + List of OpenSearch filter clauses (term/terms/range objects) Placeholder values with "__IMPOSSIBLE_VALUE__" are ignored """ if not filter_obj: @@ -1825,6 +1840,8 @@ def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]: field, vals = next(iter(f["terms"].items())) if isinstance(vals, list) and len(vals) > 0: explicit_clauses.append(f) + elif self._is_valid_range_clause(f): + explicit_clauses.append(f) return explicit_clauses # Case B: convert context-style maps into clauses @@ -2368,6 +2385,11 @@ def search(self, query: str | None = None) -> list[dict[str, Any]]: } } } + # Filter inside each kNN clause so the nearest neighbours are chosen + # from matching documents; the outer bool filter alone would only + # post-filter the global top-k and miss narrow ranges. + if filter_clauses: + base_query["knn"][selected_field]["filter"] = {"bool": {"filter": filter_clauses}} if use_num_candidates: query_with_candidates = copy.deepcopy(base_query) diff --git a/flows/ingestion_flow.json b/flows/ingestion_flow.json index 9db862dcb7..24cf650884 100644 --- a/flows/ingestion_flow.json +++ b/flows/ingestion_flow.json @@ -4248,7 +4248,7 @@ "last_updated": "2026-08-18T22:46:55.693Z", "legacy": false, "metadata": { - "code_hash": "1a2a5b8d0ca7", + "code_hash": "4b1d8605aa31", "dependencies": { "dependencies": [ { @@ -4409,7 +4409,7 @@ "show": true, "title_case": false, "type": "code", - "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" + "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n @staticmethod\n def _is_valid_range_clause(clause: Any) -> bool:\n \"\"\"Return True for a single-field range clause with at least one bound.\"\"\"\n if not isinstance(clause, dict) or not isinstance(clause.get(\"range\"), dict):\n return False\n range_obj = clause[\"range\"]\n if len(range_obj) != 1:\n return False\n bounds = next(iter(range_obj.values()))\n if not isinstance(bounds, dict):\n return False\n return any(\n key in bounds and bounds[key] not in (None, \"\") for key in (\"gte\", \"gt\", \"lte\", \"lt\")\n )\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms/range objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n elif self._is_valid_range_clause(f):\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n # Filter inside each kNN clause so the nearest neighbours are chosen\n # from matching documents; the outer bool filter alone would only\n # post-filter the global top-k and miss narrow ranges.\n if filter_clauses:\n base_query[\"knn\"][selected_field][\"filter\"] = {\"bool\": {\"filter\": filter_clauses}}\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" }, "docs_metadata": { "_input_type": "TableInput", diff --git a/flows/openrag_agent.json b/flows/openrag_agent.json index a269c4e627..5436557fad 100644 --- a/flows/openrag_agent.json +++ b/flows/openrag_agent.json @@ -2738,7 +2738,7 @@ "last_updated": "2026-09-16T13:50:51.381Z", "legacy": false, "metadata": { - "code_hash": "1a2a5b8d0ca7", + "code_hash": "4b1d8605aa31", "dependencies": { "dependencies": [ { @@ -2865,7 +2865,7 @@ "show": true, "title_case": false, "type": "code", - "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" + "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n @staticmethod\n def _is_valid_range_clause(clause: Any) -> bool:\n \"\"\"Return True for a single-field range clause with at least one bound.\"\"\"\n if not isinstance(clause, dict) or not isinstance(clause.get(\"range\"), dict):\n return False\n range_obj = clause[\"range\"]\n if len(range_obj) != 1:\n return False\n bounds = next(iter(range_obj.values()))\n if not isinstance(bounds, dict):\n return False\n return any(\n key in bounds and bounds[key] not in (None, \"\") for key in (\"gte\", \"gt\", \"lte\", \"lt\")\n )\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms/range objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n elif self._is_valid_range_clause(f):\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n # Filter inside each kNN clause so the nearest neighbours are chosen\n # from matching documents; the outer bool filter alone would only\n # post-filter the global top-k and miss narrow ranges.\n if filter_clauses:\n base_query[\"knn\"][selected_field][\"filter\"] = {\"bool\": {\"filter\": filter_clauses}}\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" }, "docs_metadata": { "_input_type": "TableInput", @@ -4018,4 +4018,4 @@ "assistants", "agents" ] -} \ No newline at end of file +} diff --git a/flows/openrag_nudges.json b/flows/openrag_nudges.json index 111086f0d0..d4f80e9f20 100644 --- a/flows/openrag_nudges.json +++ b/flows/openrag_nudges.json @@ -1460,7 +1460,7 @@ "last_updated": "2026-08-05T16:19:12.028Z", "legacy": false, "metadata": { - "code_hash": "1a2a5b8d0ca7", + "code_hash": "4b1d8605aa31", "dependencies": { "dependencies": [ { @@ -1623,7 +1623,7 @@ "show": true, "title_case": false, "type": "code", - "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" + "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n @staticmethod\n def _is_valid_range_clause(clause: Any) -> bool:\n \"\"\"Return True for a single-field range clause with at least one bound.\"\"\"\n if not isinstance(clause, dict) or not isinstance(clause.get(\"range\"), dict):\n return False\n range_obj = clause[\"range\"]\n if len(range_obj) != 1:\n return False\n bounds = next(iter(range_obj.values()))\n if not isinstance(bounds, dict):\n return False\n return any(\n key in bounds and bounds[key] not in (None, \"\") for key in (\"gte\", \"gt\", \"lte\", \"lt\")\n )\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms/range objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n elif self._is_valid_range_clause(f):\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n # Filter inside each kNN clause so the nearest neighbours are chosen\n # from matching documents; the outer bool filter alone would only\n # post-filter the global top-k and miss narrow ranges.\n if filter_clauses:\n base_query[\"knn\"][selected_field][\"filter\"] = {\"bool\": {\"filter\": filter_clauses}}\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" }, "docs_metadata": { "_input_type": "TableInput", @@ -2899,4 +2899,4 @@ "assistants", "agents" ] -} \ No newline at end of file +} diff --git a/flows/openrag_url_mcp.json b/flows/openrag_url_mcp.json index 78d8d60eed..042fdf4342 100644 --- a/flows/openrag_url_mcp.json +++ b/flows/openrag_url_mcp.json @@ -5925,7 +5925,7 @@ "last_updated": "2026-08-18T22:49:22.128Z", "legacy": false, "metadata": { - "code_hash": "1a2a5b8d0ca7", + "code_hash": "4b1d8605aa31", "dependencies": { "dependencies": [ { @@ -6086,7 +6086,7 @@ "show": true, "title_case": false, "type": "code", - "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" + "value": "from __future__ import annotations\n\nimport copy\nimport json\nimport unicodedata\nimport uuid\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport httpx\nfrom lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store\nfrom lfx.base.vectorstores.vector_store_connection_decorator import vector_store_connection\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MultilineInput,\n Output,\n SecretStrInput,\n StrInput,\n TableInput,\n)\nfrom lfx.log import logger\nfrom lfx.schema.data import Data\nfrom opensearchpy import OpenSearch, helpers\nfrom opensearchpy.exceptions import OpenSearchException, RequestError\n\nREQUEST_TIMEOUT = 60\nMAX_RETRIES = 5\nEMBEDDING_SPACE_PAGE_SIZE = 100\nSELECTED_EMBEDDING_PROVIDER_VAR = \"SELECTED_EMBEDDING_PROVIDER\"\nLEGACY_EMBEDDING_ROUTE_PREFIX = \"legacy:\"\nINDEXED_EMBEDDING_ROUTE_PREFIX = \"space:\"\n\n# Some PDF fonts map common kanji (日, 月, 木, 金, 土, ...) to the look-alike\n# Kangxi Radical code points in their /ToUnicode CMap, and Docling copies them\n# through unchanged. Map just those blocks back to CJK Unified Ideographs;\n# full NFKC would also rewrite full-width digits, ㈱, ①, half-width katakana.\n# Mirrors normalize_cjk_radicals in src/utils/document_processing.py.\nCJK_RADICAL_TABLE = {\n cp: unicodedata.normalize(\"NFKC\", chr(cp))\n for cp in range(0x2E80, 0x2FE0) # CJK Radicals Supplement + Kangxi Radicals\n if unicodedata.normalize(\"NFKC\", chr(cp)) != chr(cp)\n}\n\n\n@dataclass(frozen=True)\nclass EmbeddingSpace:\n \"\"\"One indexed vector space and the model route needed to query it.\"\"\"\n\n space_id: str\n route_model: str\n field_identity: str\n legacy: bool = False\n\n\ndef build_embedding_space_aggregation(\n *,\n size: int,\n qualified_after: dict[str, Any] | None = None,\n legacy_after: dict[str, Any] | None = None,\n include_qualified: bool = True,\n include_legacy: bool = True,\n) -> dict[str, Any]:\n \"\"\"Build pageable discovery aggregations for exact and legacy vector spaces.\"\"\"\n aggregations: dict[str, Any] = {}\n if include_qualified:\n composite: dict[str, Any] = {\n \"size\": size,\n \"sources\": [{\"space_id\": {\"terms\": {\"field\": \"embedding_space_id\"}}}],\n }\n if qualified_after:\n composite[\"after\"] = qualified_after\n aggregations[\"embedding_spaces\"] = {\"composite\": composite}\n\n if include_legacy:\n composite = {\n \"size\": size,\n \"sources\": [{\"model\": {\"terms\": {\"field\": \"embedding_model\"}}}],\n }\n if legacy_after:\n composite[\"after\"] = legacy_after\n aggregations[\"legacy_embedding_models\"] = {\"composite\": composite}\n return aggregations\n\n\ndef embedding_space_after_keys(\n result: dict[str, Any],\n) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:\n \"\"\"Return composite pagination cursors for exact and legacy aggregations.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_after = aggregations.get(\"embedding_spaces\", {}).get(\"after_key\")\n legacy_after = aggregations.get(\"legacy_embedding_models\", {}).get(\"after_key\")\n return qualified_after, legacy_after\n\n\ndef embedding_spaces_from_aggregation(result: dict[str, Any]) -> list[EmbeddingSpace]:\n \"\"\"Convert an OpenSearch aggregation response into stable retrieval identities.\"\"\"\n aggregations = result.get(\"aggregations\", {})\n qualified_buckets = aggregations.get(\"embedding_spaces\", {}).get(\"buckets\", [])\n legacy_buckets = aggregations.get(\"legacy_embedding_models\", {}).get(\"buckets\", [])\n\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n for bucket in qualified_buckets:\n key = bucket.get(\"key\")\n raw_space_id = key.get(\"space_id\") if isinstance(key, dict) else key\n space_id = str(raw_space_id or \"\").strip()\n if not space_id or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=f\"{INDEXED_EMBEDDING_ROUTE_PREFIX}{space_id}\",\n field_identity=space_id,\n )\n )\n\n for bucket in legacy_buckets:\n key = bucket.get(\"key\")\n raw_model = key.get(\"model\") if isinstance(key, dict) else key\n model = str(raw_model or \"\").strip()\n space_id = f\"{LEGACY_EMBEDDING_ROUTE_PREFIX}{model}\"\n if not model or space_id in seen:\n continue\n seen.add(space_id)\n spaces.append(\n EmbeddingSpace(\n space_id=space_id,\n route_model=space_id,\n field_identity=model,\n legacy=True,\n )\n )\n return spaces\n\n\n# watsonx.ai surfaces rate-limit state via these (mostly non-standard) response\n# headers. The IBM SDK acts on the x-requests-limit-* family directly; we log\n# them on a failed embedding call to aid plan/region tuning.\n_WATSONX_RATE_LIMIT_HEADERS = (\n \"x-requests-limit-rate\",\n \"x-requests-limit-remaining\",\n \"x-requests-limit-reset\",\n \"Retry-After\",\n)\n\n\ndef _log_watsonx_rate_limit_headers(error: Exception) -> None:\n \"\"\"Best-effort diagnostic: log watsonx rate-limit headers from a failed call.\n\n The watsonx SDK raises ``ApiRequestFailure``, which carries the originating\n httpx/requests ``Response`` as ``.response``. On a 429 exhaustion we surface\n the documented rate-limit headers so operators can tune throughput.\n \"\"\"\n try:\n response = getattr(error, \"response\", None)\n headers = getattr(response, \"headers\", None)\n if not headers:\n return\n status = getattr(response, \"status_code\", \"unknown\")\n observed = {\n h: headers.get(h) for h in _WATSONX_RATE_LIMIT_HEADERS if headers.get(h) is not None\n }\n if str(status) == \"429\" or observed:\n logger.warning(f\"watsonx rate-limit response (status={status}): {observed}\")\n except Exception as log_error: # never let diagnostics mask the real error\n logger.debug(f\"Could not extract watsonx rate-limit headers: {log_error}\")\n\n\ndef normalize_model_name(model_name: str) -> str:\n \"\"\"Normalize embedding model name for use as field suffix.\n\n Converts model names to valid OpenSearch field names by replacing\n special characters and ensuring alphanumeric format.\n\n Args:\n model_name: Original embedding model name (e.g., \"text-embedding-3-small\")\n\n Returns:\n Normalized field suffix (e.g., \"text_embedding_3_small\")\n \"\"\"\n normalized = model_name.lower()\n # Replace common separators with underscores\n normalized = normalized.replace(\"-\", \"_\").replace(\":\", \"_\").replace(\"/\", \"_\").replace(\".\", \"_\")\n # Remove any non-alphanumeric characters except underscores\n normalized = \"\".join(c if c.isalnum() or c == \"_\" else \"_\" for c in normalized)\n # Remove duplicate underscores\n while \"__\" in normalized:\n normalized = normalized.replace(\"__\", \"_\")\n return normalized.strip(\"_\")\n\n\ndef get_embedding_field_name(model_name: str) -> str:\n \"\"\"Get the dynamic embedding field name for a model.\n\n Args:\n model_name: Embedding model name\n\n Returns:\n Field name in format: chunk_embedding_{normalized_model_name}\n \"\"\"\n logger.info(f\"chunk_embedding_{normalize_model_name(model_name)}\")\n return f\"chunk_embedding_{normalize_model_name(model_name)}\"\n\n\nUNTRUSTED_CHUNK_FENCE_START = \"<<>>\"\nUNTRUSTED_CHUNK_FENCE_END = \"<<>>\"\n\n\ndef fence_untrusted_text(text: str) -> str:\n \"\"\"Wrap retrieved document text in delimiters marking it as untrusted data.\n\n Retrieved chunk text can contain attacker-controlled natural-language\n instructions (indirect prompt injection, VULN-13906). Fencing it lets the\n system prompt tell the LLM to treat anything between these markers as\n data only, never as instructions or tool directives.\n\n Any literal fence markers already present in `text` are escaped first, so\n a poisoned chunk can't embed a fake end-of-fence marker to break out of\n the untrusted section and have its continuation misread as trusted.\n \"\"\"\n if not text:\n return text\n escaped = text.replace(UNTRUSTED_CHUNK_FENCE_START, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_START).replace(\n UNTRUSTED_CHUNK_FENCE_END, \"\\\\\" + UNTRUSTED_CHUNK_FENCE_END\n )\n return f\"{UNTRUSTED_CHUNK_FENCE_START}\\n{escaped}\\n{UNTRUSTED_CHUNK_FENCE_END}\"\n\n\n@vector_store_connection\nclass OpenSearchVectorStoreComponentMultimodalMultiEmbedding(LCVectorStoreComponent):\n \"\"\"OpenSearch Vector Store Component with Multi-Model Hybrid Search Capabilities.\n\n This component provides vector storage and retrieval using OpenSearch, combining semantic\n similarity search (KNN) with keyword-based search for optimal results. It supports:\n - Multiple embedding models per index with dynamic field names\n - Automatic detection and querying of all available embedding models\n - Parallel embedding generation for multi-model search\n - Document ingestion with model tracking\n - Advanced filtering and aggregations\n - Flexible authentication options\n\n Features:\n - Multi-model vector storage with dynamic fields (chunk_embedding_{model_name})\n - Hybrid search combining multiple KNN queries (dis_max) + keyword matching\n - Auto-detection of available models in the index\n - Parallel query embedding generation for all detected models\n - Vector storage with configurable engines (jvector, nmslib, faiss, lucene)\n - Flexible authentication (Basic auth, JWT tokens)\n\n Model Name Resolution:\n - Priority: deployment > model > model_name attributes\n - This ensures correct matching between embedding objects and index fields\n - When multiple embeddings are provided, specify embedding_model_name to select which one to use\n - During search, each detected model in the index is matched to its corresponding embedding object\n \"\"\"\n\n display_name: str = \"OpenSearch (Multi-Model Multi-Embedding)\"\n icon: str = \"OpenSearch\"\n description: str = (\n \"Store and search documents using OpenSearch with multi-model hybrid semantic and keyword search. \"\n \"To search use the tools search_documents and raw_search. \"\n \"Search documents takes a query for vector search, for example\\n\"\n ' {search_query: \"components in openrag\"}'\n )\n\n # Keys we consider baseline\n default_keys: list[str] = [\n \"opensearch_url\",\n \"index_name\",\n *[i.name for i in LCVectorStoreComponent.inputs], # search_query, add_documents, etc.\n \"embedding\",\n \"embedding_model_name\",\n \"embedding_provider_name\",\n \"vector_field\",\n \"number_of_results\",\n \"auth_mode\",\n \"username\",\n \"password\",\n \"jwt_token\",\n \"jwt_header\",\n \"bearer_prefix\",\n \"use_ssl\",\n \"verify_certs\",\n \"filter_expression\",\n \"engine\",\n \"space_type\",\n \"ef_construction\",\n \"m\",\n \"num_candidates\",\n \"docs_metadata\",\n \"request_timeout\",\n \"max_retries\",\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ]\n _openrag_ingest_global_placeholders = {\n \"openrag_ingest_url\": \"OPENRAG_INGEST_URL\",\n \"openrag_ingest_token\": \"OPENRAG_INGEST_TOKEN\",\n \"openrag_ingest_run_id\": \"OPENRAG_INGEST_RUN_ID\",\n }\n\n inputs = [\n TableInput(\n name=\"docs_metadata\",\n display_name=\"Document Metadata\",\n info=(\n \"Additional metadata key-value pairs to be added to all ingested documents. \"\n \"Useful for tagging documents with source information, categories, or other custom attributes.\"\n ),\n table_schema=[\n {\n \"name\": \"key\",\n \"display_name\": \"Key\",\n \"type\": \"str\",\n \"description\": \"Key name\",\n },\n {\n \"name\": \"value\",\n \"display_name\": \"Value\",\n \"type\": \"str\",\n \"description\": \"Value of the metadata\",\n },\n ],\n value=[],\n input_types=[\"Data\", \"JSON\"],\n ),\n StrInput(\n name=\"opensearch_url\",\n display_name=\"OpenSearch URL\",\n value=\"http://localhost:9200\",\n info=(\n \"The connection URL for your OpenSearch cluster \"\n \"(e.g., http://localhost:9200 for local development or your cloud endpoint).\"\n ),\n ),\n StrInput(\n name=\"index_name\",\n display_name=\"Index Name\",\n value=\"langflow\",\n info=(\n \"The OpenSearch index name where documents will be stored and searched. \"\n \"Will be created automatically if it doesn't exist.\"\n ),\n ),\n DropdownInput(\n name=\"engine\",\n display_name=\"Vector Engine\",\n options=[\"nmslib\", \"faiss\", \"lucene\", \"jvector\"],\n value=\"jvector\",\n info=(\n \"Vector search engine for similarity calculations. 'nmslib' works with standard \"\n \"OpenSearch. 'jvector' requires OpenSearch 2.9+. 'lucene' requires index.knn: true. \"\n \"Amazon OpenSearch Serverless only supports 'nmslib' or 'faiss'.\"\n ),\n advanced=True,\n ),\n DropdownInput(\n name=\"space_type\",\n display_name=\"Distance Metric\",\n options=[\"l2\", \"l1\", \"cosinesimil\", \"linf\", \"innerproduct\"],\n value=\"l2\",\n info=(\n \"Distance metric for calculating vector similarity. 'l2' (Euclidean) is most common, \"\n \"'cosinesimil' for cosine similarity, 'innerproduct' for dot product.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"ef_construction\",\n display_name=\"EF Construction\",\n value=512,\n info=(\n \"Size of the dynamic candidate list during index construction. \"\n \"Higher values improve recall but increase indexing time and memory usage.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"m\",\n display_name=\"M Parameter\",\n value=16,\n info=(\n \"Number of bidirectional connections for each vector in the HNSW graph. \"\n \"Higher values improve search quality but increase memory usage and indexing time.\"\n ),\n advanced=True,\n ),\n IntInput(\n name=\"num_candidates\",\n display_name=\"Candidate Pool Size\",\n value=1000,\n info=(\n \"Number of approximate neighbors to consider for each KNN query. \"\n \"Some OpenSearch deployments do not support this parameter; set to 0 to disable.\"\n ),\n advanced=True,\n ),\n *LCVectorStoreComponent.inputs, # includes search_query, add_documents, etc.\n HandleInput(\n name=\"embedding\", display_name=\"Embedding\", input_types=[\"Embeddings\"], is_list=True\n ),\n StrInput(\n name=\"embedding_model_name\",\n display_name=\"Embedding Model Name\",\n value=\"\",\n info=(\n \"Name of the embedding model to use for ingestion. This selects which embedding from the list \"\n \"will be used to embed documents. Matches on deployment, model, model_id, or model_name. \"\n \"For duplicate deployments, use combined format: 'deployment:model' \"\n \"(e.g., 'text-embedding-ada-002:text-embedding-3-large'). \"\n \"Leave empty to use the first embedding. Error message will show all available identifiers.\"\n ),\n advanced=False,\n ),\n StrInput(\n name=\"embedding_provider_name\",\n display_name=\"Embedding Provider\",\n value=SELECTED_EMBEDDING_PROVIDER_VAR,\n info=(\n \"Exact OpenRAG provider key used for new ingestion metadata and vector fields. \"\n \"Bound to SELECTED_EMBEDDING_PROVIDER in managed OpenRAG flows.\"\n ),\n load_from_db=True,\n advanced=True,\n ),\n StrInput(\n name=\"vector_field\",\n display_name=\"Legacy Vector Field Name\",\n value=\"chunk_embedding\",\n advanced=True,\n info=(\n \"Legacy field name for backward compatibility. New documents use dynamic fields \"\n \"(chunk_embedding_{model_name}) based on the embedding_model_name.\"\n ),\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Default Result Limit\",\n value=10,\n advanced=True,\n info=(\n \"Default maximum number of search results to return when no limit is \"\n \"specified in the filter expression.\"\n ),\n ),\n MultilineInput(\n name=\"filter_expression\",\n display_name=\"Search Filters (JSON)\",\n value=\"\",\n info=(\n \"Optional JSON configuration for search filtering, result limits, and score thresholds.\\n\\n\"\n \"Format 1 - Explicit filters:\\n\"\n '{\"filter\": [{\"term\": {\"filename\":\"doc.pdf\"}}, '\n '{\"terms\":{\"owner\":[\"user1\",\"user2\"]}}], \"limit\": 10, \"score_threshold\": 1.6}\\n\\n'\n \"Format 2 - Context-style mapping:\\n\"\n '{\"data_sources\":[\"file.pdf\"], \"document_types\":[\"application/pdf\"], \"owners\":[\"user123\"]}\\n\\n'\n \"Use __IMPOSSIBLE_VALUE__ as placeholder to ignore specific filters.\"\n ),\n ),\n # ----- Auth controls (dynamic) -----\n DropdownInput(\n name=\"auth_mode\",\n display_name=\"Authentication Mode\",\n value=\"openrag\",\n options=[\"basic\", \"jwt\", \"openrag\"],\n info=(\n \"Authentication method: 'basic' for username/password authentication, \"\n \"'jwt' for JSON Web Token (Bearer) authentication, or 'openrag' to \"\n \"delegate writes to the OpenRAG backend ingest callback (no direct \"\n \"OpenSearch credentials required — only OPENRAG_* fields).\"\n ),\n real_time_refresh=True,\n advanced=False,\n ),\n StrInput(\n name=\"username\",\n display_name=\"Username\",\n value=\"admin\",\n show=True,\n ),\n SecretStrInput(\n name=\"password\",\n display_name=\"OpenSearch Password\",\n value=\"OPENSEARCH_PASSWORD\",\n load_from_db=True,\n show=True,\n ),\n SecretStrInput(\n name=\"jwt_token\",\n display_name=\"JWT Token\",\n value=\"JWT\",\n load_from_db=False,\n show=False,\n info=(\n \"Valid JSON Web Token for authentication. \"\n \"Will be sent in the Authorization header (with optional 'Bearer ' prefix).\"\n ),\n required=False,\n ),\n StrInput(\n name=\"jwt_header\",\n display_name=\"JWT Header Name\",\n value=\"Authorization\",\n show=False,\n advanced=True,\n ),\n BoolInput(\n name=\"bearer_prefix\",\n display_name=\"Prefix 'Bearer '\",\n value=False,\n show=False,\n advanced=True,\n ),\n # ----- TLS -----\n BoolInput(\n name=\"use_ssl\",\n display_name=\"Use SSL/TLS\",\n value=True,\n advanced=True,\n info=\"Enable SSL/TLS encryption for secure connections to OpenSearch.\",\n ),\n BoolInput(\n name=\"verify_certs\",\n display_name=\"Verify SSL Certificates\",\n value=False,\n advanced=True,\n info=(\n \"Verify SSL certificates when connecting. \"\n \"Disable for self-signed certificates in development environments.\"\n ),\n ),\n # ----- Timeout / Retry -----\n StrInput(\n name=\"request_timeout\",\n display_name=\"Request Timeout (seconds)\",\n value=\"60\",\n advanced=True,\n info=(\n \"Time in seconds to wait for a response from OpenSearch. \"\n \"Increase for large bulk ingestion or complex hybrid queries.\"\n ),\n ),\n StrInput(\n name=\"max_retries\",\n display_name=\"Max Retries\",\n value=\"3\",\n advanced=True,\n info=\"Number of retries for failed connections before raising an error.\",\n ),\n StrInput(\n name=\"openrag_ingest_url\",\n display_name=\"OpenRAG Ingest URL\",\n value=\"OPENRAG_INGEST_URL\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Internal OpenRAG callback URL for backend-owned document indexing.\",\n ),\n SecretStrInput(\n name=\"openrag_ingest_token\",\n display_name=\"OpenRAG Ingest Token\",\n value=\"OPENRAG_INGEST_TOKEN\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n info=\"Short-lived token used only for OpenRAG ingest callbacks.\",\n ),\n StrInput(\n name=\"openrag_ingest_run_id\",\n display_name=\"OpenRAG Ingest Run ID\",\n value=\"OPENRAG_INGEST_RUN_ID\",\n load_from_db=True,\n input_types=[\"Text\", \"Message\"],\n advanced=True,\n ),\n IntInput(\n name=\"openrag_ingest_batch_size\",\n display_name=\"OpenRAG Ingest Batch Size\",\n value=100,\n advanced=True,\n ),\n ]\n outputs = [\n Output(\n display_name=\"Search Results\",\n name=\"search_results\",\n method=\"search_documents\",\n ),\n Output(display_name=\"Raw Search\", name=\"raw_search\", method=\"raw_search\"),\n ]\n\n def raw_search(self, query: str | dict | None = None) -> Data:\n \"\"\"Execute a raw OpenSearch query against the target index.\n\n Args:\n query (dict[str, Any]): The OpenSearch query DSL dictionary.\n\n Returns:\n Data: Search results as a Data object.\n\n Raises:\n ValueError: If 'query' is not a valid OpenSearch query (must be a non-empty dict).\n \"\"\"\n raw_query = query if query is not None else self.search_query\n\n if raw_query is None or (isinstance(raw_query, str) and not raw_query.strip()):\n self.log(\"No query provided for raw search - returning empty results\")\n return Data(data={})\n\n if isinstance(raw_query, dict):\n query_body = copy.deepcopy(raw_query)\n elif isinstance(raw_query, str):\n s = raw_query.strip()\n\n # First, optimistically try to parse as JSON DSL\n try:\n query_body = json.loads(s)\n except json.JSONDecodeError:\n # Fallback: treat as a basic text query over common fields\n query_body = {\n \"query\": {\n \"multi_match\": {\n \"query\": s,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n }\n }\n }\n else:\n msg = f\"Unsupported raw_search query type: {type(raw_query)!r}\"\n raise TypeError(msg)\n\n filter_obj = self._parse_filter_expression()\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n if filter_clauses:\n if \"query\" in query_body:\n original_query = query_body[\"query\"]\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [original_query],\n \"filter\": filter_clauses,\n }\n }\n else:\n query_body[\"query\"] = {\n \"bool\": {\n \"must\": [{\"match_all\": {}}],\n \"filter\": filter_clauses,\n }\n }\n\n if filter_obj:\n # Apply limit if not already set in the raw query\n if \"size\" not in query_body:\n limit = self._resolve_limit(filter_obj, default_limit=None)\n if limit is not None:\n query_body[\"size\"] = limit\n\n # Apply score_threshold / scoreThreshold as min_score if not already set\n if \"min_score\" not in query_body:\n score_threshold = self._resolve_score_threshold(filter_obj)\n if score_threshold is not None:\n query_body[\"min_score\"] = score_threshold\n\n client = self.build_client()\n logger.info(f\"query: {query_body}\")\n resp = client.search(\n index=self.index_name,\n body=query_body,\n params={\"terminate_after\": 0},\n )\n # Remove any _source keys whose value is a list of floats (embedding vectors)\n # Minimum length threshold to identify embedding vectors\n min_vector_length = 100\n\n def is_vector(val):\n # Accepts if it's a list of numbers (float or int) and has reasonable vector length\n return (\n isinstance(val, list)\n and len(val) > min_vector_length\n and all(isinstance(x, (float, int)) for x in val)\n )\n\n if \"hits\" in resp and \"hits\" in resp[\"hits\"]:\n for hit in resp[\"hits\"][\"hits\"]:\n source = hit.get(\"_source\")\n if isinstance(source, dict):\n keys_to_remove = [k for k, v in source.items() if is_vector(v)]\n for k in keys_to_remove:\n source.pop(k)\n if \"text\" in source:\n source[\"text\"] = fence_untrusted_text(source[\"text\"])\n logger.info(f\"Raw search response (all embedding vectors removed): {resp}\")\n return Data(**resp)\n\n def _get_embedding_model_name(self, embedding_obj=None) -> str:\n \"\"\"Get the embedding model name from component config or embedding object.\n\n Priority: deployment > model > model_id > model_name\n This ensures we use the actual model being deployed, not just the configured model.\n Supports multiple embedding providers (OpenAI, Watsonx, Cohere, etc.)\n\n Args:\n embedding_obj: Specific embedding object to get name from (optional)\n\n Returns:\n Embedding model name\n\n Raises:\n ValueError: If embedding model name cannot be determined\n \"\"\"\n # First try explicit embedding_model_name input\n if hasattr(self, \"embedding_model_name\") and self.embedding_model_name:\n return self.embedding_model_name.strip()\n\n # Try to get from provided embedding object\n if embedding_obj:\n # Priority: deployment > model > model_id > model_name\n if hasattr(embedding_obj, \"deployment\") and embedding_obj.deployment:\n return str(embedding_obj.deployment)\n if hasattr(embedding_obj, \"model\") and embedding_obj.model:\n return str(embedding_obj.model)\n if hasattr(embedding_obj, \"model_id\") and embedding_obj.model_id:\n return str(embedding_obj.model_id)\n if hasattr(embedding_obj, \"model_name\") and embedding_obj.model_name:\n return str(embedding_obj.model_name)\n\n # Try to get from embedding component (legacy single embedding)\n if hasattr(self, \"embedding\") and self.embedding:\n # Handle list of embeddings\n if isinstance(self.embedding, list) and len(self.embedding) > 0:\n first_emb = self.embedding[0]\n if hasattr(first_emb, \"deployment\") and first_emb.deployment:\n return str(first_emb.deployment)\n if hasattr(first_emb, \"model\") and first_emb.model:\n return str(first_emb.model)\n if hasattr(first_emb, \"model_id\") and first_emb.model_id:\n return str(first_emb.model_id)\n if hasattr(first_emb, \"model_name\") and first_emb.model_name:\n return str(first_emb.model_name)\n # Handle single embedding\n elif not isinstance(self.embedding, list):\n if hasattr(self.embedding, \"deployment\") and self.embedding.deployment:\n return str(self.embedding.deployment)\n if hasattr(self.embedding, \"model\") and self.embedding.model:\n return str(self.embedding.model)\n if hasattr(self.embedding, \"model_id\") and self.embedding.model_id:\n return str(self.embedding.model_id)\n if hasattr(self.embedding, \"model_name\") and self.embedding.model_name:\n return str(self.embedding.model_name)\n\n msg = (\n \"Could not determine embedding model name. \"\n \"Please set the 'embedding_model_name' field or ensure the embedding component \"\n \"has a 'deployment', 'model', 'model_id', or 'model_name' attribute.\"\n )\n raise ValueError(msg)\n\n # ---------- helper functions for index management ----------\n def _default_text_mapping(\n self,\n dim: int,\n engine: str = \"jvector\",\n space_type: str = \"l2\",\n ef_search: int = 512,\n ef_construction: int = 100,\n m: int = 16,\n vector_field: str = \"vector_field\",\n ) -> dict[str, Any]:\n \"\"\"Create the default OpenSearch index mapping for vector search.\n\n This method generates the index configuration with k-NN settings optimized\n for approximate nearest neighbor search using the specified vector engine.\n Includes the embedding_model keyword field for tracking which model was used.\n\n Args:\n dim: Dimensionality of the vector embeddings\n engine: Vector search engine (jvector, nmslib, faiss, lucene)\n space_type: Distance metric for similarity calculation\n ef_search: Size of dynamic list used during search\n ef_construction: Size of dynamic list used during index construction\n m: Number of bidirectional links for each vector\n vector_field: Name of the field storing vector embeddings\n\n Returns:\n Dictionary containing OpenSearch index mapping configuration\n \"\"\"\n return {\n \"settings\": {\"index\": {\"knn\": True, \"knn.algo_param.ef_search\": ef_search}},\n \"mappings\": {\n \"properties\": {\n vector_field: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n \"embedding_model\": {\"type\": \"keyword\"}, # Track which model was used\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n },\n }\n\n def _ensure_embedding_field_mapping(\n self,\n client: OpenSearch,\n index_name: str,\n field_name: str,\n dim: int,\n engine: str,\n space_type: str,\n ef_construction: int,\n m: int,\n ) -> None:\n \"\"\"Lazily add a dynamic embedding field to the index if it doesn't exist.\n\n This allows adding new embedding models without recreating the entire index.\n Also ensures the embedding_model tracking field exists.\n\n Note: Some OpenSearch versions/configurations have issues with dynamically adding\n knn_vector mappings (NullPointerException). This method checks if the field\n already exists before attempting to add it, and gracefully skips if the field\n is already properly configured.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index name\n field_name: Dynamic field name for this embedding model\n dim: Vector dimensionality\n engine: Vector search engine\n space_type: Distance metric\n ef_construction: Construction parameter\n m: HNSW parameter\n \"\"\"\n # First, check if the field already exists and is properly mapped\n properties = self._get_index_properties(client)\n if self._is_knn_vector_field(properties, field_name):\n # Field already exists as knn_vector - verify dimensions match\n existing_dim = self._get_field_dimension(properties, field_name)\n if existing_dim is not None and existing_dim != dim:\n logger.warning(\n f\"Field '{field_name}' exists with dimension {existing_dim}, \"\n f\"but current embedding has dimension {dim}. Using existing mapping.\"\n )\n else:\n logger.info(\n f\"[OpenSearchMultimodel] Field '{field_name}' already exists\"\n f\"as knn_vector with matching dimensions - skipping mapping update\"\n )\n return\n\n # Field doesn't exist, try to add the mapping\n try:\n mapping = {\n \"properties\": {\n field_name: {\n \"type\": \"knn_vector\",\n \"dimension\": dim,\n \"method\": {\n \"name\": \"disk_ann\",\n \"space_type\": space_type,\n \"engine\": engine,\n \"parameters\": {\"ef_construction\": ef_construction, \"m\": m},\n },\n },\n # Also ensure the embedding_model tracking field exists as keyword\n \"embedding_model\": {\"type\": \"keyword\"},\n \"embedding_provider\": {\"type\": \"keyword\"},\n \"embedding_space_id\": {\"type\": \"keyword\"},\n \"embedding_dimensions\": {\"type\": \"integer\"},\n }\n }\n client.indices.put_mapping(index=index_name, body=mapping)\n logger.info(f\"Added/updated embedding field mapping: {field_name}\")\n except RequestError as e:\n error_str = str(e).lower()\n if \"invalid engine\" in error_str and \"jvector\" in error_str:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to OpenSearch 2.9+.\"\n )\n raise ValueError(msg) from e\n if \"index.knn\" in error_str:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from e\n raise\n except Exception as e:\n # Check if this is the known OpenSearch k-NN NullPointerException issue\n error_str = str(e).lower()\n if \"null\" in error_str or \"nullpointerexception\" in error_str:\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}\"\n f\"due to OpenSearch k-NN plugin issue: {e}. \"\n f\"This is a known issue with some OpenSearch versions. \"\n f\"[OpenSearchMultimodel] Skipping mapping update. \"\n f\"Please ensure the index has the correct mapping for KNN search to work.\"\n )\n # Skip and continue - ingestion will proceed, but KNN search may fail if mapping doesn't exist\n return\n logger.warning(\n f\"[OpenSearchMultimodel] Could not add embedding field mapping for {field_name}: {e}\"\n )\n raise\n\n # Verify the field was added correctly\n properties = self._get_index_properties(client)\n if not self._is_knn_vector_field(properties, field_name):\n msg = f\"Field '{field_name}' is not mapped as knn_vector. Current mapping: {properties.get(field_name)}\"\n logger.error(msg)\n raise ValueError(msg)\n\n def _validate_aoss_with_engines(self, *, is_aoss: bool, engine: str) -> None:\n \"\"\"Validate engine compatibility with Amazon OpenSearch Serverless (AOSS).\n\n Amazon OpenSearch Serverless has restrictions on which vector engines\n can be used. This method ensures the selected engine is compatible.\n\n Args:\n is_aoss: Whether the connection is to Amazon OpenSearch Serverless\n engine: The selected vector search engine\n\n Raises:\n ValueError: If AOSS is used with an incompatible engine\n \"\"\"\n if is_aoss and engine not in {\"nmslib\", \"faiss\"}:\n msg = \"Amazon OpenSearch Service Serverless only supports `nmslib` or `faiss` engines\"\n raise ValueError(msg)\n\n def _is_aoss_enabled(self, http_auth: Any) -> bool:\n \"\"\"Determine if Amazon OpenSearch Serverless (AOSS) is being used.\n\n Args:\n http_auth: The HTTP authentication object\n\n Returns:\n True if AOSS is enabled, False otherwise\n \"\"\"\n return (\n http_auth is not None and hasattr(http_auth, \"service\") and http_auth.service == \"aoss\"\n )\n\n @staticmethod\n def _openrag_input_to_str(value: Any) -> str:\n if value is None:\n return \"\"\n if hasattr(value, \"get_secret_value\"):\n value = value.get_secret_value()\n if hasattr(value, \"text\"):\n value = value.text\n return str(value or \"\").strip()\n\n def _openrag_callback_value(self, attr_name: str) -> str:\n value = self._openrag_input_to_str(getattr(self, attr_name, \"\"))\n if value == self._openrag_ingest_global_placeholders.get(attr_name):\n return \"\"\n return value\n\n def _openrag_ingest_callback_config(self) -> tuple[str, str, str] | None:\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n\n if not url and not token and not ingest_run_id:\n return None\n if not url or not token or not ingest_run_id:\n msg = \"OpenRAG ingest callback requires url, token, and ingest_run_id.\"\n raise ValueError(msg)\n\n logger.debug(\n \"[OpenRAG callback config] configured for ingest_run_id=%s (url_len=%d, token_len=%d)\",\n ingest_run_id,\n len(url),\n len(token),\n )\n return url, token, ingest_run_id\n\n def _post_openrag_ingest_batches(\n self,\n *,\n requests: list[dict],\n vector_field: str,\n text_field: str,\n ) -> None:\n callback_config = self._openrag_ingest_callback_config()\n if callback_config is None:\n return\n\n url, token, ingest_run_id = callback_config\n batch_size = max(self._parse_int_param(\"openrag_ingest_batch_size\", 100), 1)\n timeout = self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT)\n headers = {\"Authorization\": f\"Bearer {token}\"}\n\n logger.debug(\n f\"[OpenRAG ingest POST] total_chunks={len(requests)} batch_size={batch_size} timeout_s={timeout}\"\n )\n\n with httpx.Client(timeout=timeout) as client:\n total_batches = (len(requests) + batch_size - 1) // batch_size\n for batch_number, start in enumerate(range(0, len(requests), batch_size), start=1):\n batch = requests[start : start + batch_size]\n final = batch_number == total_batches\n payload = {\n \"ingest_run_id\": ingest_run_id,\n \"batch_id\": batch_number,\n \"final\": final,\n \"chunks\": [\n self._openrag_chunk_payload(\n request,\n vector_field=vector_field,\n text_field=text_field,\n )\n for request in batch\n ],\n }\n logger.debug(\n f\"[OpenRAG ingest POST] -> batch={batch_number}/{total_batches} \"\n f\"chunks={len(payload['chunks'])} final={final}\"\n )\n response = client.post(url, json=payload, headers=headers)\n logger.debug(\n f\"[OpenRAG ingest POST resp] batch={batch_number} status={response.status_code}\"\n )\n if response.status_code >= 400:\n msg = (\n \"OpenRAG ingest callback failed \"\n f\"(batch={batch_number}, status={response.status_code})\"\n )\n logger.warning(msg)\n raise RuntimeError(msg)\n\n logger.debug(f\"Posted {len(requests)} chunks to OpenRAG backend ingest callback.\")\n\n @staticmethod\n def _openrag_chunk_payload(\n request: dict,\n *,\n vector_field: str,\n text_field: str,\n ) -> dict:\n metadata = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\", vector_field, text_field}\n }\n page = metadata.get(\"page\")\n if isinstance(page, str) and page.isdigit():\n page = int(page)\n return {\n \"id\": request.get(\"_id\") or request.get(\"id\"),\n \"text\": request.get(text_field, \"\"),\n \"vector\": request[vector_field],\n \"page\": page if isinstance(page, int) else None,\n \"metadata\": metadata,\n }\n\n def _bulk_ingest_embeddings(\n self,\n client: OpenSearch,\n index_name: str,\n embeddings: list[list[float]],\n texts: list[str],\n metadatas: list[dict] | None = None,\n ids: list[str] | None = None,\n vector_field: str = \"vector_field\",\n text_field: str = \"text\",\n embedding_model: str = \"unknown\",\n embedding_provider: str | None = None,\n embedding_space_id: str | None = None,\n mapping: dict | None = None,\n max_chunk_bytes: int | None = 1 * 1024 * 1024,\n *,\n is_aoss: bool = False,\n ) -> list[str]:\n \"\"\"Efficiently ingest multiple documents with embeddings into OpenSearch.\n\n This method uses bulk operations to insert documents with their vector\n embeddings and metadata into the specified OpenSearch index. Each document\n is tagged with the embedding_model name for tracking.\n\n Args:\n client: OpenSearch client instance\n index_name: Target index for document storage\n embeddings: List of vector embeddings for each document\n texts: List of document texts\n metadatas: Optional metadata dictionaries for each document\n ids: Optional document IDs (UUIDs generated if not provided)\n vector_field: Field name for storing vector embeddings\n text_field: Field name for storing document text\n embedding_model: Name of the embedding model used\n embedding_provider: Exact provider key, when known\n embedding_space_id: Provider-qualified vector-space identity\n mapping: Optional index mapping configuration\n max_chunk_bytes: Maximum size per bulk request chunk\n is_aoss: Whether using Amazon OpenSearch Serverless\n\n Returns:\n List of document IDs that were successfully ingested\n \"\"\"\n logger.debug(f\"[OpenSearchMultimodel] Bulk ingesting embeddings for {index_name}\")\n if not mapping:\n mapping = {}\n\n requests = []\n return_ids = []\n vector_dimensions = len(embeddings[0]) if embeddings else None\n\n for i, text in enumerate(texts):\n metadata = metadatas[i] if metadatas else {}\n if vector_dimensions is not None and \"embedding_dimensions\" not in metadata:\n metadata = {**metadata, \"embedding_dimensions\": vector_dimensions}\n\n # Normalize ACL fields that may arrive as JSON strings from flows\n for key in (\"allowed_users\", \"allowed_groups\", \"allowed_principals\"):\n value = metadata.get(key)\n if isinstance(value, str):\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list):\n metadata[key] = parsed\n except (json.JSONDecodeError, TypeError):\n # Leave value as-is if it isn't valid JSON\n pass\n\n metadata_document_id = str(metadata.get(\"document_id\") or \"\").strip()\n if metadata_document_id and metadata_document_id.lower() != \"none\":\n generated_id = f\"{metadata_document_id}_{i}\"\n else:\n generated_id = str(uuid.uuid4())\n _id = ids[i] if ids else generated_id\n request = {\n \"_op_type\": \"index\",\n \"_index\": index_name,\n vector_field: embeddings[i],\n text_field: text,\n \"embedding_model\": embedding_model, # Track which model was used\n **metadata,\n }\n if embedding_provider:\n request[\"embedding_provider\"] = embedding_provider\n if embedding_space_id:\n request[\"embedding_space_id\"] = embedding_space_id\n if is_aoss:\n request[\"id\"] = _id\n else:\n request[\"_id\"] = _id\n requests.append(request)\n return_ids.append(_id)\n if metadatas:\n self.log(f\"Sample metadata: {metadatas[0] if metadatas else {}}\")\n if self._openrag_ingest_callback_config() is not None:\n self._post_openrag_ingest_batches(\n requests=requests,\n vector_field=vector_field,\n text_field=text_field,\n )\n return return_ids\n try:\n helpers.bulk(client, requests, max_chunk_bytes=max_chunk_bytes)\n except Exception as bulk_error:\n if \"Unsupported request type for filter level DLS\" not in str(bulk_error):\n raise\n logger.warning(\n \"[OpenSearchMultimodel] Bulk ingest is blocked by filter-level DLS; \"\n \"falling back to per-document index requests.\"\n )\n self._index_embeddings_individually(client, requests)\n return return_ids\n\n def _index_embeddings_individually(\n self,\n client: OpenSearch,\n requests: list[dict],\n ) -> None:\n \"\"\"Index documents one at a time when OpenSearch DLS rejects bulk writes.\"\"\"\n for request in requests:\n document_id = request.get(\"_id\") or request.get(\"id\")\n body = {\n key: value\n for key, value in request.items()\n if key not in {\"_op_type\", \"_index\", \"_id\", \"id\"}\n }\n client.index(index=request[\"_index\"], id=document_id, body=body)\n\n def _log_index_admin_skip(self, operation: str, error: Exception) -> None:\n \"\"\"Log index-admin operations that may be blocked under filter-level DLS.\"\"\"\n logger.warning(\n f\"[OpenSearchMultimodel] Could not run index-admin operation '{operation}': {error}. \"\n \"Assuming the backend pre-created the required index/mapping and continuing.\"\n )\n\n # ---------- param helpers ----------\n def _parse_int_param(self, attr_name: str, default: int) -> int:\n \"\"\"Parse a string attribute to int, returning *default* on failure.\"\"\"\n raw = getattr(self, attr_name, None)\n if raw is None or str(raw).strip() == \"\":\n return default\n try:\n value = int(str(raw).strip())\n except ValueError:\n logger.warning(\n f\"Invalid integer value '{raw}' for {attr_name}, using default {default}\"\n )\n return default\n\n if value < 0:\n logger.warning(f\"Negative value '{raw}' for {attr_name}, using default {default}\")\n return default\n\n return value\n\n # ---------- auth / client ----------\n def _build_auth_kwargs(self) -> dict[str, Any]:\n \"\"\"Build authentication configuration for OpenSearch client.\n\n Constructs the appropriate authentication parameters based on the\n selected auth mode (basic username/password or JWT token).\n\n Returns:\n Dictionary containing authentication configuration\n\n Raises:\n ValueError: If required authentication parameters are missing\n \"\"\"\n mode = (self.auth_mode or \"basic\").strip().lower()\n if mode == \"jwt\":\n token = (self.jwt_token or \"\").strip()\n if not token:\n msg = \"Auth Mode is 'jwt' but no jwt_token was provided.\"\n raise ValueError(msg)\n header_name = (self.jwt_header or \"Authorization\").strip()\n header_value = f\"Bearer {token}\" if self.bearer_prefix else token\n return {\"headers\": {header_name: header_value}}\n if mode == \"openrag\":\n # Writes are delegated to the OpenRAG backend ingest callback,\n # so no direct OpenSearch credentials are needed. Only the\n # OPENRAG_* fields are required for ingestion to function.\n url = self._openrag_callback_value(\"openrag_ingest_url\")\n token = self._openrag_callback_value(\"openrag_ingest_token\")\n ingest_run_id = self._openrag_callback_value(\"openrag_ingest_run_id\")\n missing = [\n name\n for name, value in (\n (\"openrag_ingest_url\", url),\n (\"openrag_ingest_token\", token),\n (\"openrag_ingest_run_id\", ingest_run_id),\n )\n if not (value or \"\").strip()\n ]\n if missing:\n msg = (\n \"Auth Mode is 'openrag' but required OPENRAG_* fields are \"\n f\"missing: {', '.join(missing)}.\"\n )\n raise ValueError(msg)\n return {}\n user = (self.username or \"\").strip()\n pwd = (self._openrag_input_to_str(self.password) or \"\").strip()\n if pwd == \"OPENSEARCH_PASSWORD\" or not pwd:\n from config.settings import get_opensearch_password\n\n pwd = get_opensearch_password() or \"\"\n if not user or not pwd:\n msg = \"Auth Mode is 'basic' but username/password are missing.\"\n raise ValueError(msg)\n return {\"http_auth\": (user, pwd)}\n\n def build_client(self) -> OpenSearch:\n \"\"\"Create and configure an OpenSearch client instance.\n\n Returns:\n Configured OpenSearch client ready for operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel] Building OpenSearch client\")\n auth_kwargs = self._build_auth_kwargs()\n return OpenSearch(\n hosts=[self.opensearch_url],\n use_ssl=self.use_ssl,\n verify_certs=self.verify_certs,\n ssl_assert_hostname=False,\n ssl_show_warn=False,\n timeout=self._parse_int_param(\"request_timeout\", REQUEST_TIMEOUT),\n max_retries=self._parse_int_param(\"max_retries\", MAX_RETRIES),\n retry_on_timeout=True,\n **auth_kwargs,\n )\n\n @check_cached_vector_store\n def build_vector_store(self) -> OpenSearch:\n # Return raw OpenSearch client as our \"vector store.\"\n client = self.build_client()\n\n # Check if we're in ingestion-only mode (no search query)\n has_search_query = bool((self.search_query or \"\").strip())\n if not has_search_query:\n logger.debug(\n \"[OpenSearchMultimodel] Ingestion-only mode activated: search operations will be skipped\"\n )\n logger.debug(\"[OpenSearchMultimodel] Starting ingestion mode...\")\n\n logger.debug(f\"[OpenSearchMultimodel] Embedding: {self.embedding}\")\n self._add_documents_to_vector_store(client=client)\n return client\n\n # ---------- ingest ----------\n def _add_documents_to_vector_store(self, client: OpenSearch) -> None:\n \"\"\"Process and ingest documents into the OpenSearch vector store.\n\n This method handles the complete document ingestion pipeline:\n - Prepares document data and metadata\n - Generates vector embeddings using the selected model\n - Creates appropriate index mappings with dynamic field names\n - Bulk inserts documents with vectors and model tracking\n\n Args:\n client: OpenSearch client for performing operations\n \"\"\"\n logger.debug(\"[OpenSearchMultimodel][INGESTION] _add_documents_to_vector_store called\")\n # Convert DataFrame to Data if needed using parent's method\n self.ingest_data = self._prepare_ingest_data()\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data type: \"\n f\"{type(self.ingest_data)}, length: {len(self.ingest_data) if self.ingest_data else 0}\"\n )\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] ingest_data content: \"\n f\"{self.ingest_data[:2] if self.ingest_data and len(self.ingest_data) > 0 else 'empty'}\"\n )\n\n docs = self.ingest_data or []\n if not docs:\n logger.debug(\"Ingestion complete: No documents provided\")\n return\n\n if not self.embedding:\n msg = \"Embedding handle is required to embed documents.\"\n raise ValueError(msg)\n\n # Normalize embedding to list first\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n\n # Filter out None values (fail-safe mode) - do this BEFORE checking if empty\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n # NOW check if we have any valid embeddings left after filtering\n if not embeddings_list:\n logger.warning(\n \"All embeddings returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n self.log(\n \"Embedding returned None (fail-safe mode enabled). Skipping document ingestion.\"\n )\n return\n\n logger.debug(\n f\"[OpenSearchMultimodel][INGESTION] Valid embeddings after filtering: {len(embeddings_list)}\"\n )\n self.log(\n f\"[OpenSearchMultimodel][INGESTION] Available embedding models: {len(embeddings_list)}\"\n )\n\n # Select the embedding to use for ingestion\n selected_embedding = None\n embedding_model = None\n\n # If embedding_model_name is specified, find matching embedding\n if (\n hasattr(self, \"embedding_model_name\")\n and self.embedding_model_name\n and self.embedding_model_name.strip()\n ):\n target_model_name = self.embedding_model_name.strip()\n self.log(f\"Looking for embedding model: {target_model_name}\")\n\n for emb_obj in embeddings_list:\n # Check all possible model identifiers (deployment, model, model_id, model_name)\n # Also check available_models list from EmbeddingsWithModels\n possible_names = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n if deployment:\n possible_names.append(str(deployment))\n if model:\n possible_names.append(str(model))\n if model_id:\n possible_names.append(str(model_id))\n if model_name:\n possible_names.append(str(model_name))\n\n # Also add combined identifier\n if deployment and model and deployment != model:\n possible_names.append(f\"{deployment}:{model}\")\n\n # Add all models from available_models dict\n if available_models_attr and isinstance(available_models_attr, dict):\n possible_names.extend(\n str(model_key).strip()\n for model_key in available_models_attr\n if model_key and str(model_key).strip()\n )\n\n # Match if target matches any of the possible names\n if target_model_name in possible_names:\n # Check if target is in available_models dict - use dedicated instance\n if (\n available_models_attr\n and isinstance(available_models_attr, dict)\n and target_model_name in available_models_attr\n ):\n # Use the dedicated embedding instance from the dict\n selected_embedding = available_models_attr[target_model_name]\n embedding_model = target_model_name\n self.log(\n f\"Found dedicated embedding instance for '{embedding_model}' in available_models dict\"\n )\n else:\n # Traditional identifier match\n selected_embedding = emb_obj\n embedding_model = self._get_embedding_model_name(emb_obj)\n self.log(\n f\"Found matching embedding model: {embedding_model} (matched on: {target_model_name})\"\n )\n break\n\n if not selected_embedding:\n # Build detailed list of available embeddings with all their identifiers\n available_info = []\n for idx, emb in enumerate(embeddings_list):\n emb_type = type(emb).__name__\n identifiers = []\n deployment = getattr(emb, \"deployment\", None)\n model = getattr(emb, \"model\", None)\n model_id = getattr(emb, \"model_id\", None)\n model_name = getattr(emb, \"model_name\", None)\n available_models_attr = getattr(emb, \"available_models\", None)\n\n if deployment:\n identifiers.append(f\"deployment='{deployment}'\")\n if model:\n identifiers.append(f\"model='{model}'\")\n if model_id:\n identifiers.append(f\"model_id='{model_id}'\")\n if model_name:\n identifiers.append(f\"model_name='{model_name}'\")\n\n # Add combined identifier as an option\n if deployment and model and deployment != model:\n identifiers.append(f\"combined='{deployment}:{model}'\")\n\n # Add available_models dict if present\n if available_models_attr and isinstance(available_models_attr, dict):\n identifiers.append(f\"available_models={list(available_models_attr.keys())}\")\n\n available_info.append(\n f\" [{idx}] {emb_type}: {', '.join(identifiers) if identifiers else 'No identifiers'}\"\n )\n\n msg = (\n f\"Embedding model '{target_model_name}' not found in available embeddings.\\n\\n\"\n f\"Available embeddings:\\n\" + \"\\n\".join(available_info) + \"\\n\\n\"\n \"Please set 'embedding_model_name' to one of the identifier values shown above \"\n \"(use the value after the '=' sign, without quotes).\\n\"\n \"For duplicate deployments, use the 'combined' format.\\n\"\n \"Or leave it empty to use the first embedding.\"\n )\n raise ValueError(msg)\n else:\n # Use first embedding if no model name specified\n selected_embedding = embeddings_list[0]\n embedding_model = self._get_embedding_model_name(selected_embedding)\n self.log(f\"No embedding_model_name specified, using first embedding: {embedding_model}\")\n\n embedding_provider = self._openrag_input_to_str(\n getattr(self, \"embedding_provider_name\", \"\")\n ).lower()\n if embedding_provider == SELECTED_EMBEDDING_PROVIDER_VAR.lower():\n embedding_provider = \"\"\n embedding_space_id = (\n f\"{embedding_provider}:{embedding_model}\" if embedding_provider else None\n )\n dynamic_field_name = get_embedding_field_name(embedding_space_id or embedding_model)\n\n logger.info(f\"Selected embedding model for ingestion: '{embedding_model}'\")\n self.log(f\"Using embedding model for ingestion: {embedding_model}\")\n self.log(f\"Dynamic vector field: {dynamic_field_name}\")\n\n # Log embedding details for debugging\n if hasattr(selected_embedding, \"deployment\"):\n logger.info(f\"Embedding deployment: {selected_embedding.deployment}\")\n if hasattr(selected_embedding, \"model\"):\n logger.info(f\"Embedding model: {selected_embedding.model}\")\n if hasattr(selected_embedding, \"model_id\"):\n logger.info(f\"Embedding model_id: {selected_embedding.model_id}\")\n if hasattr(selected_embedding, \"dimensions\"):\n logger.info(f\"Embedding dimensions: {selected_embedding.dimensions}\")\n if hasattr(selected_embedding, \"available_models\"):\n logger.info(f\"Embedding available_models: {selected_embedding.available_models}\")\n\n # No model switching needed - each model in available_models has its own dedicated instance\n # The selected_embedding is already configured correctly for the target model\n logger.info(\n f\"Using embedding instance for '{embedding_model}' - pre-configured and ready to use\"\n )\n\n # Extract texts and metadata from documents\n texts = []\n metadatas = []\n # Process docs_metadata table input into a dict\n additional_metadata = {}\n logger.debug(f\"[LF] Docs metadata {self.docs_metadata}\")\n if hasattr(self, \"docs_metadata\") and self.docs_metadata:\n logger.info(f\"[LF] Docs metadata {self.docs_metadata}\")\n if isinstance(self.docs_metadata[-1], Data):\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n self.docs_metadata = self.docs_metadata[-1].data\n logger.info(f\"[LF] Docs metadata is a Data object {self.docs_metadata}\")\n additional_metadata.update(self.docs_metadata)\n else:\n for item in self.docs_metadata:\n if isinstance(item, dict) and \"key\" in item and \"value\" in item:\n additional_metadata[item[\"key\"]] = item[\"value\"]\n # Replace string \"None\" values with actual None\n for key, value in additional_metadata.items():\n if value == \"None\":\n additional_metadata[key] = None\n logger.info(f\"[LF] Additional metadata {additional_metadata}\")\n for doc_obj in docs:\n data_copy = json.loads(doc_obj.model_dump_json())\n text = data_copy.pop(doc_obj.text_key, doc_obj.default_value)\n if isinstance(text, str):\n text = text.translate(CJK_RADICAL_TABLE)\n texts.append(text)\n\n # Merge additional metadata from table input\n data_copy.update(additional_metadata)\n\n metadatas.append(data_copy)\n self.log(metadatas)\n\n # Determine whether the selected embedding is watsonx/IBM. The watsonx\n # SDK ships its own rate-limit machinery (input batching, proactive\n # x-requests-limit-* TokenBucket throttling, and jittered exponential\n # backoff on 429), so we lean on it instead of retrying on top of it.\n # The type-name check also covers watsonx-hosted, non-\"ibm/\" models\n # (e.g. intfloat/multilingual-e5-large).\n is_ibm = (embedding_model and \"ibm\" in str(embedding_model).lower()) or (\n selected_embedding and \"watsonx\" in type(selected_embedding).__name__.lower()\n )\n logger.debug(f\"Is IBM/watsonx embedding: {is_ibm}\")\n\n if is_ibm:\n # Hand the full batch to the SDK and let it batch/throttle/retry.\n # Retry attempts and base backoff are tunable via the SDK's own\n # WATSONX_MAX_RETRIES / WATSONX_DELAY_TIME environment variables.\n logger.info(\n f\"Embedding {len(texts)} chunks via watsonx SDK batch (SDK-managed throttle + 429 retry)\"\n )\n try:\n vectors: list[list[float]] = selected_embedding.embed_documents(texts)\n logger.info(f\"Successfully embedded {len(vectors)} chunks via watsonx SDK\")\n except Exception as embed_error:\n _log_watsonx_rate_limit_headers(embed_error)\n logger.error(\n f\"Failed to embed {len(texts)} chunks via watsonx SDK. Error: {embed_error}\",\n )\n raise\n\n else:\n # Non-watsonx providers (OpenAI, Ollama) lack the watsonx SDK's\n # built-in rate-limit handling, so embed per chunk in parallel with\n # a generic rate-limit-aware tenacity retry.\n vectors: list[list[float]] = [None] * len(texts)\n from tenacity import (\n retry,\n retry_if_exception,\n stop_after_attempt,\n wait_exponential,\n )\n\n def is_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a rate limit error (429).\"\"\"\n error_str = str(exception).lower()\n return \"429\" in error_str or \"rate_limit\" in error_str or \"rate limit\" in error_str\n\n def is_other_retryable_error(exception: Exception) -> bool:\n \"\"\"Check if exception is a transient network error worth retrying.\"\"\"\n if is_rate_limit_error(exception):\n return False\n return isinstance(exception, (ConnectionError, TimeoutError, OSError))\n\n # Retry decorator for rate limit errors (longer backoff)\n retry_on_rate_limit = retry(\n retry=retry_if_exception(is_rate_limit_error),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=2, min=2, max=30),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Rate limit hit for chunk (attempt {retry_state.attempt_number}/5), \"\n f\"backing off for {retry_state.next_action.sleep:.1f}s\"\n ),\n )\n\n # Retry decorator for other errors (shorter backoff)\n retry_on_other_errors = retry(\n retry=retry_if_exception(is_other_retryable_error),\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=1, min=1, max=8),\n reraise=True,\n before_sleep=lambda retry_state: logger.warning(\n f\"Error embedding chunk (attempt {retry_state.attempt_number}/3), \"\n f\"retrying in {retry_state.next_action.sleep:.1f}s: {retry_state.outcome.exception()}\"\n ),\n )\n\n def embed_chunk_with_retry(chunk_text: str, chunk_idx: int) -> list[float]:\n \"\"\"Embed a single chunk with rate-limit-aware retry logic.\"\"\"\n\n @retry_on_rate_limit\n @retry_on_other_errors\n def _embed(text: str) -> list[float]:\n return selected_embedding.embed_documents([text])[0]\n\n try:\n return _embed(chunk_text)\n except Exception as e:\n logger.error(\n f\"Failed to embed chunk {chunk_idx} after all retries: {e}\",\n error=str(e),\n )\n raise\n\n max_workers = min(max(len(texts), 1), 8)\n logger.debug(f\"Using parallel processing with {max_workers} workers\")\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n futures = {\n executor.submit(embed_chunk_with_retry, chunk, idx): idx\n for idx, chunk in enumerate(texts)\n }\n for future in as_completed(futures):\n idx = futures[future]\n vectors[idx] = future.result()\n\n if not vectors:\n self.log(f\"No vectors generated from documents for model {embedding_model}.\")\n return\n\n # Get vector dimension for mapping\n dim = len(vectors[0]) if vectors else 768 # default fallback\n use_openrag_ingest_callback = self._openrag_ingest_callback_config() is not None\n\n is_aoss = False\n mapping: dict | None = None\n\n engine = getattr(self, \"engine\", \"jvector\")\n\n if use_openrag_ingest_callback:\n self.log(\"Using OpenRAG backend ingest callback; skipping direct OpenSearch writes.\")\n else:\n # Check for AOSS\n auth_kwargs = self._build_auth_kwargs()\n is_aoss = self._is_aoss_enabled(auth_kwargs.get(\"http_auth\"))\n\n # Validate engine with AOSS\n self._validate_aoss_with_engines(is_aoss=is_aoss, engine=engine)\n\n # Create mapping with proper KNN settings\n space_type = getattr(self, \"space_type\", \"l2\")\n ef_construction = getattr(self, \"ef_construction\", 512)\n m = getattr(self, \"m\", 16)\n\n mapping = self._default_text_mapping(\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n vector_field=dynamic_field_name, # Use dynamic field name\n )\n\n # Ensure index exists with baseline mapping (index.knn: true is required for vector search)\n index_exists = True\n try:\n index_exists = bool(client.indices.exists(index=self.index_name))\n except OpenSearchException as exists_error:\n self._log_index_admin_skip(\"indices.exists\", exists_error)\n\n try:\n if not index_exists:\n self.log(f\"Creating index '{self.index_name}' with base mapping\")\n client.indices.create(index=self.index_name, body=mapping)\n except RequestError as creation_error:\n if creation_error.error == \"resource_already_exists_exception\":\n pass # Index was created concurrently\n else:\n error_msg = str(creation_error).lower()\n if \"invalid engine\" in error_msg or \"illegal_argument\" in error_msg:\n if \"jvector\" in error_msg:\n msg = (\n \"The 'jvector' engine is not available in your OpenSearch installation. \"\n \"Use 'nmslib' or 'faiss' for standard OpenSearch, or upgrade to 2.9+.\"\n )\n raise ValueError(msg) from creation_error\n if \"index.knn\" in error_msg:\n msg = (\n \"The index has index.knn: false. Delete the existing index and let the \"\n \"component recreate it, or create a new index with a different name.\"\n )\n raise ValueError(msg) from creation_error\n logger.warning(f\"Failed to create index '{self.index_name}': {creation_error}\")\n raise\n\n # Ensure the dynamic field exists in the index\n self._ensure_embedding_field_mapping(\n client=client,\n index_name=self.index_name,\n field_name=dynamic_field_name,\n dim=dim,\n engine=engine,\n space_type=space_type,\n ef_construction=ef_construction,\n m=m,\n )\n\n self.log(\n f\"Indexing {len(texts)} documents into '{self.index_name}' with model '{embedding_model}'...\"\n )\n logger.info(f\"Will store embeddings in field: {dynamic_field_name}\")\n logger.info(f\"Will tag documents with embedding_model: {embedding_model}\")\n\n # Use the bulk ingestion with model tracking\n return_ids = self._bulk_ingest_embeddings(\n client=client,\n index_name=self.index_name,\n embeddings=vectors,\n texts=texts,\n metadatas=metadatas,\n vector_field=dynamic_field_name, # Use dynamic field name\n text_field=\"text\",\n embedding_model=embedding_model, # Track the model\n embedding_provider=embedding_provider or None,\n embedding_space_id=embedding_space_id,\n mapping=mapping,\n is_aoss=is_aoss,\n )\n self.log(metadatas)\n\n logger.info(\n f\"Ingestion complete: Successfully indexed {len(return_ids)} documents with model '{embedding_model}'\"\n )\n self.log(f\"Successfully indexed {len(return_ids)} documents with model {embedding_model}.\")\n\n # ---------- helpers for filters ----------\n def _is_placeholder_term(self, term_obj: dict) -> bool:\n # term_obj like {\"filename\": \"__IMPOSSIBLE_VALUE__\"}\n return any(v == \"__IMPOSSIBLE_VALUE__\" for v in term_obj.values())\n\n @staticmethod\n def _is_valid_range_clause(clause: Any) -> bool:\n \"\"\"Return True for a single-field range clause with at least one bound.\"\"\"\n if not isinstance(clause, dict) or not isinstance(clause.get(\"range\"), dict):\n return False\n range_obj = clause[\"range\"]\n if len(range_obj) != 1:\n return False\n bounds = next(iter(range_obj.values()))\n if not isinstance(bounds, dict):\n return False\n return any(\n key in bounds and bounds[key] not in (None, \"\") for key in (\"gte\", \"gt\", \"lte\", \"lt\")\n )\n\n def _coerce_filter_clauses(self, filter_obj: dict | None) -> list[dict]:\n \"\"\"Convert filter expressions into OpenSearch-compatible filter clauses.\n\n This method accepts two filter formats and converts them to standardized\n OpenSearch query clauses:\n\n Format A - Explicit filters:\n {\"filter\": [{\"term\": {\"field\": \"value\"}}, {\"terms\": {\"field\": [\"val1\", \"val2\"]}}],\n \"limit\": 10, \"score_threshold\": 1.5}\n\n Format B - Context-style mapping:\n {\"data_sources\": [\"file1.pdf\"], \"document_types\": [\"pdf\"], \"owners\": [\"user1\"]}\n\n Args:\n filter_obj: Filter configuration dictionary or None\n\n Returns:\n List of OpenSearch filter clauses (term/terms/range objects)\n Placeholder values with \"__IMPOSSIBLE_VALUE__\" are ignored\n \"\"\"\n if not filter_obj:\n return []\n\n # If it is a string, try to parse it once\n if isinstance(filter_obj, str):\n try:\n filter_obj = json.loads(filter_obj)\n except json.JSONDecodeError:\n # Not valid JSON - treat as no filters\n return []\n\n # Case A: already an explicit list/dict under \"filter\"\n if \"filter\" in filter_obj:\n raw = filter_obj[\"filter\"]\n if isinstance(raw, dict):\n raw = [raw]\n explicit_clauses: list[dict] = []\n for f in raw or []:\n if (\n \"term\" in f\n and isinstance(f[\"term\"], dict)\n and not self._is_placeholder_term(f[\"term\"])\n ):\n explicit_clauses.append(f)\n elif \"terms\" in f and isinstance(f[\"terms\"], dict):\n field, vals = next(iter(f[\"terms\"].items()))\n if isinstance(vals, list) and len(vals) > 0:\n explicit_clauses.append(f)\n elif self._is_valid_range_clause(f):\n explicit_clauses.append(f)\n return explicit_clauses\n\n # Case B: convert context-style maps into clauses\n field_mapping = {\n \"data_sources\": \"filename\",\n \"document_types\": \"mimetype\",\n \"owners\": \"owner\",\n }\n context_clauses: list[dict] = []\n for k, values in filter_obj.items():\n if not isinstance(values, list):\n continue\n field = field_mapping.get(k, k)\n if len(values) == 0:\n # Match-nothing placeholder (kept to mirror your tool semantics)\n context_clauses.append({\"term\": {field: \"__IMPOSSIBLE_VALUE__\"}})\n elif len(values) == 1:\n if values[0] != \"__IMPOSSIBLE_VALUE__\":\n context_clauses.append({\"term\": {field: values[0]}})\n else:\n context_clauses.append({\"terms\": {field: values}})\n return context_clauses\n\n def _parse_filter_expression(self) -> dict | None:\n \"\"\"Parse and validate optional filter_expression JSON.\n\n Returns:\n Parsed JSON object as a dict, or None when unset/blank.\n\n Raises:\n ValueError: If JSON is invalid or does not decode to an object.\n \"\"\"\n filter_expression = getattr(self, \"filter_expression\", \"\")\n if not isinstance(filter_expression, str) or not filter_expression.strip():\n return None\n try:\n filter_obj = json.loads(filter_expression)\n except json.JSONDecodeError as e:\n msg = f\"Invalid filter_expression JSON: {e}\"\n raise ValueError(msg) from e\n\n if not isinstance(filter_obj, dict):\n msg = \"Invalid filter_expression JSON type: expected a JSON object.\"\n raise TypeError(msg)\n return filter_obj\n\n def _resolve_limit(self, filter_obj: dict | None, default_limit: int | None) -> int | None:\n \"\"\"Resolve an integer result limit from filter settings.\"\"\"\n if not filter_obj:\n return default_limit\n raw_limit = filter_obj.get(\"limit\", default_limit)\n if raw_limit is None:\n return None\n if isinstance(raw_limit, bool):\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise TypeError(msg)\n try:\n limit = int(raw_limit)\n except (TypeError, ValueError) as e:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg) from e\n if limit <= 0:\n msg = \"Invalid filter_expression.limit: expected a positive integer.\"\n raise ValueError(msg)\n return limit\n\n def _resolve_score_threshold(self, filter_obj: dict | None) -> float | None:\n \"\"\"Resolve optional positive min score from filter settings.\"\"\"\n if not filter_obj:\n return None\n score_threshold = filter_obj.get(\"score_threshold\")\n if score_threshold is None:\n score_threshold = filter_obj.get(\"scoreThreshold\")\n if not isinstance(score_threshold, (int, float)) or score_threshold <= 0:\n return None\n return float(score_threshold)\n\n def _detect_available_models(\n self, client: OpenSearch, filter_clauses: list[dict] | None = None\n ) -> list[EmbeddingSpace]:\n \"\"\"Detect which embedding models have documents in the index.\n\n Uses aggregation to find all unique embedding_model values, optionally\n filtered to only documents matching the user's filter criteria.\n\n Args:\n client: OpenSearch client instance\n filter_clauses: Optional filter clauses to scope model detection\n\n Returns:\n List of embedding model names found in the index\n \"\"\"\n try:\n spaces: list[EmbeddingSpace] = []\n seen: set[str] = set()\n result: dict[str, Any] = {}\n\n # Composite aggregations must be top-level. Exact and legacy\n # spaces use different document filters, so paginate them in two\n # independent passes rather than nesting composite under filter.\n for legacy in (False, True):\n after = None\n while True:\n agg_query = {\n \"size\": 0,\n \"aggs\": build_embedding_space_aggregation(\n size=EMBEDDING_SPACE_PAGE_SIZE,\n qualified_after=None if legacy else after,\n legacy_after=after if legacy else None,\n include_qualified=not legacy,\n include_legacy=legacy,\n ),\n }\n bool_query: dict[str, Any] = {}\n if filter_clauses:\n bool_query[\"filter\"] = filter_clauses\n if legacy:\n bool_query[\"must_not\"] = [{\"exists\": {\"field\": \"embedding_space_id\"}}]\n if bool_query:\n agg_query[\"query\"] = {\"bool\": bool_query}\n\n logger.debug(f\"Model detection query: {agg_query}\")\n result = client.search(\n index=self.index_name,\n body=agg_query,\n params={\"terminate_after\": 0},\n )\n for space in embedding_spaces_from_aggregation(result):\n if space.space_id not in seen:\n seen.add(space.space_id)\n spaces.append(space)\n\n qualified_after, legacy_after = embedding_space_after_keys(result)\n next_after = legacy_after if legacy else qualified_after\n if not next_after or next_after == after:\n break\n after = next_after\n\n # Log detailed bucket info for debugging\n logger.info(\n f\"Detected embedding spaces in corpus: {[space.space_id for space in spaces]}\"\n + (f\" (with {len(filter_clauses)} filters)\" if filter_clauses else \"\")\n )\n if not spaces:\n total_hits = result.get(\"hits\", {}).get(\"total\", {})\n total_count = (\n total_hits.get(\"value\", 0) if isinstance(total_hits, dict) else total_hits\n )\n logger.warning(\n f\"No embedding_model values found in index '{self.index_name}'. \"\n f\"Total docs in index: {total_count}. \"\n f\"This may indicate documents were indexed without the embedding_model field.\"\n )\n except (OpenSearchException, KeyError, ValueError) as e:\n logger.warning(f\"Failed to detect embedding models: {e}\")\n # Provider provenance cannot be inferred safely after discovery\n # fails. Keyword search remains available without guessing.\n return []\n else:\n return spaces\n\n def _get_index_properties(self, client: OpenSearch) -> dict[str, Any] | None:\n \"\"\"Retrieve flattened mapping properties for the current index.\"\"\"\n try:\n mapping = client.indices.get_mapping(index=self.index_name)\n except OpenSearchException as e:\n logger.warning(\n f\"Failed to fetch mapping for index '{self.index_name}': {e}. Proceeding without mapping metadata.\"\n )\n return None\n\n properties: dict[str, Any] = {}\n for index_data in mapping.values():\n props = index_data.get(\"mappings\", {}).get(\"properties\", {})\n if isinstance(props, dict):\n properties.update(props)\n return properties\n\n def _is_knn_vector_field(self, properties: dict[str, Any] | None, field_name: str) -> bool:\n \"\"\"Check whether the field is mapped as a knn_vector.\"\"\"\n if not field_name:\n return False\n if properties is None:\n logger.warning(\n f\"Mapping metadata unavailable; assuming field '{field_name}' is usable.\"\n )\n return True\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return False\n if field_def.get(\"type\") == \"knn_vector\":\n return True\n\n nested_props = field_def.get(\"properties\")\n return bool(isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\")\n\n def _get_field_dimension(\n self, properties: dict[str, Any] | None, field_name: str\n ) -> int | None:\n \"\"\"Get the dimension of a knn_vector field from the index mapping.\n\n Args:\n properties: Index properties from mapping\n field_name: Name of the vector field\n\n Returns:\n Dimension of the field, or None if not found\n \"\"\"\n if not field_name or properties is None:\n return None\n\n field_def = properties.get(field_name)\n if not isinstance(field_def, dict):\n return None\n\n # Check direct knn_vector field\n if field_def.get(\"type\") == \"knn_vector\":\n return field_def.get(\"dimension\")\n\n # Check nested properties\n nested_props = field_def.get(\"properties\")\n if isinstance(nested_props, dict) and nested_props.get(\"type\") == \"knn_vector\":\n return nested_props.get(\"dimension\")\n\n return None\n\n def _get_filename_agg_field(self, index_properties: dict[str, Any] | None) -> str:\n \"\"\"Choose the appropriate field for filename aggregations.\"\"\"\n if not index_properties:\n return \"filename.keyword\"\n\n filename_def = index_properties.get(\"filename\")\n if not isinstance(filename_def, dict):\n return \"filename.keyword\"\n\n field_type = filename_def.get(\"type\")\n fields_def = filename_def.get(\"fields\", {})\n\n # Top-level keyword with no subfields\n if field_type == \"keyword\" and not isinstance(fields_def, dict):\n return \"filename\"\n\n # Text field with keyword subfield\n if isinstance(fields_def, dict) and \"keyword\" in fields_def:\n return \"filename.keyword\"\n\n # Fallback: aggregate on filename directly\n return \"filename\"\n\n # ---------- search (multi-model hybrid) ----------\n def search(self, query: str | None = None) -> list[dict[str, Any]]:\n \"\"\"Perform multi-model hybrid search combining multiple vector similarities and keyword matching.\n\n This method executes a sophisticated search that:\n 1. Auto-detects all embedding models present in the index\n 2. Generates query embeddings for ALL detected models in parallel\n 3. Combines multiple KNN queries using dis_max (picks best match)\n 4. Adds keyword search with fuzzy matching (30% weight)\n 5. Applies optional filtering and score thresholds\n 6. Returns aggregations for faceted search\n\n Search weights:\n - Semantic search (dis_max across all models): 70%\n - Keyword search: 30%\n\n Args:\n query: Search query string (used for both vector embedding and keyword search)\n\n Returns:\n List of search results with page_content, metadata, and relevance scores\n\n Raises:\n ValueError: If filter JSON is invalid\n \"\"\"\n logger.info(self.ingest_data)\n client = self.build_client()\n q = (query or \"\").strip()\n\n # Parse optional filter expression\n filter_obj = self._parse_filter_expression()\n\n # Build filter clauses first so we can use them in model detection\n filter_clauses = self._coerce_filter_clauses(filter_obj)\n\n # Detect available embedding models in the index (scoped by filters)\n available_spaces = self._detect_available_models(client, filter_clauses)\n\n if not available_spaces:\n logger.warning(\"No embedding spaces found in index; using keyword-only search\")\n\n # Generate embeddings for ALL detected models\n query_embeddings = {}\n\n # Normalize embedding to list\n embeddings_list = self.embedding if isinstance(self.embedding, list) else [self.embedding]\n # Filter out None values (fail-safe mode)\n embeddings_list = [e for e in embeddings_list if e is not None]\n\n if not embeddings_list:\n logger.warning(\n \"No valid embedding adapter is available; continuing with keyword-only search.\"\n )\n\n # Create a comprehensive map of model names to embedding objects\n # Check all possible identifiers (deployment, model, model_id, model_name)\n # Also leverage available_models list from EmbeddingsWithModels\n # Handle duplicate identifiers by creating combined keys\n embedding_by_model = {}\n identifier_conflicts = {} # Track which identifiers have conflicts\n\n for idx, emb_obj in enumerate(embeddings_list):\n # Get all possible identifiers for this embedding\n identifiers = []\n deployment = getattr(emb_obj, \"deployment\", None)\n model = getattr(emb_obj, \"model\", None)\n model_id = getattr(emb_obj, \"model_id\", None)\n model_name = getattr(emb_obj, \"model_name\", None)\n dimensions = getattr(emb_obj, \"dimensions\", None)\n available_models_attr = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Embedding object {idx}: deployment={deployment}, model={model}, \"\n f\"model_id={model_id}, model_name={model_name}, dimensions={dimensions}, \"\n f\"available_models={available_models_attr}\"\n )\n\n # If this embedding has available_models dict, map all models to their dedicated instances\n if available_models_attr and isinstance(available_models_attr, dict):\n logger.info(\n f\"Embedding object {idx} provides {len(available_models_attr)} models via available_models dict\"\n )\n for model_name_key, dedicated_embedding in available_models_attr.items():\n if model_name_key and str(model_name_key).strip():\n model_str = str(model_name_key).strip()\n if model_str not in embedding_by_model:\n # Use the dedicated embedding instance from the dict\n embedding_by_model[model_str] = dedicated_embedding\n logger.info(\n f\"Mapped available model '{model_str}' to dedicated embedding instance\"\n )\n else:\n # Conflict detected - track it\n if model_str not in identifier_conflicts:\n identifier_conflicts[model_str] = [embedding_by_model[model_str]]\n identifier_conflicts[model_str].append(dedicated_embedding)\n logger.warning(\n f\"Available model '{model_str}' has conflict - used by multiple embeddings\"\n )\n\n # Also map traditional identifiers (for backward compatibility)\n if deployment:\n identifiers.append(str(deployment))\n if model:\n identifiers.append(str(model))\n if model_id:\n identifiers.append(str(model_id))\n if model_name:\n identifiers.append(str(model_name))\n\n # Map all identifiers to this embedding object\n for identifier in identifiers:\n if identifier not in embedding_by_model:\n embedding_by_model[identifier] = emb_obj\n logger.info(f\"Mapped identifier '{identifier}' to embedding object {idx}\")\n else:\n # Conflict detected - track it\n if identifier not in identifier_conflicts:\n identifier_conflicts[identifier] = [embedding_by_model[identifier]]\n identifier_conflicts[identifier].append(emb_obj)\n logger.warning(\n f\"Identifier '{identifier}' has conflict - used by multiple embeddings\"\n )\n\n # For embeddings with model+deployment, create combined identifier\n # This helps when deployment is the same but model differs\n if deployment and model and deployment != model:\n combined_id = f\"{deployment}:{model}\"\n if combined_id not in embedding_by_model:\n embedding_by_model[combined_id] = emb_obj\n logger.info(\n f\"Created combined identifier '{combined_id}' for embedding object {idx}\"\n )\n\n # Log conflicts\n if identifier_conflicts:\n logger.warning(\n f\"Found {len(identifier_conflicts)} conflicting identifiers. \"\n f\"Consider using combined format 'deployment:model' or specifying unique model names.\"\n )\n for conflict_id, emb_list in identifier_conflicts.items():\n logger.warning(\n f\" Conflict on '{conflict_id}': {len(emb_list)} embeddings use this identifier\"\n )\n\n logger.info(f\"Generating embeddings for {len(available_spaces)} spaces in index\")\n logger.info(f\"Available embedding identifiers: {list(embedding_by_model.keys())}\")\n self.log(\n f\"[SEARCH] Embedding spaces detected in index: \"\n f\"{[space.space_id for space in available_spaces]}\"\n )\n self.log(f\"[SEARCH] Available embedding identifiers: {list(embedding_by_model.keys())}\")\n\n # Track matching status for debugging\n matched_models = []\n unmatched_models = []\n\n for space in available_spaces:\n model_name = space.space_id\n try:\n emb_obj = None\n\n # The OpenRAG proxy adapter can create an immutable child for\n # any provider-qualified or legacy route. Resolve through it\n # first so a legacy OpenAI vector is never queried through the\n # currently selected Azure/Watsonx provider by accident.\n for candidate in embeddings_list:\n resolver = getattr(candidate, \"for_model\", None)\n if callable(resolver):\n emb_obj = resolver(space.route_model)\n break\n\n # Backward compatibility for custom flows that still wire one\n # concrete embedding adapter per model.\n if emb_obj is None:\n for identifier in (space.space_id, space.field_identity):\n if identifier in embedding_by_model:\n emb_obj = embedding_by_model[identifier]\n break\n\n if emb_obj is not None:\n emb_deployment = getattr(emb_obj, \"deployment\", None)\n emb_model = getattr(emb_obj, \"model\", None)\n emb_model_id = getattr(emb_obj, \"model_id\", None)\n emb_dimensions = getattr(emb_obj, \"dimensions\", None)\n emb_available_models = getattr(emb_obj, \"available_models\", None)\n\n logger.info(\n f\"Using embedding object for model '{model_name}': \"\n f\"deployment={emb_deployment}, model={emb_model}, model_id={emb_model_id}, \"\n f\"dimensions={emb_dimensions}\"\n )\n\n # Check if this is a dedicated instance from available_models dict\n if emb_available_models and isinstance(emb_available_models, dict):\n logger.info(\n f\"Model '{model_name}' using dedicated instance from available_models dict \"\n f\"(pre-configured with correct model and dimensions)\"\n )\n\n # Use the embedding instance directly - no model switching needed!\n vec = emb_obj.embed_query(q)\n query_embeddings[space.field_identity] = vec\n matched_models.append(model_name)\n logger.info(\n f\"Generated embedding for model: {model_name} (actual dimensions: {len(vec)})\"\n )\n self.log(f\"[MATCH] Model '{model_name}' - generated {len(vec)}-dim embedding\")\n else:\n # No matching embedding found for this model\n unmatched_models.append(model_name)\n logger.warning(\n f\"No matching embedding found for model '{model_name}'. \"\n f\"This model will be skipped. Available identifiers: {list(embedding_by_model.keys())}\"\n )\n self.log(\n f\"[NO MATCH] Model '{model_name}' - available: {list(embedding_by_model.keys())}\"\n )\n except Exception as e: # provider clients expose different exception families\n logger.warning(f\"Failed to generate embedding for {model_name}: {e}\")\n self.log(f\"[ERROR] Embedding generation failed for '{model_name}': {e}\")\n\n # Log summary of model matching\n logger.info(\n f\"Model matching summary: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n self.log(\n f\"[SUMMARY] Model matching: {len(matched_models)} matched, {len(unmatched_models)} unmatched\"\n )\n if unmatched_models:\n self.log(f\"[WARN] Unmatched models in index: {unmatched_models}\")\n\n if not query_embeddings:\n self.log(\n \"[WARN] No indexed embedding space could be resolved; \"\n \"continuing with keyword-only search\"\n )\n\n index_properties = self._get_index_properties(client)\n # Build KNN queries for each model\n embedding_fields: list[str] = []\n knn_queries_with_candidates = []\n knn_queries_without_candidates = []\n\n raw_num_candidates = getattr(self, \"num_candidates\", 1000)\n try:\n num_candidates = int(raw_num_candidates) if raw_num_candidates is not None else 0\n except (TypeError, ValueError):\n num_candidates = 0\n use_num_candidates = num_candidates > 0\n\n for model_name, embedding_vector in query_embeddings.items():\n field_name = get_embedding_field_name(model_name)\n selected_field = field_name\n vector_dim = len(embedding_vector)\n\n # Only use the expected dynamic field - no legacy fallback\n # This prevents dimension mismatches between models\n if not self._is_knn_vector_field(index_properties, selected_field):\n logger.warning(\n f\"Skipping model {model_name}: field '{field_name}' is not mapped as knn_vector. \"\n f\"Documents must be indexed with this embedding model before querying.\"\n )\n self.log(\n f\"[SKIP] Field '{selected_field}' not a knn_vector - skipping model '{model_name}'\"\n )\n continue\n\n # Validate vector dimensions match the field dimensions\n field_dim = self._get_field_dimension(index_properties, selected_field)\n if field_dim is not None and field_dim != vector_dim:\n logger.error(\n f\"Dimension mismatch for model '{model_name}': \"\n f\"Query vector has {vector_dim} dimensions but field '{selected_field}' expects {field_dim}. \"\n f\"Skipping this model to prevent search errors.\"\n )\n self.log(\n f\"[DIM MISMATCH] Model '{model_name}': query={vector_dim} vs field={field_dim} - skipping\"\n )\n continue\n\n logger.info(\n f\"Adding KNN query for model '{model_name}': field='{selected_field}', \"\n f\"query_dims={vector_dim}, field_dims={field_dim or 'unknown'}\"\n )\n embedding_fields.append(selected_field)\n base_query = {\n \"knn\": {\n selected_field: {\n \"vector\": embedding_vector,\n \"k\": 50,\n }\n }\n }\n # Filter inside each kNN clause so the nearest neighbours are chosen\n # from matching documents; the outer bool filter alone would only\n # post-filter the global top-k and miss narrow ranges.\n if filter_clauses:\n base_query[\"knn\"][selected_field][\"filter\"] = {\"bool\": {\"filter\": filter_clauses}}\n\n if use_num_candidates:\n query_with_candidates = copy.deepcopy(base_query)\n query_with_candidates[\"knn\"][selected_field][\"num_candidates\"] = num_candidates\n else:\n query_with_candidates = base_query\n\n knn_queries_with_candidates.append(query_with_candidates)\n knn_queries_without_candidates.append(base_query)\n\n if not knn_queries_with_candidates:\n logger.warning(\n \"No valid knn_vector fields found for embedding models. \"\n \"This may indicate an empty index or missing field mappings. \"\n \"Continuing with keyword-only search.\"\n )\n self.log(\n f\"[WARN] No valid KNN queries could be built. \"\n f\"Query embeddings generated: {list(query_embeddings.keys())}, \"\n f\"but no matching knn_vector fields found in index; using keyword search.\"\n )\n\n # Get limit and score threshold\n limit = self._resolve_limit(filter_obj, default_limit=self.number_of_results)\n score_threshold = self._resolve_score_threshold(filter_obj)\n\n # Determine the best aggregation field for filename based on index mapping\n filename_agg_field = self._get_filename_agg_field(index_properties)\n\n keyword_query = {\n \"multi_match\": {\n \"query\": q,\n \"fields\": [\"text^2\", \"filename^1.5\"],\n \"type\": \"best_fields\",\n \"fuzziness\": \"AUTO\",\n \"boost\": 0.3 if knn_queries_with_candidates else 1.0,\n }\n }\n should_queries: list[dict[str, Any]] = [keyword_query]\n if knn_queries_with_candidates:\n should_queries.insert(\n 0,\n {\n \"dis_max\": {\n \"tie_breaker\": 0.0,\n \"boost\": 0.7,\n \"queries\": knn_queries_with_candidates,\n }\n },\n )\n\n # Keyword search intentionally has no embedding-field existence filter:\n # documents in an unresolved legacy space must remain discoverable.\n body = {\n \"query\": {\n \"bool\": {\n \"should\": should_queries,\n \"minimum_should_match\": 1,\n \"filter\": filter_clauses,\n }\n },\n \"aggs\": {\n \"data_sources\": {\"terms\": {\"field\": filename_agg_field, \"size\": 20}},\n \"document_types\": {\"terms\": {\"field\": \"mimetype\", \"size\": 10}},\n \"owners\": {\"terms\": {\"field\": \"owner\", \"size\": 10}},\n \"embedding_models\": {\"terms\": {\"field\": \"embedding_model\", \"size\": 10}},\n },\n \"_source\": [\n \"filename\",\n \"mimetype\",\n \"page\",\n \"text\",\n \"source_url\",\n \"owner\",\n \"embedding_model\",\n \"embedding_provider\",\n \"embedding_space_id\",\n \"parser\",\n \"chunk_size\",\n \"chunk_overlap\",\n \"allowed_users\",\n \"allowed_groups\",\n \"allowed_principals\",\n ],\n \"size\": limit,\n }\n\n if score_threshold is not None:\n body[\"min_score\"] = score_threshold\n\n logger.info(\n f\"Executing hybrid search with {len(knn_queries_with_candidates)} resolved embedding spaces: \"\n f\"{list(query_embeddings.keys())}\"\n )\n self.log(\n f\"[EXEC] Executing search with {len(knn_queries_with_candidates)} KNN queries, limit={limit}\"\n )\n self.log(f\"[EXEC] Embedding models used: {list(query_embeddings.keys())}\")\n self.log(f\"[EXEC] KNN fields being queried: {embedding_fields}\")\n\n try:\n resp = client.search(index=self.index_name, body=body, params={\"terminate_after\": 0})\n except RequestError as e:\n error_message = str(e)\n lowered = error_message.lower()\n if knn_queries_with_candidates and use_num_candidates and \"num_candidates\" in lowered:\n logger.warning(\n \"Retrying search without num_candidates parameter due to cluster capabilities\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n try:\n fallback_body[\"query\"][\"bool\"][\"should\"][0][\"dis_max\"][\"queries\"] = (\n knn_queries_without_candidates\n )\n except (KeyError, IndexError, TypeError) as inner_err:\n raise e from inner_err\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n elif knn_queries_with_candidates and (\n \"knn_vector\" in lowered or (\"field\" in lowered and \"knn\" in lowered)\n ):\n logger.warning(\n \"KNN search failed for an indexed embedding field; retrying keyword-only\",\n error=error_message,\n )\n fallback_body = copy.deepcopy(body)\n fallback_body[\"query\"][\"bool\"][\"should\"] = [keyword_query]\n resp = client.search(\n index=self.index_name,\n body=fallback_body,\n params={\"terminate_after\": 0},\n )\n else:\n raise\n hits = resp.get(\"hits\", {}).get(\"hits\", [])\n\n logger.info(f\"Found {len(hits)} results\")\n self.log(f\"[RESULT] Search complete: {len(hits)} results found\")\n\n if len(hits) == 0:\n self.log(\n f\"[EMPTY] Debug info: \"\n f\"spaces_in_index={[space.space_id for space in available_spaces]}, \"\n f\"matched_models={matched_models}, \"\n f\"knn_fields={embedding_fields}, \"\n f\"filters={len(filter_clauses)} clauses\"\n )\n\n return [\n {\n \"page_content\": fence_untrusted_text(hit[\"_source\"].get(\"text\", \"\")),\n \"metadata\": {\n **{k: v for k, v in hit[\"_source\"].items() if k != \"text\"},\n \"chunk_id\": hit.get(\"_id\"),\n \"id\": hit.get(\"_id\"),\n \"score\": hit.get(\"_score\"),\n },\n \"score\": hit.get(\"_score\"),\n }\n for hit in hits\n ]\n\n def search_documents(self) -> list[Data]:\n \"\"\"Search documents and return results as a list of Data objects.\n\n This is the main interface method that performs the multi-model search using the\n configured search_query and returns results in a format that downstream Parser\n components and tool wrappers can consume directly without DataFrame string truncation.\n\n Always builds the vector store (triggering ingestion if needed), then performs\n search only if a query is provided.\n\n Returns:\n List containing search results with text and metadata\n\n Raises:\n Exception: If search operation fails\n \"\"\"\n\n try:\n # Always build/cache the vector store to ensure ingestion happens\n logger.info(f\"Search query: {self.search_query}\")\n if self._cached_vector_store is None:\n self.build_vector_store()\n\n # Only perform search if query is provided\n search_query = (self.search_query or \"\").strip()\n if not search_query:\n self.log(\"No search query provided - ingestion completed, returning empty results\")\n\n return []\n\n # Perform search with the provided query\n raw = self.search(search_query)\n raw_list = [Data(text=hit[\"page_content\"], **hit[\"metadata\"]) for hit in raw]\n return raw_list\n\n except Exception as e:\n self.log(f\"search_documents error: {e}\")\n raise\n\n # -------- dynamic UI handling (auth switch) --------\n async def update_build_config(\n self, build_config: dict, field_value: str, field_name: str | None = None\n ) -> dict:\n \"\"\"Dynamically update component configuration based on field changes.\n\n This method handles real-time UI updates, particularly for authentication\n mode changes that show/hide relevant input fields.\n\n Args:\n build_config: Current component configuration\n field_value: New value for the changed field\n field_name: Name of the field that changed\n\n Returns:\n Updated build configuration with appropriate field visibility\n \"\"\"\n try:\n if field_name == \"auth_mode\":\n mode = (field_value or \"basic\").strip().lower()\n is_basic = mode == \"basic\"\n is_jwt = mode == \"jwt\"\n is_openrag = mode == \"openrag\"\n\n build_config[\"username\"][\"show\"] = is_basic\n build_config[\"password\"][\"show\"] = is_basic\n\n build_config[\"jwt_token\"][\"show\"] = is_jwt\n build_config[\"jwt_header\"][\"show\"] = is_jwt\n build_config[\"bearer_prefix\"][\"show\"] = is_jwt\n\n build_config[\"username\"][\"required\"] = is_basic\n build_config[\"password\"][\"required\"] = is_basic\n\n # build_config[\"jwt_token\"][\"required\"] = is_jwt\n build_config[\"jwt_header\"][\"required\"] = is_jwt\n build_config[\"bearer_prefix\"][\"required\"] = False\n\n # In 'openrag' mode, expose the OPENRAG_* fields up front\n # since they are the only credentials required.\n for openrag_field in (\n \"openrag_ingest_url\",\n \"openrag_ingest_token\",\n \"openrag_ingest_run_id\",\n \"openrag_ingest_batch_size\",\n ):\n if openrag_field in build_config:\n build_config[openrag_field][\"advanced\"] = not is_openrag\n build_config[openrag_field][\"required\"] = (\n is_openrag and openrag_field != \"openrag_ingest_batch_size\"\n )\n\n if is_basic or is_openrag:\n build_config[\"jwt_token\"][\"value\"] = \"\"\n\n return build_config\n\n except (KeyError, ValueError) as e:\n self.log(f\"update_build_config error: {e}\")\n\n return build_config\n" }, "docs_metadata": { "_input_type": "TableInput", diff --git a/frontend/app/api/queries/useGetSearchQuery.test.tsx b/frontend/app/api/queries/useGetSearchQuery.test.tsx index d29fadb5b0..42780baf99 100644 --- a/frontend/app/api/queries/useGetSearchQuery.test.tsx +++ b/frontend/app/api/queries/useGetSearchQuery.test.tsx @@ -61,10 +61,17 @@ function serveSearch(body: Record, status = 200) { return seen; } -async function runSearch(query: string, data?: ParsedQueryData | null) { - const rendered = renderHook(() => useGetSearchQuery(query, data), { - wrapper: createQueryWrapper(), - }); +async function runSearch( + query: string, + data?: ParsedQueryData | null, + dateRange?: { createdAfter?: string; createdBefore?: string }, +) { + const rendered = renderHook( + () => useGetSearchQuery(query, data, undefined, dateRange), + { + wrapper: createQueryWrapper(), + }, + ); await waitFor(() => expect( rendered.result.current.isSuccess || rendered.result.current.isError, @@ -389,4 +396,41 @@ describe("useGetSearchQuery", () => { expect(calls).toBe(1); }); }); + + describe("date bounds", () => { + const AFTER = "2026-10-07T07:00:00.000Z"; + const BEFORE = "2026-10-11T06:59:59.999999Z"; + const saved = () => + queryData({ + filters: { + data_sources: [], + document_types: ["application/pdf"], + owners: [], + connector_types: [], + created_after: AFTER, + created_before: BEFORE, + }, + }); + + it("sends the effective bounds beside the PDF filter", async () => { + const seen = serveSearch({ results: [] }); + await runSearch("term", saved(), { + createdAfter: AFTER, + createdBefore: BEFORE, + }); + expect(seen.payload?.filters).toEqual({ + document_types: ["application/pdf"], + }); + expect(seen.payload?.created_after).toBe(AFTER); + expect(seen.payload?.created_before).toBe(BEFORE); + }); + + it("does not fall back to saved dates when the range was cleared", async () => { + const seen = serveSearch({ results: [] }); + await runSearch("term", saved(), {}); + expect(seen.payload?.created_after).toBeUndefined(); + expect(seen.payload?.created_before).toBeUndefined(); + expect(JSON.stringify(seen.payload)).not.toContain("created_"); + }); + }); }); diff --git a/frontend/app/api/queries/useGetSearchQuery.ts b/frontend/app/api/queries/useGetSearchQuery.ts index da4a0dd2e7..1d06382298 100644 --- a/frontend/app/api/queries/useGetSearchQuery.ts +++ b/frontend/app/api/queries/useGetSearchQuery.ts @@ -18,6 +18,8 @@ export interface SearchPayload { owners?: string[]; connector_types?: string[]; }; + created_after?: string; + created_before?: string; } export interface ChunkResult { @@ -46,6 +48,7 @@ export interface ChunkResult { index?: number; allowed_users?: string[]; allowed_groups?: string[]; + indexed_time?: string; } export interface File { @@ -82,6 +85,7 @@ export interface File { chunks?: ChunkResult[]; allowed_users?: string[]; allowed_groups?: string[]; + indexed_time?: string; } // Non-fatal signal from the backend — e.g. an embedding provider was removed @@ -124,6 +128,9 @@ export const useGetSearchQuery = ( UseQueryOptions, "queryKey" | "queryFn" >, + // Effective UTC bounds from the caller. Saved dates inside `queryData` are + // deliberately not sent so a cleared temporary range cannot fall back to them. + dateRange?: { createdAfter?: string; createdBefore?: string }, ) => { const queryClient = useQueryClient(); @@ -159,7 +166,15 @@ export const useGetSearchQuery = ( }; if (queryData?.filters) { searchPayload.filters = - buildSearchPayloadFilters(queryData.filters) ?? undefined; + buildSearchPayloadFilters(queryData.filters, { + includeDates: false, + }) ?? undefined; + } + if (dateRange?.createdAfter) { + searchPayload.created_after = dateRange.createdAfter; + } + if (dateRange?.createdBefore) { + searchPayload.created_before = dateRange.createdBefore; } const response = await apiClient.post(`/search`, searchPayload); @@ -193,6 +208,7 @@ export const useGetSearchQuery = ( embedding_dimensions?: number; allowed_users?: string[]; allowed_groups?: string[]; + indexed_time?: string; } >(); @@ -233,6 +249,7 @@ export const useGetSearchQuery = ( embedding_dimensions: chunk.embedding_dimensions, allowed_users: chunk.allowed_users || [], allowed_groups: chunk.allowed_groups || [], + indexed_time: chunk.indexed_time, }); } }); @@ -253,6 +270,7 @@ export const useGetSearchQuery = ( chunks: file.chunks, allowed_users: file.allowed_users || [], allowed_groups: file.allowed_groups || [], + indexed_time: file.indexed_time, })); const warnings: SearchWarning[] = Array.isArray(data.warnings) @@ -272,7 +290,7 @@ export const useGetSearchQuery = ( return useQuery( { - queryKey: ["search", queryData, query], + queryKey: ["search", queryData, query, dateRange], placeholderData: (prev) => prev, staleTime: 0, queryFn: getFiles, diff --git a/frontend/app/api/queries/useListFiles.test.ts b/frontend/app/api/queries/useListFiles.test.ts new file mode 100644 index 0000000000..bad4b6262f --- /dev/null +++ b/frontend/app/api/queries/useListFiles.test.ts @@ -0,0 +1,162 @@ +/** + * Tests for useListFiles. + * + * Verifies URL parameter serialization (including the new createdAfter / + * createdBefore date-range params), response parsing, and error handling. + */ +import { renderHook, waitFor } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { describe, expect, it } from "vitest"; +import { server } from "@/test-utils/msw/server"; +import { createQueryWrapper } from "@/test-utils/render"; +import { useListFiles } from "./useListFiles"; + +function makeFileResponse(overrides: Record = {}) { + return { + files: [], + total: 0, + is_approximate: false, + page: 1, + page_size: 25, + after_key: null, + ...overrides, + }; +} + +describe("useListFiles", () => { + it("sends created_after param when createdAfter is provided", async () => { + let capturedUrl: string | undefined; + server.use( + http.get("/api/files", ({ request }) => { + capturedUrl = request.url; + return HttpResponse.json(makeFileResponse()); + }), + ); + + const { result } = renderHook( + () => + useListFiles({ + createdAfter: "2024-01-01T00:00:00.000Z", + }), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + const url = new URL(capturedUrl!); + expect(url.searchParams.get("created_after")).toBe( + "2024-01-01T00:00:00.000Z", + ); + }); + + it("sends created_before param when createdBefore is provided", async () => { + let capturedUrl: string | undefined; + server.use( + http.get("/api/files", ({ request }) => { + capturedUrl = request.url; + return HttpResponse.json(makeFileResponse()); + }), + ); + + const { result } = renderHook( + () => + useListFiles({ + createdBefore: "2024-01-31T23:59:59.999Z", + }), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + const url = new URL(capturedUrl!); + expect(url.searchParams.get("created_before")).toBe( + "2024-01-31T23:59:59.999Z", + ); + }); + + it("sends both date params together", async () => { + let capturedUrl: string | undefined; + server.use( + http.get("/api/files", ({ request }) => { + capturedUrl = request.url; + return HttpResponse.json(makeFileResponse()); + }), + ); + + const { result } = renderHook( + () => + useListFiles({ + createdAfter: "2024-03-01T00:00:00.000Z", + createdBefore: "2024-03-31T23:59:59.999Z", + }), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + const url = new URL(capturedUrl!); + expect(url.searchParams.get("created_after")).toBe( + "2024-03-01T00:00:00.000Z", + ); + expect(url.searchParams.get("created_before")).toBe( + "2024-03-31T23:59:59.999Z", + ); + }); + + it("omits date params when not provided", async () => { + let capturedUrl: string | undefined; + server.use( + http.get("/api/files", ({ request }) => { + capturedUrl = request.url; + return HttpResponse.json(makeFileResponse()); + }), + ); + + const { result } = renderHook(() => useListFiles({}), { + wrapper: createQueryWrapper(), + }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + const url = new URL(capturedUrl!); + expect(url.searchParams.has("created_after")).toBe(false); + expect(url.searchParams.has("created_before")).toBe(false); + }); + + it("maps indexed_time from the response", async () => { + server.use( + http.get("/api/files", () => + HttpResponse.json( + makeFileResponse({ + files: [ + { + filename: "doc.pdf", + mimetype: "application/pdf", + chunk_count: 5, + source_url: "", + owner: "user1", + owner_name: "User One", + owner_email: "user@example.com", + file_size: 1024, + connector_type: "local", + allowed_users: [], + allowed_groups: [], + indexed_time: "2024-06-15T10:30:00.000Z", + }, + ], + }), + ), + ), + ); + + const { result } = renderHook(() => useListFiles({}), { + wrapper: createQueryWrapper(), + }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(result.current.data?.files[0]?.indexed_time).toBe( + "2024-06-15T10:30:00.000Z", + ); + }); +}); diff --git a/frontend/app/api/queries/useListFiles.ts b/frontend/app/api/queries/useListFiles.ts index a14306a104..8be1428122 100644 --- a/frontend/app/api/queries/useListFiles.ts +++ b/frontend/app/api/queries/useListFiles.ts @@ -17,6 +17,8 @@ export interface ListFilesParams { dataSources?: string[]; search?: string; afterKey?: Record | null; //composite pagination cursor, could be undefined in the beginning + createdAfter?: string; + createdBefore?: string; } export interface ListFilesResponse { @@ -50,6 +52,10 @@ export const useListFiles = ( if (params.search) searchParams.set("search", params.search); if (params.afterKey) searchParams.set("after_key", JSON.stringify(params.afterKey)); + if (params.createdAfter) + searchParams.set("created_after", params.createdAfter); + if (params.createdBefore) + searchParams.set("created_before", params.createdBefore); const url = `/files?${searchParams.toString()}`; //internal (cookie auth) @@ -79,6 +85,7 @@ export const useListFiles = ( embedding_dimensions: f.embedding_dimensions as number | undefined, allowed_users: (f.allowed_users as string[]) || [], allowed_groups: (f.allowed_groups as string[]) || [], + indexed_time: (f.indexed_time as string) || undefined, status: "active" as const, }), ); diff --git a/frontend/app/api/queries/useUploadHistogram.test.ts b/frontend/app/api/queries/useUploadHistogram.test.ts new file mode 100644 index 0000000000..99a9e0b17d --- /dev/null +++ b/frontend/app/api/queries/useUploadHistogram.test.ts @@ -0,0 +1,73 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { describe, expect, it } from "vitest"; +import { server } from "@/test-utils/msw/server"; +import { createQueryWrapper } from "@/test-utils/render"; +import { useUploadHistogram } from "./useUploadHistogram"; + +const baseParams = { + start: "2024-05-01T00:00:00.000Z", + end: "2024-06-20T23:59:59.999Z", + tz: "America/New_York", +}; + +describe("useUploadHistogram", () => { + it("sends the window, time zone and filters", async () => { + let capturedUrl: string | undefined; + server.use( + http.get("/api/files/upload-histogram", ({ request }) => { + capturedUrl = request.url; + return HttpResponse.json({ total: 3, days: { "2024-06-01": 3 } }); + }), + ); + + const { result } = renderHook( + () => + useUploadHistogram({ + ...baseParams, + connectorType: ["local"], + mimetype: ["application/pdf"], + owner: ["u1", "u2"], + dataSources: ["a.pdf"], + }), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + const url = new URL(capturedUrl!); + expect(url.searchParams.get("start")).toBe(baseParams.start); + expect(url.searchParams.get("end")).toBe(baseParams.end); + expect(url.searchParams.get("tz")).toBe("America/New_York"); + expect(url.searchParams.getAll("connector_type")).toEqual(["local"]); + expect(url.searchParams.getAll("mimetype")).toEqual(["application/pdf"]); + expect(url.searchParams.getAll("owner")).toEqual(["u1", "u2"]); + expect(url.searchParams.getAll("data_sources")).toEqual(["a.pdf"]); + expect(result.current.data).toEqual({ + total: 3, + days: { "2024-06-01": 3 }, + }); + }); + + it("defaults missing fields", async () => { + server.use( + http.get("/api/files/upload-histogram", () => HttpResponse.json({})), + ); + const { result } = renderHook(() => useUploadHistogram(baseParams), { + wrapper: createQueryWrapper(), + }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data).toEqual({ total: 0, days: {} }); + }); + + it("surfaces server errors", async () => { + server.use( + http.get("/api/files/upload-histogram", () => + HttpResponse.json({ error: "boom" }, { status: 500 }), + ), + ); + const { result } = renderHook(() => useUploadHistogram(baseParams), { + wrapper: createQueryWrapper(), + }); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); +}); diff --git a/frontend/app/api/queries/useUploadHistogram.ts b/frontend/app/api/queries/useUploadHistogram.ts new file mode 100644 index 0000000000..e76b46b848 --- /dev/null +++ b/frontend/app/api/queries/useUploadHistogram.ts @@ -0,0 +1,69 @@ +import { + type UseQueryOptions, + useQuery, + useQueryClient, +} from "@tanstack/react-query"; +import { apiClient } from "@/lib/api-client"; + +export interface UploadHistogramParams { + start: string; + end: string; + tz: string; + connectorType?: string[]; + mimetype?: string[]; + owner?: string[]; + dataSources?: string[]; +} + +export interface UploadHistogram { + /** Distinct files across all time, under the same non-date filters. */ + total: number; + /** Distinct files uploaded per local day, keyed yyyy-MM-dd. */ + days: Record; +} + +export const useUploadHistogram = ( + params: UploadHistogramParams, + options?: Omit, "queryKey" | "queryFn">, +) => { + const queryClient = useQueryClient(); + + async function fetchHistogram(): Promise { + const searchParams = new URLSearchParams({ + start: params.start, + end: params.end, + tz: params.tz, + }); + for (const v of params.connectorType ?? []) + searchParams.append("connector_type", v); + for (const v of params.mimetype ?? []) searchParams.append("mimetype", v); + for (const v of params.owner ?? []) searchParams.append("owner", v); + for (const v of params.dataSources ?? []) + searchParams.append("data_sources", v); + + const response = await apiClient.get( + `/files/upload-histogram?${searchParams.toString()}`, + ); + if (response.status < 200 || response.status >= 300) { + throw new Error( + response.data?.error || + `Failed to load upload histogram: ${response.status}`, + ); + } + return { + total: response.data?.total ?? 0, + days: response.data?.days ?? {}, + }; + } + + return useQuery( + { + queryKey: ["uploadHistogram", params], + placeholderData: (prev: UploadHistogram | undefined) => prev, + queryFn: fetchHistogram, + retry: false, + ...options, + }, + queryClient, + ); +}; diff --git a/frontend/app/chat/_types/types.ts b/frontend/app/chat/_types/types.ts index fc475d039a..adc18cf1b1 100644 --- a/frontend/app/chat/_types/types.ts +++ b/frontend/app/chat/_types/types.ts @@ -129,6 +129,8 @@ export interface RequestBody { document_types?: string[]; owners?: string[]; connector_types?: string[]; + created_after?: string; + created_before?: string; }; filter_id?: string; limit?: number; diff --git a/frontend/app/knowledge/page-date-pagination.test.tsx b/frontend/app/knowledge/page-date-pagination.test.tsx new file mode 100644 index 0000000000..30bc360cf6 --- /dev/null +++ b/frontend/app/knowledge/page-date-pagination.test.tsx @@ -0,0 +1,217 @@ +/** + * Changing the effective date range must restart pagination: the next /files + * request carries the new bounds, page 1, and no cursor from the old window. + */ + +import { act, fireEvent } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { authPresets } from "@/test-utils/fixtures/auth"; +import { renderWithProviders, screen, waitFor } from "@/test-utils/render"; +import { setMockLocation } from "@/test-utils/router"; +import ProtectedSearchPage from "./page"; + +const OCT = { + after: "2026-10-07T07:00:00.000Z", + before: "2026-10-11T06:59:59.999999Z", +}; +const SEPT = { + after: "2026-09-01T07:00:00.000Z", + before: "2026-09-30T06:59:59.999999Z", +}; + +function savedFilter(id: string, range: { after: string; before: string }) { + return { + id, + name: id, + description: "", + owner: "u", + created_at: "", + updated_at: "", + query_data: JSON.stringify({ + query: "", + filters: { + data_sources: ["*"], + document_types: ["*"], + owners: ["*"], + connector_types: ["*"], + created_after: range.after, + created_before: range.before, + }, + limit: 10, + scoreThreshold: 0, + color: "zinc", + icon: "filter", + }), + }; +} + +const OCT_FILTER = savedFilter("oct", OCT); +const SEPT_FILTER = savedFilter("sept", SEPT); + +interface FilesRequest { + page: string | null; + afterKey: string | null; + createdAfter: string | null; + createdBefore: string | null; +} + +let requests: FilesRequest[] = []; + +function Controls() { + const { setSelectedFilter, setDateRange, clearFilter } = useKnowledgeFilter(); + return ( + <> + + + + + + + ); +} + +async function click(el: HTMLElement) { + await act(async () => { + fireEvent.click(el); + }); +} + +const last = () => requests[requests.length - 1]; + +async function renderPage() { + setMockLocation({ pathname: "/knowledge" }); + renderWithProviders( + <> + + + , + { + providers: "all", + auth: authPresets.admin, + handlers: [ + // 3 pages of 25; every page hands out a cursor for the next one. + http.get("/api/files", ({ request }) => { + const p = new URL(request.url).searchParams; + const page = Number(p.get("page") ?? "1"); + requests.push({ + page: p.get("page"), + afterKey: p.get("after_key"), + createdAfter: p.get("created_after"), + createdBefore: p.get("created_before"), + }); + return HttpResponse.json({ + files: [], + total: 75, + is_approximate: false, + page, + page_size: 25, + after_key: { cursor: `after-page-${page}` }, + }); + }), + http.post("/api/search", () => + HttpResponse.json({ files: [], warnings: [] }), + ), + http.get("/api/upload_options", () => HttpResponse.json({})), + http.get("/api/connectors", () => + HttpResponse.json({ connectors: [] }), + ), + http.get("/api/connectors/:name/defaults", () => HttpResponse.json({})), + ], + }, + ); + await waitFor(() => expect(requests.length).toBeGreaterThan(0)); +} + +/** Select a saved filter, then page forward to page 2 with its cursor in play. */ +async function reachPageTwo(filterButton: string) { + await click(screen.getByRole("button", { name: filterButton })); + await waitFor(() => expect(last()?.createdAfter).not.toBeNull()); + await waitFor(() => + expect(screen.getByRole("button", { name: "Next page" })).toBeEnabled(), + ); + await click(screen.getByRole("button", { name: "Next page" })); + await waitFor(() => expect(last().page).toBe("2")); + expect(last().afterKey).not.toBeNull(); + await screen.findByText(/Page 2 of 3/); +} + +// These render the whole page and page through it; the default 5s limit is +// too tight when the full suite runs in parallel. +vi.setConfig({ testTimeout: 20_000 }); + +describe("Knowledge page date range pagination", () => { + beforeEach(() => { + requests = []; + }); + + it("restarts at page 1 without a cursor when the range changes", async () => { + await renderPage(); + await reachPageTwo("select-oct"); + expect(last().createdAfter).toBe(OCT.after); + + await click(screen.getByRole("button", { name: "change-range" })); + + await waitFor(() => expect(last().createdAfter).not.toBe(OCT.after)); + await waitFor(() => expect(last().page).toBe("1")); + expect(last().afterKey).toBeNull(); + expect(last().createdAfter).toBeTruthy(); + expect(last().createdBefore).toBeTruthy(); + await screen.findByText(/Page 1 of 3/); + }); + + it("restarts at page 1 without a cursor, and without dates, when the range is cleared", async () => { + await renderPage(); + await reachPageTwo("select-oct"); + + await click(screen.getByRole("button", { name: "clear-range" })); + + await waitFor(() => expect(last().createdAfter).toBeNull()); + await waitFor(() => expect(last().page).toBe("1")); + expect(last().afterKey).toBeNull(); + // Cleared means cleared: neither bound falls back to the saved filter's. + expect(last().createdBefore).toBeNull(); + await screen.findByText(/Page 1 of 3/); + }); + + it("restarts at page 1 without a cursor when switching saved filters", async () => { + await renderPage(); + await reachPageTwo("select-oct"); + + await click(screen.getByRole("button", { name: "select-sept" })); + + await waitFor(() => expect(last().createdAfter).toBe(SEPT.after)); + await waitFor(() => expect(last().page).toBe("1")); + expect(last().afterKey).toBeNull(); + expect(last().createdBefore).toBe(SEPT.before); + await screen.findByText(/Page 1 of 3/); + }); + + it("restarts at page 1 without dates when the saved filter is cleared", async () => { + await renderPage(); + await reachPageTwo("select-oct"); + + await click(screen.getByRole("button", { name: "clear-filter" })); + + await waitFor(() => expect(last().createdAfter).toBeNull()); + await waitFor(() => expect(last().page).toBe("1")); + expect(last().afterKey).toBeNull(); + expect(last().createdBefore).toBeNull(); + }); +}); diff --git a/frontend/app/knowledge/page.tsx b/frontend/app/knowledge/page.tsx index 624db4f1bc..77c7d26306 100644 --- a/frontend/app/knowledge/page.tsx +++ b/frontend/app/knowledge/page.tsx @@ -14,6 +14,7 @@ import { AgGridReact, type CustomCellRendererProps } from "ag-grid-react"; import { AlertTriangle, Cloud, FileIcon, Globe, RefreshCw } from "lucide-react"; import { useRouter } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { KnowledgeDateRangeFilter } from "@/components/knowledge-date-range-filter"; import { KnowledgeDropdown } from "@/components/knowledge-dropdown"; import { ProtectedRoute } from "@/components/protected-route"; import { Banner, BannerIcon, BannerTitle } from "@/components/ui/banner"; @@ -128,6 +129,7 @@ const AG_FIELD_TO_SORT_BY: Record = { mimetype: "mimetype", owner: "owner", chunkCount: "chunk_count", + indexed_time: "indexed_time", status: "status", }; @@ -424,6 +426,8 @@ function SearchPage() { queryOverride, selectedFilter, setSelectedSources, + dateRange, + effectiveDateRange, } = useKnowledgeFilter(); const [selectedRows, setSelectedRows] = useState([]); @@ -608,13 +612,19 @@ function SearchPage() { const effectiveSearchText = queryOverride.trim() || parsedFilterData?.query?.trim() || ""; const hasActiveFilters = parsedFilterData?.filters - ? buildSearchPayloadFilters(parsedFilterData.filters) !== undefined + ? buildSearchPayloadFilters(parsedFilterData.filters, { + includeDates: false, + }) !== undefined : false; const isWildcardQuery = (effectiveSearchText === "" || effectiveSearchText === "*") && !hasActiveFilters; const filterPageResetKey = buildFilterPageResetKey(parsedFilterData); + // Retrieval uses the effective UTC bounds as-is; `dateRange` is display-only. + const createdAfter = effectiveDateRange.createdAfter; + const createdBefore = effectiveDateRange.createdBefore; + const { data: listFilesData, isLoading: isListFilesLoading, @@ -638,6 +648,8 @@ function SearchPage() { dataSources: listFilesFilterValues( parsedFilterData?.filters?.data_sources, ), + createdAfter, + createdBefore, }, { refetchInterval: 5000, @@ -650,9 +662,12 @@ function SearchPage() { isLoading: isSearchLoading, error: searchError, isError: isSearchError, - } = useGetSearchQuery(queryOverride, parsedFilterData, { - enabled: !isWildcardQuery, - }); + } = useGetSearchQuery( + queryOverride, + parsedFilterData, + { enabled: !isWildcardQuery }, + { createdAfter, createdBefore }, + ); const { files: searchFiles, warnings: searchWarnings } = searchData as SearchResult; @@ -737,6 +752,8 @@ function SearchPage() { effectiveData, taskFiles, Boolean(selectedFilter), + dateRange, + { createdAfter, createdBefore }, ); const serverTotal = isWildcardQuery @@ -748,7 +765,7 @@ function SearchPage() { useEffect(() => { cursorCacheRef.current = new Map(); setCurrentPage(1); - }, [effectiveSearchText, filterPageResetKey]); + }, [effectiveSearchText, filterPageResetKey, createdAfter, createdBefore]); // when the server responds with an after_key for page N, cache it as the cursor for page N+1 useEffect(() => { @@ -930,6 +947,26 @@ function SearchPage() { comparator: serverSideComparator, }; + const colIndexedTime: ColDef = { + field: "indexed_time", + headerName: "Date Ingested", + ...(isCloudBrand ? { flex: 1.2, minWidth: 140 } : { minWidth: 130 }), + sortable: true, + comparator: serverSideComparator, + cellClass: isCloudBrand ? "text-muted-foreground" : undefined, + valueFormatter: ({ value }: ValueFormatterParams) => { + if (!value) return "—"; + const d = new Date(value); + return Number.isNaN(d.getTime()) + ? "—" + : d.toLocaleDateString(undefined, { + year: "numeric", + month: "short", + day: "numeric", + }); + }, + }; + const colChunks: ColDef = { field: "chunkCount", headerName: "Chunks", @@ -1027,6 +1064,7 @@ function SearchPage() { colSize, colType, colOwner, + colIndexedTime, colChunks, colAvgScore, colStatus, @@ -1259,6 +1297,7 @@ function SearchPage() { Delete )} +
@@ -1325,12 +1364,25 @@ function SearchPage() { rowHeight={64} noRowsOverlayComponent={() => (
-
- No knowledge -
-
- Add files from local or your preferred cloud. -
+ {dateRange && serverTotal === 0 ? ( + <> +
+ No results for the selected date range +
+
+ Try a wider range or clear the date filter. +
+ + ) : ( + <> +
+ No knowledge +
+
+ Add files from local or your preferred cloud. +
+ + )}
)} /> @@ -1359,12 +1411,25 @@ function SearchPage() { onSortChanged={onSortChanged} noRowsOverlayComponent={() => (
-
- No knowledge -
-
- Add files from local or your preferred cloud. -
+ {dateRange && serverTotal === 0 ? ( + <> +
+ No results for the selected date range +
+
+ Try a wider range or clear the date filter. +
+ + ) : ( + <> +
+ No knowledge +
+
+ Add files from local or your preferred cloud. +
+ + )}
)} /> diff --git a/frontend/components/knowledge-date-range-filter.test.tsx b/frontend/components/knowledge-date-range-filter.test.tsx new file mode 100644 index 0000000000..f252042bc3 --- /dev/null +++ b/frontend/components/knowledge-date-range-filter.test.tsx @@ -0,0 +1,484 @@ +import { act, fireEvent } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { useEffect } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { renderWithProviders, screen, waitFor } from "@/test-utils/render"; +import { KnowledgeDateRangeFilter } from "./knowledge-date-range-filter"; + +let histogramRequests: URL[] = []; +let mockRunMode: string | null = "oss"; + +vi.mock("@/contexts/auth-context", async (importOriginal) => ({ + ...(await importOriginal()), + useAuth: () => ({ runMode: mockRunMode }), +})); + +vi.mock("@/contexts/brand-context", async (importOriginal) => ({ + ...(await importOriginal()), + useIsCloudBrand: () => false, +})); + +const histogramHandler = http.get( + "/api/files/upload-histogram", + ({ request }) => { + histogramRequests.push(new URL(request.url)); + return HttpResponse.json({ + total: 10, + days: { + "2024-06-20": 1, + "2024-06-19": 3, + "2024-06-10": 6, + "2024-05-15": 2, + }, + }); + }, +); + +// Radix Popover ignores userEvent pointer sequences in JSDOM; fireEvent.click opens it. +async function click(el: HTMLElement) { + await act(async () => { + fireEvent.click(el); + }); +} + +function renderFilter() { + renderWithProviders(, { + providers: ["knowledgeFilter"], + handlers: [histogramHandler], + }); +} + +async function openPopover() { + await click(screen.getByRole("button", { name: /date range/i })); + await waitFor(() => + expect(screen.getByRole("button", { name: /any time/i })).toHaveTextContent( + "10", + ), + ); +} + +function preset(name: RegExp) { + return screen.getByRole("button", { name }); +} + +function juneDay(day: number) { + const cell = screen + .getAllByRole("gridcell") + .find( + (c) => + c.tagName === "BUTTON" && + !c.className.includes("day-outside") && + c.querySelector("[data-day-number]")?.textContent === String(day), + ); + if (!cell) throw new Error(`June ${day} not found`); + return cell; +} + +const fmt = (d: Date) => + d.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + year: "numeric", + }); +const short = (d: Date) => + d.toLocaleDateString(undefined, { month: "short", day: "numeric" }); + +describe("KnowledgeDateRangeFilter", () => { + // Only Date is faked so act()/Radix timers keep running normally. + beforeEach(() => { + mockRunMode = "oss"; + histogramRequests = []; + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(new Date(2024, 5, 20, 12)); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + it.each(["saas", null])("is hidden when the run mode is %s", (runMode) => { + mockRunMode = runMode; + renderFilter(); + expect( + screen.queryByRole("button", { name: /filter by date range/i }), + ).not.toBeInTheDocument(); + expect(histogramRequests).toHaveLength(0); + }); + + it("renders the trigger with the default label", () => { + renderFilter(); + expect( + screen.getByRole("button", { name: /filter by date range/i }), + ).toBeInTheDocument(); + expect(screen.getByText("Date range")).toBeInTheDocument(); + }); + + it("only fetches upload counts while the popover is open", async () => { + renderFilter(); + expect(histogramRequests).toHaveLength(0); + await openPopover(); + expect(histogramRequests.length).toBeGreaterThan(0); + expect(histogramRequests[0].searchParams.get("tz")).toBeTruthy(); + }); + + it("recomputes today when the picker is reopened after midnight", async () => { + renderFilter(); + await openPopover(); + await click(preset(/^today/i)); + await click(screen.getByRole("button", { name: /apply/i })); + + act(() => { + vi.setSystemTime(new Date(2024, 5, 21, 0, 5)); + }); + await click( + screen.getByRole("button", { + name: new RegExp(fmt(new Date(2024, 5, 20))), + }), + ); + await waitFor(() => + expect(screen.getByRole("button", { name: /any time/i })).toBeVisible(), + ); + + expect(juneDay(21)).not.toBeDisabled(); + await click(preset(/^today/i)); + await click(screen.getByRole("button", { name: /apply/i })); + expect( + screen.getByRole("button", { + name: new RegExp(fmt(new Date(2024, 5, 21))), + }), + ).toBeInTheDocument(); + }); + + it("lists quick ranges with their upload counts", async () => { + renderFilter(); + await openPopover(); + expect(preset(/any time/i)).toHaveTextContent("10"); + expect(preset(/^today/i)).toHaveTextContent("1"); + expect(preset(/yesterday/i)).toHaveTextContent("3"); + expect(preset(/last 7 days/i)).toHaveTextContent("4"); + expect(preset(/last 30 days/i)).toHaveTextContent("10"); + expect(preset(/this month/i)).toHaveTextContent("10"); + expect(preset(/last month/i)).toHaveTextContent("2"); + expect(preset(/any time/i)).toHaveAttribute("aria-pressed", "true"); + }); + + it("marks each day with its upload density", async () => { + renderFilter(); + await openPopover(); + const density = (day: number) => + juneDay(day) + .querySelector("[data-testid='upload-density']") + ?.getAttribute("data-density"); + expect(density(20)).toBe("low"); + expect(density(19)).toBe("medium"); + expect(density(10)).toBe("high"); + expect(density(11)).toBeUndefined(); + }); + + it("previews a preset in the footer and applies it on Apply", async () => { + renderFilter(); + await openPopover(); + await click(preset(/last 7 days/i)); + + expect(preset(/last 7 days/i)).toHaveAttribute("aria-pressed", "true"); + expect( + screen.getByText( + `${short(new Date(2024, 5, 14))} – ${short(new Date(2024, 5, 20))}`, + ), + ).toBeInTheDocument(); + expect(screen.getByText("4 files")).toBeInTheDocument(); + // Nothing is applied until the user confirms. + expect( + screen.queryByRole("button", { name: /clear date range/i }), + ).not.toBeInTheDocument(); + + await click(screen.getByRole("button", { name: "Apply filter" })); + const expected = `${fmt(new Date(2024, 5, 14))} – ${fmt(new Date(2024, 5, 20))}`; + expect(screen.getByRole("button", { name: expected })).toBeInTheDocument(); + expect( + screen.queryByRole("button", { name: /last 7 days/i }), + ).not.toBeInTheDocument(); + }); + + it("discards the draft on Cancel", async () => { + renderFilter(); + await openPopover(); + await click(preset(/yesterday/i)); + await click(screen.getByRole("button", { name: "Cancel" })); + expect(screen.getByText("Date range")).toBeInTheDocument(); + expect( + screen.queryByRole("button", { name: /clear date range/i }), + ).not.toBeInTheDocument(); + }); + + it("selects a custom range by clicking two calendar days", async () => { + renderFilter(); + await openPopover(); + await click(juneDay(10)); + await click(juneDay(15)); + + expect( + screen.getByText( + `${short(new Date(2024, 5, 10))} – ${short(new Date(2024, 5, 15))}`, + ), + ).toBeInTheDocument(); + expect(screen.getByText("6 files")).toBeInTheDocument(); + expect( + screen.queryByRole("button", { pressed: true }), + ).not.toBeInTheDocument(); + + await click(screen.getByRole("button", { name: "Apply filter" })); + const expected = `${fmt(new Date(2024, 5, 10))} – ${fmt(new Date(2024, 5, 15))}`; + expect(screen.getByRole("button", { name: expected })).toHaveTextContent( + expected, + ); + }); + + it("reopens with the applied range selected", async () => { + renderFilter(); + await openPopover(); + await click(preset(/last month/i)); + await click(screen.getByRole("button", { name: "Apply filter" })); + await click( + screen.getByRole("button", { + name: `${fmt(new Date(2024, 4, 1))} – ${fmt(new Date(2024, 4, 31))}`, + }), + ); + expect(preset(/last month/i)).toHaveAttribute("aria-pressed", "true"); + expect(screen.getByText("May 2024")).toBeInTheDocument(); + }); + + it("clears an applied range from the external clear button", async () => { + renderFilter(); + await openPopover(); + await click(preset(/last 7 days/i)); + await click(screen.getByRole("button", { name: "Apply filter" })); + await click(screen.getByRole("button", { name: /clear date range/i })); + expect(screen.getByText("Date range")).toBeInTheDocument(); + }); + + it("clears an applied range by applying Any time", async () => { + renderFilter(); + await openPopover(); + await click(preset(/last 7 days/i)); + await click(screen.getByRole("button", { name: "Apply filter" })); + await click( + screen.getByRole("button", { + name: `${fmt(new Date(2024, 5, 14))} – ${fmt(new Date(2024, 5, 20))}`, + }), + ); + await click(preset(/any time/i)); + expect(screen.getByText("10 files")).toBeInTheDocument(); + await click(screen.getByRole("button", { name: "Apply filter" })); + expect(screen.getByText("Date range")).toBeInTheDocument(); + }); + + it("selects the last N days from a typed number and applies it", async () => { + renderFilter(); + await openPopover(); + await act(async () => { + fireEvent.change( + screen.getByRole("textbox", { name: "Number of days" }), + { + target: { value: "45" }, + }, + ); + }); + + expect( + screen.getByText( + `${short(new Date(2024, 4, 7))} – ${short(new Date(2024, 5, 20))}`, + ), + ).toBeInTheDocument(); + expect(screen.getByText("12 files")).toBeInTheDocument(); + + await click(screen.getByRole("button", { name: "Apply filter" })); + expect( + screen.getByRole("button", { + name: `${fmt(new Date(2024, 4, 7))} – ${fmt(new Date(2024, 5, 20))}`, + }), + ).toBeInTheDocument(); + }); + + it("ignores non-numeric input and clears the draft when emptied", async () => { + renderFilter(); + await openPopover(); + const input = screen.getByRole("textbox", { name: "Number of days" }); + await act(async () => { + fireEvent.change(input, { target: { value: "7a" } }); + }); + expect(input).toHaveValue("7"); + expect(preset(/last 7 days/i)).toHaveAttribute("aria-pressed", "true"); + + await act(async () => { + fireEvent.change(input, { target: { value: "" } }); + }); + expect(preset(/any time/i)).toHaveAttribute("aria-pressed", "true"); + }); + + it("rolls today over at midnight while the picker stays open", async () => { + // The rollover is a setTimeout armed on open, so fake it before rendering. + vi.useRealTimers(); + vi.useFakeTimers({ + toFake: ["Date", "setTimeout", "clearTimeout"], + shouldAdvanceTime: true, + }); + vi.setSystemTime(new Date(2024, 5, 20, 23, 59, 0)); + renderFilter(); + await click(screen.getByRole("button", { name: /date range/i })); + await waitFor(() => + expect( + screen.getByRole("button", { name: /any time/i }), + ).toHaveTextContent("10"), + ); + expect(juneDay(20)).not.toBeDisabled(); + expect( + screen + .getAllByRole("gridcell") + .some( + (c) => c.querySelector("[data-day-number]")?.textContent === "21", + ), + ).toBe(true); + // Before midnight the 21st is in the future and disabled. + expect(juneDay(21)).toBeDisabled(); + + await act(async () => { + vi.setSystemTime(new Date(2024, 5, 21, 0, 0, 5)); + await vi.advanceTimersByTimeAsync(2 * 60 * 1000); + }); + + expect(juneDay(21)).not.toBeDisabled(); + }); + + describe("saved ranges", () => { + const LA = "America/Los_Angeles"; + const originalTZ = process.env.TZ; + + // Pin a viewer zone that differs from the saved one, whatever the host's is. + beforeEach(() => { + process.env.TZ = "Europe/London"; + }); + afterEach(() => { + if (originalTZ === undefined) delete process.env.TZ; + else process.env.TZ = originalTZ; + }); + + function savedFilter(filters: Record, context?: object) { + return { + id: "f1", + name: "f1", + description: "", + owner: "u", + created_at: "", + updated_at: "", + query_data: JSON.stringify({ + query: "", + filters: { + data_sources: ["*"], + document_types: ["*"], + owners: ["*"], + connector_types: ["*"], + ...filters, + }, + ...(context && { date_range_context: context }), + }), + }; + } + + let effective: ReturnType["effectiveDateRange"]; + + function Probe({ filter }: { filter: ReturnType }) { + const ctx = useKnowledgeFilter(); + effective = ctx.effectiveDateRange; + // biome-ignore lint/correctness/useExhaustiveDependencies: select once + useEffect(() => ctx.setSelectedFilter(filter), []); + return ; + } + + function renderSaved(filter: ReturnType) { + renderWithProviders(, { + providers: ["knowledgeFilter"], + handlers: [histogramHandler], + }); + } + + async function applyUnchanged() { + await click(screen.getByRole("button", { name: /Jun|\d{4}/ })); + await click(screen.getByRole("button", { name: /apply filter/i })); + } + + it("keeps a created_after-only range when applied without edits", async () => { + const after = "2024-06-10T00:00:00.000Z"; + renderSaved(savedFilter({ created_after: after })); + await waitFor(() => expect(effective.createdAfter).toBe(after)); + + await applyUnchanged(); + + expect(effective.createdAfter).toBe(after); + expect(effective.createdBefore).toBeUndefined(); + expect(screen.getByRole("button", { name: /^From / })).toBeVisible(); + }); + + it("keeps a created_before-only range when applied without edits", async () => { + const before = "2024-06-15T23:59:59.999999Z"; + renderSaved(savedFilter({ created_before: before })); + await waitFor(() => expect(effective.createdBefore).toBe(before)); + + await applyUnchanged(); + + expect(effective.createdBefore).toBe(before); + expect(effective.createdAfter).toBeUndefined(); + expect(screen.getByRole("button", { name: /^Until / })).toBeVisible(); + }); + + it("counts an unchanged saved range in its saved timezone", async () => { + renderSaved( + savedFilter( + { + created_after: "2024-06-10T07:00:00.000Z", + created_before: "2024-06-14T06:59:59.999999Z", + }, + { from: "2024-06-10", to: "2024-06-13", timezone: LA }, + ), + ); + await waitFor(() => expect(effective.createdAfter).toBeTruthy()); + histogramRequests = []; + + await click(screen.getByRole("button", { name: /Jun/ })); + await waitFor(() => expect(histogramRequests.length).toBeGreaterThan(0)); + + const first = histogramRequests[0].searchParams; + expect(first.get("tz")).toBe(LA); + // The window reaches the saved bounds, so every counted day is in range. + expect(new Date(first.get("start") ?? "").getTime()).toBeLessThanOrEqual( + new Date("2024-06-10T07:00:00.000Z").getTime(), + ); + expect(new Date(first.get("end") ?? "").getTime()).toBeGreaterThanOrEqual( + new Date("2024-06-14T06:59:59.999999Z").getTime(), + ); + }); + + it("switches to the viewer's timezone once the range is edited", async () => { + renderSaved( + savedFilter( + { + created_after: "2024-06-10T07:00:00.000Z", + created_before: "2024-06-14T06:59:59.999999Z", + }, + { from: "2024-06-10", to: "2024-06-13", timezone: LA }, + ), + ); + await waitFor(() => expect(effective.createdAfter).toBeTruthy()); + await click(screen.getByRole("button", { name: /Jun/ })); + await waitFor(() => expect(histogramRequests.length).toBeGreaterThan(0)); + histogramRequests = []; + + await click(preset(/last 7 days/i)); + + await waitFor(() => expect(histogramRequests.length).toBeGreaterThan(0)); + expect(histogramRequests.at(-1)?.searchParams.get("tz")).toBe( + "Europe/London", + ); + }); + }); +}); diff --git a/frontend/components/knowledge-date-range-filter.tsx b/frontend/components/knowledge-date-range-filter.tsx new file mode 100644 index 0000000000..31f91b1ad4 --- /dev/null +++ b/frontend/components/knowledge-date-range-filter.tsx @@ -0,0 +1,626 @@ +"use client"; + +import { + addDays, + endOfMonth, + isAfter, + isSameDay, + min, + startOfDay, + startOfMonth, + subDays, + subMonths, +} from "date-fns"; +import { CalendarIcon, X } from "lucide-react"; +import { useEffect, useMemo, useState } from "react"; +import type { DateRange, DayContentProps } from "react-day-picker"; +import { useUploadHistogram } from "@/app/api/queries/useUploadHistogram"; +import { Button } from "@/components/ui/button"; +import { Calendar } from "@/components/ui/calendar"; +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover"; +import { useAuth } from "@/contexts/auth-context"; +import { useIsCloudBrand } from "@/contexts/brand-context"; +import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { isDateRangeFilterEnabled } from "@/lib/date-range-filter"; +import { + type DateRangeValue, + foreignTimeZone, + type SavedDateRange, + viewerTimeZone, +} from "@/lib/saved-date-range"; +import { + toLocalDayEndISOString, + toLocalDayKey, + toLocalDayStartISOString, +} from "@/lib/time-utils"; +import { cn } from "@/lib/utils"; + +type Range = { from: Date; to: Date }; + +const EMPTY_DAYS: Record = {}; +const MAX_CUSTOM_DAYS = 3650; + +interface Preset { + id: string; + label: string; + range: (today: Date) => Range | null; +} + +const PRESETS: Preset[] = [ + { id: "any", label: "Any time", range: () => null }, + { id: "today", label: "Today", range: (t) => ({ from: t, to: t }) }, + { + id: "yesterday", + label: "Yesterday", + range: (t) => ({ from: subDays(t, 1), to: subDays(t, 1) }), + }, + { + id: "last-7", + label: "Last 7 days", + range: (t) => ({ from: subDays(t, 6), to: t }), + }, + { + id: "last-30", + label: "Last 30 days", + range: (t) => ({ from: subDays(t, 29), to: t }), + }, + { + id: "this-month", + label: "This month", + range: (t) => ({ from: startOfMonth(t), to: t }), + }, + { + id: "last-month", + label: "Last month", + range: (t) => { + const prev = subMonths(t, 1); + return { from: startOfMonth(prev), to: endOfMonth(prev) }; + }, + }, +]; + +function formatDateLabel(date: Date): string { + return date.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + year: "numeric", + }); +} + +function formatShortLabel(date: Date, today: Date): string { + return date.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + ...(date.getFullYear() !== today.getFullYear() && { year: "numeric" }), + }); +} + +/** "A – B", a single day, or "From A" / "Until B" for a one-sided saved range. */ +function describeRange( + range: DateRangeValue | null, + format: (date: Date) => string, +): string | null { + if (range?.from && range.to) { + return isSameDay(range.from, range.to) + ? format(range.from) + : `${format(range.from)} – ${format(range.to)}`; + } + if (range?.from) return `From ${format(range.from)}`; + if (range?.to) return `Until ${format(range.to)}`; + return null; +} + +function earlierInstant(iso: string, other?: string): string { + return other && new Date(other) < new Date(iso) ? other : iso; +} + +function laterInstant(iso: string, other?: string): string { + return other && new Date(other) > new Date(iso) ? other : iso; +} + +function filterValues(values?: string[]) { + const filtered = values?.filter((value) => value !== "*"); + return filtered && filtered.length > 0 ? filtered : undefined; +} + +function countInRange(days: Record, { from, to }: Range) { + let total = 0; + for (let d = startOfDay(from); !isAfter(d, to); d = addDays(d, 1)) { + total += days[toLocalDayKey(d)] ?? 0; + } + return total; +} + +function DensityMark({ + count, + inverted, +}: { + count: number; + inverted: boolean; +}) { + if (count <= 0) return ; + return ( + = 5 ? "high" : count >= 2 ? "medium" : "low"} + className={cn( + "h-1 rounded-full", + count >= 5 ? "w-5" : count >= 2 ? "w-3" : "w-1", + inverted ? "bg-primary-foreground" : "bg-primary", + )} + /> + ); +} + +/** Toolbar picker: edits the temporary range of the Knowledge page. */ +export function KnowledgeDateRangeFilter() { + const { dateRange, setDateRange, parsedFilterData, effectiveDateRange } = + useKnowledgeFilter(); + const enabled = useDateRangeFilterEnabled(); + if (!enabled) return null; + return ( + + ); +} + +export function useDateRangeFilterEnabled(): boolean { + const { runMode } = useAuth(); + const isCloudBrand = useIsCloudBrand(); + return isDateRangeFilterEnabled(runMode, { isCloudBrand }); +} + +function useToday(open: boolean) { + // Refreshed while open so a page left running past midnight never keeps + // yesterday as "today". + const [today, setToday] = useState(() => startOfDay(new Date())); + + useEffect(() => { + if (!open) return; + const rollover = () => { + const now = startOfDay(new Date()); + setToday((prev) => (isSameDay(prev, now) ? prev : now)); + }; + const nextMidnight = addDays(startOfDay(new Date()), 1); + const timer = setTimeout( + rollover, + nextMidnight.getTime() - Date.now() + 1000, + ); + return () => clearTimeout(timer); + }, [open, today]); + + return [today, setToday] as const; +} + +function QuickRanges({ + today, + activePresetId, + presetCount, + customDays, + onPreset, + onCustomDays, +}: { + today: Date; + activePresetId: string | undefined; + presetCount: (range: Range | null) => number | undefined; + customDays: string; + onPreset: (preset: Preset) => void; + onCustomDays: (value: string) => void; +}) { + return ( +
+

+ Quick ranges +

+ {/* A scrolling row on narrow screens, a list beside the calendar otherwise. */} +
+ {PRESETS.map((preset) => { + const count = presetCount(preset.range(today)); + const active = preset.id === activePresetId; + return ( + + ); + })} +
+
+ + onCustomDays(e.target.value)} + className="h-8 w-16 rounded-md border bg-background px-2 text-sm tabular-nums" + /> + days +
+
+ ); +} + +function CalendarPanel({ + draft, + month, + today, + DayContent, + onSelect, + onMonthChange, +}: { + draft: DateRange | undefined; + month: Date; + today: Date; + DayContent: (props: DayContentProps) => React.JSX.Element; + onSelect: (next: DateRange | undefined) => void; + onMonthChange: (month: Date) => void; +}) { + return ( +
+ +
+ + 1 upload + + + + 2–4 + + + + 5+ + +
+
+ ); +} + +function PopoverFooter({ + footerCount, + footerRange, + onCancel, + onApply, +}: { + footerCount: number | undefined; + footerRange: string; + onCancel: () => void; + onApply: () => void; +}) { + return ( +
+

+ {footerCount !== undefined && ( + <> + + {footerCount} {footerCount === 1 ? "file" : "files"} + + {" · "} + + )} + {footerRange} +

+
+ + +
+
+ ); +} + +export interface DateRangePickerProps { + value: DateRangeValue | null; + onChange: (range: DateRangeValue | null) => void; + /** Facet selections that scope the upload-density marks (not the range itself). */ + facetFilters?: { + data_sources?: string[]; + document_types?: string[]; + owners?: string[]; + connector_types?: string[]; + }; + /** + * Bounds + timezone context behind `value`. While the user has not edited the + * range, upload counts use this timezone so they agree with retrieval. + */ + savedRange?: SavedDateRange; +} + +/** + * Controlled date-range picker. It owns only the popover draft; the caller + * decides what applying or clearing means (temporary toolbar range, or an + * unsaved draft inside the Knowledge Filter editor). + */ +export function DateRangePicker({ + value: dateRange, + onChange, + facetFilters, + savedRange, +}: DateRangePickerProps) { + const [open, setOpen] = useState(false); + const [today, setToday] = useToday(open); + const [draft, setDraft] = useState(); + const [month, setMonth] = useState(today); + const [customDays, setCustomDays] = useState(""); + // False until the user changes the draft; an unchanged Apply keeps the range as saved. + const [edited, setEdited] = useState(false); + const savedTimeZone = foreignTimeZone(savedRange); + const countInSavedZone = !edited && savedTimeZone !== null; + + // Cover the shown month, every preset (last month, or 29 days back when that + // is earlier), and the draft start so typed/custom ranges count correctly. + const windowStart = min([ + startOfMonth(month), + startOfMonth(subMonths(today, 1)), + subDays(today, 29), + ...(draft?.from ? [startOfDay(draft.from)] : []), + ]); + // An unedited saved range is counted in the zone it was saved in, over a + // window that includes its exact bounds; edits switch to the viewer's zone. + const windowStartISO = toLocalDayStartISOString(windowStart); + const windowEndISO = toLocalDayEndISOString(today); + const { data: histogram } = useUploadHistogram( + { + start: countInSavedZone + ? earlierInstant(windowStartISO, savedRange?.createdAfter) + : windowStartISO, + end: countInSavedZone + ? laterInstant(windowEndISO, savedRange?.createdBefore) + : windowEndISO, + tz: countInSavedZone && savedTimeZone ? savedTimeZone : viewerTimeZone(), + connectorType: filterValues(facetFilters?.connector_types), + mimetype: filterValues(facetFilters?.document_types), + owner: filterValues(facetFilters?.owners), + dataSources: filterValues(facetFilters?.data_sources), + }, + { enabled: open }, + ); + const days = histogram?.days ?? EMPTY_DAYS; + + const handleOpenChange = (next: boolean) => { + if (next) { + const freshToday = startOfDay(new Date()); + setToday(freshToday); + setDraft( + dateRange?.from || dateRange?.to + ? { from: dateRange.from, to: dateRange.to } + : undefined, + ); + setMonth(dateRange?.to ?? dateRange?.from ?? freshToday); + setCustomDays(""); + setEdited(false); + } + setOpen(next); + }; + + const draftRange: Range | null = draft?.from + ? { from: draft.from, to: draft.to ?? draft.from } + : null; + + const presetCount = (range: Range | null) => + histogram + ? range + ? countInRange(days, range) + : histogram.total + : undefined; + + // A saved range with only one bound is shown as such until it is edited. + const openEndedUnedited = + !edited && Boolean(dateRange) && !(dateRange?.from && dateRange?.to); + + const activePresetId = openEndedUnedited + ? undefined + : PRESETS.find((preset) => { + const presetRange = preset.range(today); + if (!presetRange || !draftRange) return !presetRange && !draftRange; + return ( + isSameDay(presetRange.from, draftRange.from) && + isSameDay(presetRange.to, draftRange.to) + ); + })?.id; + + const handlePreset = (preset: Preset) => { + const range = preset.range(today); + setEdited(true); + setDraft(range ?? undefined); + setCustomDays(""); + if (range) setMonth(range.to); + }; + + const handleCustomDays = (value: string) => { + const digits = value.replace(/\D/g, "").slice(0, 4); + setEdited(true); + setCustomDays(digits); + const n = Math.min(Number.parseInt(digits, 10), MAX_CUSTOM_DAYS); + if (!n) { + setDraft(undefined); + return; + } + const range = { from: subDays(today, n - 1), to: today }; + setDraft(range); + setMonth(range.to); + }; + + const handleCalendarSelect = (next: DateRange | undefined) => { + setEdited(true); + setCustomDays(""); + setDraft(next); + }; + + const handleApply = () => { + // Re-applying what was shown must not reinterpret it: an open-ended range + // has no second day to fill in, and the saved bounds stay as they are. + if (edited) onChange(draftRange); + setOpen(false); + }; + + const handleClear = () => { + onChange(null); + setOpen(false); + }; + + // Rebuilt only when the counts change, so day cells don't remount per render. + const DayContent = useMemo( + () => + function DayContent({ date, activeModifiers }: DayContentProps) { + const isEnd = + Boolean(activeModifiers.range_start) || + Boolean(activeModifiers.range_end); + return ( + + {date.getDate()} + + + ); + }, + [days], + ); + + const label = describeRange(dateRange, formatDateLabel); + + // A saved one-sided range has no closed draftRange; describe it as it is. + const footerCount = openEndedUnedited ? undefined : presetCount(draftRange); + const footerRange = openEndedUnedited + ? (describeRange(dateRange, (d) => formatShortLabel(d, today)) ?? + "Any time") + : draftRange + ? isSameDay(draftRange.from, draftRange.to) + ? formatShortLabel(draftRange.from, today) + : `${formatShortLabel(draftRange.from, today)} – ${formatShortLabel(draftRange.to, today)}` + : "Any time"; + + return ( +
+ + + + + +
+ + +
+ {savedTimeZone && ( +

+ Saved in {savedTimeZone}. Days you pick here use your time zone ( + {viewerTimeZone()}). +

+ )} + setOpen(false)} + onApply={handleApply} + /> + + + {dateRange && ( + + )} +
+ ); +} diff --git a/frontend/components/knowledge-filter-date-field.tsx b/frontend/components/knowledge-filter-date-field.tsx new file mode 100644 index 0000000000..24f9656c0a --- /dev/null +++ b/frontend/components/knowledge-filter-date-field.tsx @@ -0,0 +1,263 @@ +"use client"; + +import { isAfter, isSameDay, startOfDay, subDays } from "date-fns"; +import { CalendarIcon, Check, ChevronDown } from "lucide-react"; +import { useRef, useState } from "react"; +import type { Matcher } from "react-day-picker"; +import { Button } from "@/components/ui/button"; +import { Calendar } from "@/components/ui/calendar"; +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover"; +import { + type DateRangeValue, + describeRange, + viewerTimeZone, +} from "@/lib/saved-date-range"; +import { cn } from "@/lib/utils"; + +const PRESET_DAYS = [30, 60, 90] as const; +const DEFAULT_CUSTOM_DAYS = 30; + +type Field = "from" | "to"; + +function formatFieldDate(date: Date | undefined): string { + return date + ? date.toLocaleDateString(undefined, { + month: "2-digit", + day: "2-digit", + year: "2-digit", + }) + : "Select date"; +} + +type OptionId = "all" | `${number}` | "custom"; + +function activeOption(draft: DateRangeValue | null, today: Date): OptionId { + if (!draft?.from && !draft?.to) return "all"; + const preset = PRESET_DAYS.find( + (n) => + draft.from && + draft.to && + isSameDay(draft.to, today) && + isSameDay(draft.from, subDays(today, n - 1)), + ); + return preset ? (String(preset) as OptionId) : "custom"; +} + +export interface KnowledgeFilterDateFieldProps { + value: DateRangeValue | null; + onChange: (range: DateRangeValue | null) => void; + /** Timezone the shown days were saved in, when it differs from the viewer's. */ + savedTimeZone?: string | null; +} + +/** + * "Date ingested" control for the Knowledge Filter editor: a short list of + * ranges plus From/To fields that open a single-day calendar. It edits only + * an unsaved draft; the filter's own Save decides what is persisted. + */ +export function KnowledgeFilterDateField({ + value, + onChange, + savedTimeZone = null, +}: KnowledgeFilterDateFieldProps) { + const [open, setOpen] = useState(false); + const [today, setToday] = useState(() => startOfDay(new Date())); + const [draft, setDraft] = useState(value); + const [expanded, setExpanded] = useState(null); + // False until the user changes the draft; Apply then keeps the range as saved. + const edited = useRef(false); + + const handleOpenChange = (next: boolean) => { + if (next) { + setToday(startOfDay(new Date())); + setDraft(value); + setExpanded(null); + edited.current = false; + } + setOpen(next); + }; + + const update = (next: DateRangeValue | null) => { + edited.current = true; + setDraft(next); + }; + + const active = activeOption(draft, today); + + const options: { id: OptionId; label: string; select: () => void }[] = [ + { + id: "all", + label: "All dates", + select: () => { + update(null); + setExpanded(null); + }, + }, + ...PRESET_DAYS.map((n) => ({ + id: String(n) as OptionId, + label: `Last ${n} days`, + select: () => { + update({ from: subDays(today, n - 1), to: today }); + setExpanded(null); + }, + })), + { + id: "custom", + label: "Custom date range", + select: () => { + if (!draft?.from && !draft?.to) { + update({ + from: subDays(today, DEFAULT_CUSTOM_DAYS - 1), + to: today, + }); + } + setExpanded("from"); + }, + }, + ]; + + const pick = (field: Field, date: Date) => { + if (field === "from") { + // A start after the current end drags the end along, never inverts the range. + update({ + from: date, + to: draft?.to && isAfter(date, draft.to) ? date : draft?.to, + }); + } else { + update({ from: draft?.from, to: date }); + } + setExpanded(null); + }; + + const disabledDays = (field: Field): Matcher[] => { + const matchers: Matcher[] = [{ after: today }]; + if (field === "to" && draft?.from) { + matchers.push({ before: draft.from }); + } + return matchers; + }; + + const apply = () => { + if (edited.current) onChange(draft?.from || draft?.to ? draft : null); + setOpen(false); + }; + + const label = describeRange(value); + const fields: { field: Field; title: string; day: Date | undefined }[] = [ + { field: "from", title: "From", day: draft?.from }, + { field: "to", title: "To", day: draft?.to }, + ]; + + return ( + + + + + +
+ {options.map((option) => ( + + ))} +
+ +
+ {fields.map(({ field, title, day }) => { + const isOpen = expanded === field; + return ( +
+ {title} + + {isOpen && ( +
+ date && pick(field, date)} + /> +
+ )} +
+ ); + })} +
+ + {savedTimeZone && ( +

+ Saved in {savedTimeZone}. Days you pick here use your time zone ( + {viewerTimeZone()}). +

+ )} + +
+ + +
+ + + ); +} diff --git a/frontend/components/knowledge-filter-panel.test.tsx b/frontend/components/knowledge-filter-panel.test.tsx new file mode 100644 index 0000000000..09f5cc9dfb --- /dev/null +++ b/frontend/components/knowledge-filter-panel.test.tsx @@ -0,0 +1,313 @@ +import { act, fireEvent } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { useEffect } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { authPresets, withAuth } from "@/test-utils/fixtures/auth"; +import { server } from "@/test-utils/msw/server"; +import { renderWithProviders, screen, waitFor } from "@/test-utils/render"; +import { KnowledgeFilterPanel } from "./knowledge-filter-panel"; + +vi.mock("@/lib/analytics", () => ({ trackButton: vi.fn() })); + +// The real brand provider reads localStorage; the picker only needs this flag. +vi.mock("@/contexts/brand-context", async (importOriginal) => ({ + ...(await importOriginal()), + useIsCloudBrand: () => false, +})); + +const AFTER = "2026-10-07T07:00:00.000Z"; +const BEFORE = "2026-10-11T06:59:59.999999Z"; +const CONTEXT = { + from: "2026-10-07", + to: "2026-10-10", + timezone: "America/Los_Angeles", +}; + +const auth = withAuth(authPresets.admin, { + me: { run_mode: "oss" }, + usersMe: { permissions: ["kf:create", "kf:edit:own"] }, +}); + +type OneSided = { created_after?: string; created_before?: string }; + +function saved(withDates: boolean, oneSided?: OneSided) { + return { + id: "f1", + name: "PDFs", + description: "", + owner: "user-1", + created_at: "", + updated_at: "", + query_data: JSON.stringify({ + query: "", + filters: { + data_sources: ["*"], + document_types: ["application/pdf"], + owners: ["*"], + connector_types: ["*"], + ...(oneSided + ? oneSided + : withDates && { created_after: AFTER, created_before: BEFORE }), + }, + ...((withDates || oneSided) && { + date_range_context: oneSided ? { timezone: CONTEXT.timezone } : CONTEXT, + }), + limit: 10, + scoreThreshold: 0, + color: "zinc", + icon: "filter", + }), + }; +} + +function Harness({ + withDates, + create = false, + oneSided, +}: { + withDates: boolean; + create?: boolean; + oneSided?: OneSided; +}) { + const { setSelectedFilter, dateRange, setDateRange, startCreateMode } = + useKnowledgeFilter(); + // biome-ignore lint/correctness/useExhaustiveDependencies: select once on mount + useEffect(() => { + if (!create) setSelectedFilter(saved(withDates, oneSided)); + }, []); + return ( + <> + + + {dateRange ? "set" : "none"} + + + ); +} + +let putBody: { queryData?: string } | undefined; +let postBody: { name?: string; queryData?: string } | undefined; +// A getter keeps TypeScript from narrowing the callback-assigned variable to never. +const createdBody = () => postBody; + +async function renderPanel(withDates: boolean, oneSided?: OneSided) { + putBody = undefined; + server.use( + http.post("/api/search", () => HttpResponse.json({ aggregations: {} })), + http.put("/api/knowledge-filter/f1", async ({ request }) => { + putBody = (await request.json()) as { queryData?: string }; + return HttpResponse.json({ + success: true, + filter: saved(withDates, oneSided), + }); + }), + ); + renderWithProviders(, { + providers: ["auth", "knowledgeFilter"], + auth, + }); + await screen.findByText("Date ingested"); +} + +const savedQueryData = () => JSON.parse(putBody?.queryData ?? "{}"); + +async function click(el: HTMLElement) { + await act(async () => { + fireEvent.click(el); + }); +} + +async function save() { + await click(await screen.findByRole("button", { name: /update filter/i })); + await waitFor(() => expect(putBody).toBeDefined()); + return savedQueryData(); +} + +describe("KnowledgeFilterPanel date ingested", () => { + beforeEach(() => { + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(new Date(2026, 9, 23, 12)); + }); + afterEach(() => vi.useRealTimers()); + + it("saves an unchanged dated filter with its exact bounds and context", async () => { + await renderPanel(true); + const data = await save(); + expect(data.filters.created_after).toBe(AFTER); + expect(data.filters.created_before).toBe(BEFORE); + expect(data.filters.document_types).toEqual(["application/pdf"]); + expect(data.date_range_context).toEqual(CONTEXT); + expect(data.filters.date_range_context).toBeUndefined(); + }); + + it("saves a filter without dates unrestricted", async () => { + await renderPanel(false); + const data = await save(); + expect(data.filters).not.toHaveProperty("created_after"); + expect(data.filters).not.toHaveProperty("created_before"); + expect(data).not.toHaveProperty("date_range_context"); + }); + + it("clears a saved range by choosing All dates", async () => { + await renderPanel(true); + await click(screen.getByRole("button", { name: /–/ })); + await click(await screen.findByRole("button", { name: "All dates" })); + await click(screen.getByRole("button", { name: "Apply" })); + const data = await save(); + expect(data.filters).not.toHaveProperty("created_after"); + expect(data.filters).not.toHaveProperty("created_before"); + expect(data).not.toHaveProperty("date_range_context"); + }); + + it("does not save a temporary toolbar range", async () => { + await renderPanel(false); + await click(screen.getByRole("button", { name: "toolbar-range" })); + expect(screen.getByTestId("toolbar")).toHaveTextContent("set"); + const data = await save(); + expect(data.filters).not.toHaveProperty("created_after"); + expect(data).not.toHaveProperty("date_range_context"); + }); + + it("creates a filter without dates even when a toolbar range was active", async () => { + postBody = undefined; + server.use( + http.post("/api/search", () => HttpResponse.json({ aggregations: {} })), + http.post("/api/knowledge-filter", async ({ request }) => { + postBody = (await request.json()) as typeof postBody; + return HttpResponse.json({ success: true, filter: saved(false) }); + }), + ); + renderWithProviders(, { + providers: ["auth", "knowledgeFilter"], + auth, + }); + await click(screen.getByRole("button", { name: "toolbar-range" })); + expect(screen.getByTestId("toolbar")).toHaveTextContent("set"); + + await click(screen.getByRole("button", { name: "start-create" })); + await screen.findByText("Date ingested"); + expect(screen.getByTestId("toolbar")).toHaveTextContent("none"); + expect( + screen.getByRole("button", { name: /filter by date ingested/i }), + ).toBeInTheDocument(); + + fireEvent.change(screen.getByPlaceholderText(/name/i), { + target: { value: "New filter" }, + }); + await click(screen.getByRole("button", { name: /create filter/i })); + await waitFor(() => expect(createdBody()).toBeDefined()); + + const data = JSON.parse(createdBody()?.queryData ?? "{}"); + expect(data.filters).not.toHaveProperty("created_after"); + expect(data.filters).not.toHaveProperty("created_before"); + expect(data).not.toHaveProperty("date_range_context"); + }); + + it.each([ + ["created_after", { created_after: "2026-10-07T07:00:00.000Z" }], + ["created_before", { created_before: "2026-10-11T06:59:59.999999Z" }], + ])("keeps a %s-only range when the field is applied without edits", async (key, oneSided) => { + await renderPanel(true, oneSided); + await click(screen.getByRole("button", { name: /^(From|Until) / })); + await click(await screen.findByRole("button", { name: "Apply" })); + + const data = await save(); + expect(data.filters[key]).toBe(Object.values(oneSided)[0]); + const other = key === "created_after" ? "created_before" : "created_after"; + expect(data.filters).not.toHaveProperty(other); + }); + + it("offers All dates, 30/60/90 days and a custom range", async () => { + await renderPanel(false); + await click( + screen.getByRole("button", { name: /filter by date ingested/i }), + ); + for (const name of [ + "All dates", + "Last 30 days", + "Last 60 days", + "Last 90 days", + "Custom date range", + ]) { + expect(screen.getByRole("button", { name })).toBeInTheDocument(); + } + expect(screen.getByRole("button", { name: "All dates" })).toHaveAttribute( + "aria-pressed", + "true", + ); + }); + + it("saves a Last 30 days selection as full local-day bounds", async () => { + await renderPanel(false); + await click( + screen.getByRole("button", { name: /filter by date ingested/i }), + ); + await click(screen.getByRole("button", { name: "Last 30 days" })); + await click(screen.getByRole("button", { name: "Apply" })); + + const data = await save(); + // today is faked as 2026-10-23, so 30 days is Sep 24 - Oct 23 inclusive. + expect(new Date(data.filters.created_after).getTime()).toBe( + new Date(2026, 8, 24).getTime(), + ); + expect(new Date(data.filters.created_before).getTime()).toBeGreaterThan( + new Date(2026, 9, 23, 23, 59).getTime(), + ); + expect(data.date_range_context.from).toBe("2026-09-24"); + expect(data.date_range_context.to).toBe("2026-10-23"); + }); + + it("discards the draft on Cancel", async () => { + await renderPanel(false); + await click( + screen.getByRole("button", { name: /filter by date ingested/i }), + ); + await click(screen.getByRole("button", { name: "Last 90 days" })); + await click(screen.getByRole("button", { name: "Cancel" })); + + const data = await save(); + expect(data.filters).not.toHaveProperty("created_after"); + }); + + it("sets From and To through the date fields", async () => { + await renderPanel(false); + await click( + screen.getByRole("button", { name: /filter by date ingested/i }), + ); + await click(screen.getByRole("button", { name: "Custom date range" })); + expect(screen.getByRole("button", { name: "From date" })).toHaveAttribute( + "aria-expanded", + "true", + ); + + const dayButton = (n: number) => + screen + .getAllByRole("gridcell") + .find( + (c) => + c.tagName === "BUTTON" && + !c.className.includes("day-outside") && + c.textContent === String(n), + ) as HTMLElement; + await click(dayButton(7)); + await click(screen.getByRole("button", { name: "To date" })); + // The To calendar opens on its own value's month (October). + await click(dayButton(10)); + await click(screen.getByRole("button", { name: "Apply" })); + + const data = await save(); + // Custom starts as Sep 24 - Oct 23, so the From calendar opens on September. + expect(data.date_range_context.from).toBe("2026-09-07"); + // A To before the draft's start would invert the range; the field keeps it + // disabled, so the To pick above is the first selectable day on or after From. + expect(data.date_range_context.to).toBe("2026-10-10"); + }); +}); diff --git a/frontend/components/knowledge-filter-panel.tsx b/frontend/components/knowledge-filter-panel.tsx index c64b7852c2..2426dbe2ad 100644 --- a/frontend/components/knowledge-filter-panel.tsx +++ b/frontend/components/knowledge-filter-panel.tsx @@ -10,6 +10,8 @@ import { useGetSearchAggregations, } from "@/app/api/queries/useGetSearchAggregations"; import { FilterIconPopover } from "@/components/filter-icon-popover"; +import { useDateRangeFilterEnabled } from "@/components/knowledge-date-range-filter"; +import { KnowledgeFilterDateField } from "@/components/knowledge-filter-date-field"; import { Button } from "@/components/ui/button"; import { Card, @@ -27,6 +29,15 @@ import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; import { usePermissions } from "@/hooks/use-permissions"; import { trackButton } from "@/lib/analytics"; import type { FilterColor, IconKey } from "@/lib/filter-constants"; +import { + applySelection, + displayRangeFromSaved, + foreignTimeZone, + hasDateBounds, + NO_DATE_RANGE, + type SavedDateRange, + savedDateRangeOf, +} from "@/lib/saved-date-range"; interface AvailableFacets { data_sources: FacetBucket[]; @@ -96,6 +107,9 @@ export function KnowledgeFilterPanel() { owners: ["*"] as string[], // Default to wildcard connector_types: ["*"] as string[], // Default to wildcard }); + // Unsaved draft of the filter's date range; the toolbar's temporary range is separate. + const [dateDraft, setDateDraft] = useState(NO_DATE_RANGE); + const dateRangeEnabled = useDateRangeFilterEnabled(); const [resultLimit, setResultLimit] = useState(10); const [scoreThreshold, setScoreThreshold] = useState(0); @@ -127,6 +141,7 @@ export function KnowledgeFilterPanel() { }; setSelectedFilters(processedFilters); + setDateDraft(savedDateRangeOf(parsedFilterData)); setResultLimit(parsedFilterData.limit || 10); setScoreThreshold(parsedFilterData.scoreThreshold || 0); setName(selectedFilter.name); @@ -149,6 +164,7 @@ export function KnowledgeFilterPanel() { owners: filters.owners ?? ["*"], connector_types: filters.connector_types ?? ["*"], }); + setDateDraft(savedDateRangeOf(parsedFilterData)); setResultLimit(parsedFilterData.limit || 10); setScoreThreshold(parsedFilterData.scoreThreshold || 0); setName(""); @@ -200,7 +216,17 @@ export function KnowledgeFilterPanel() { } const filterData = { query, - filters: selectedFilters, + filters: { + ...selectedFilters, + ...(dateDraft.createdAfter && { + created_after: dateDraft.createdAfter, + }), + ...(dateDraft.createdBefore && { + created_before: dateDraft.createdBefore, + }), + }, + ...(hasDateBounds(dateDraft) && + dateDraft.context && { date_range_context: dateDraft.context }), limit: resultLimit, scoreThreshold, color, @@ -431,6 +457,19 @@ export function KnowledgeFilterPanel() { /> + {dateRangeEnabled && ( +
+ + + setDateDraft((prev) => applySelection(prev, next)) + } + savedTimeZone={foreignTimeZone(dateDraft)} + /> +
+ )} + {/* Result Limit Control - exactly like search page */}
diff --git a/frontend/components/knowledge-search-bar.tsx b/frontend/components/knowledge-search-bar.tsx index ca8cc02c7a..69965ba269 100644 --- a/frontend/components/knowledge-search-bar.tsx +++ b/frontend/components/knowledge-search-bar.tsx @@ -7,6 +7,7 @@ import { useSyncAllConnectors, useSyncAllConnectorsPreview, } from "@/app/api/mutations/useSyncConnector"; +import { KnowledgeDateRangeFilter } from "@/components/knowledge-date-range-filter"; import { RequirePermission } from "@/components/require-permission"; import { Button } from "@/components/ui/button"; import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; @@ -203,6 +204,7 @@ export const KnowledgeSearchBar = () => { +
diff --git a/frontend/components/ui/calendar.test.tsx b/frontend/components/ui/calendar.test.tsx new file mode 100644 index 0000000000..3604aa8eb3 --- /dev/null +++ b/frontend/components/ui/calendar.test.tsx @@ -0,0 +1,31 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { Calendar } from "./calendar"; + +describe("Calendar", () => { + it("renders the calendar grid", () => { + render(); + expect(screen.getByRole("grid")).toBeInTheDocument(); + }); + + it("renders navigation buttons with chevron icons", () => { + render(); + // DayPicker renders two nav buttons (previous/next month) that use our custom icons + const navButtons = screen.getAllByRole("button"); + expect(navButtons.length).toBeGreaterThanOrEqual(2); + }); + + it("applies extra className to the root element", () => { + const { container } = render( + , + ); + expect(container.firstChild).toHaveClass("my-custom"); + }); + + it("renders day numbers for the current month", () => { + render(); + // At least one day button should be visible + const dayButtons = screen.getAllByRole("gridcell"); + expect(dayButtons.length).toBeGreaterThan(0); + }); +}); diff --git a/frontend/components/ui/calendar.tsx b/frontend/components/ui/calendar.tsx new file mode 100644 index 0000000000..c3bf9fb9cd --- /dev/null +++ b/frontend/components/ui/calendar.tsx @@ -0,0 +1,75 @@ +"use client"; + +import { ChevronLeft, ChevronRight } from "lucide-react"; +import * as React from "react"; +import { DayPicker } from "react-day-picker"; +import { buttonVariants } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; + +export type CalendarProps = React.ComponentProps; + +function Calendar({ + className, + classNames, + showOutsideDays = true, + components, + ...props +}: CalendarProps) { + return ( + , + IconRight: () => , + ...components, + }} + {...props} + /> + ); +} +Calendar.displayName = "Calendar"; + +export { Calendar }; diff --git a/frontend/contexts/auth-context.tsx b/frontend/contexts/auth-context.tsx index a764468f34..b99bdf29bc 100644 --- a/frontend/contexts/auth-context.tsx +++ b/frontend/contexts/auth-context.tsx @@ -291,7 +291,11 @@ export function AuthProvider({ children }: AuthProviderProps) { if (user || isNoAuthMode || isIbmAuthMode) { setPermissionsResolved(false); void fetchPermissions(); - } else { + } else if (!isLoading) { + // Only resolve with empty permissions once auth has finished loading. + // While isLoading=true the final user/mode is not known yet, so resolving + // early creates a transient state (permissionsResolved=true, roles=[]) that + // can race with the post-checkAuth effect run that actually fetches roles. resetPermissionState(); setPermissionsResolved(true); } @@ -299,6 +303,7 @@ export function AuthProvider({ children }: AuthProviderProps) { user, isNoAuthMode, isIbmAuthMode, + isLoading, fetchPermissions, resetPermissionState, ]); diff --git a/frontend/contexts/knowledge-filter-context.test.tsx b/frontend/contexts/knowledge-filter-context.test.tsx new file mode 100644 index 0000000000..cf9015020a --- /dev/null +++ b/frontend/contexts/knowledge-filter-context.test.tsx @@ -0,0 +1,134 @@ +import { act, renderHook } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { describe, expect, it } from "vitest"; +import { + KnowledgeFilterProvider, + useKnowledgeFilter, +} from "./knowledge-filter-context"; + +const AFTER = "2026-10-07T07:00:00.000Z"; +const BEFORE = "2026-10-11T06:59:59.999999Z"; + +function filter(id: string, queryData: object) { + return { + id, + name: id, + description: "", + query_data: JSON.stringify(queryData), + owner: "u", + created_at: "", + updated_at: "", + }; +} + +const withDates = filter("pdf-oct", { + filters: { + document_types: ["application/pdf"], + created_after: AFTER, + created_before: BEFORE, + }, + date_range_context: { + from: "2026-10-07", + to: "2026-10-10", + timezone: "America/Los_Angeles", + }, +}); +const withoutDates = filter("plain", { + filters: { document_types: ["application/pdf"] }, +}); + +function setup() { + return renderHook(() => useKnowledgeFilter(), { + wrapper: ({ children }: { children: ReactNode }) => ( + {children} + ), + }); +} + +describe("KnowledgeFilterProvider saved date range", () => { + it("restores saved bounds and display days when a filter is selected", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withDates)); + expect(result.current.effectiveDateRange).toMatchObject({ + createdAfter: AFTER, + createdBefore: BEFORE, + }); + expect(result.current.dateRange?.from?.getDate()).toBe(7); + expect(result.current.dateRange?.to?.getDate()).toBe(10); + }); + + it("leaves a filter without dates unrestricted", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withoutDates)); + expect(result.current.dateRange).toBeNull(); + expect(result.current.effectiveDateRange).toEqual({}); + }); + + it("keeps a temporary toolbar range separate from the saved one", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withDates)); + act(() => + result.current.setDateRange({ + from: new Date(2026, 9, 1), + to: new Date(2026, 9, 2), + }), + ); + expect(result.current.effectiveDateRange.createdAfter).not.toBe(AFTER); + expect(result.current.savedDateRange.createdAfter).toBe(AFTER); + expect(result.current.parsedFilterData?.filters.created_after).toBe(AFTER); + }); + + it("clearing a temporary range removes both bounds instead of falling back", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withDates)); + act(() => result.current.setDateRange(null)); + expect(result.current.dateRange).toBeNull(); + expect(result.current.effectiveDateRange.createdAfter).toBeUndefined(); + expect(result.current.effectiveDateRange.createdBefore).toBeUndefined(); + expect(result.current.savedDateRange.createdAfter).toBe(AFTER); + }); + + it("drops the temporary range when switching or clearing filters", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withDates)); + act(() => result.current.setDateRange(null)); + + act(() => result.current.setSelectedFilter(withDates)); + expect(result.current.effectiveDateRange.createdAfter).toBe(AFTER); + + act(() => result.current.setDateRange({ from: new Date(2026, 9, 1) })); + act(() => result.current.setSelectedFilter(withoutDates)); + expect(result.current.effectiveDateRange).toEqual({}); + + act(() => result.current.setDateRange({ from: new Date(2026, 9, 1) })); + act(() => result.current.clearFilter()); + expect(result.current.dateRange).toBeNull(); + }); + + it("re-applying the displayed days keeps the exact saved bounds", () => { + const { result } = setup(); + act(() => result.current.setSelectedFilter(withDates)); + const shown = result.current.dateRange; + act(() => result.current.setDateRange(shown)); + expect(result.current.effectiveDateRange.createdAfter).toBe(AFTER); + expect(result.current.effectiveDateRange.createdBefore).toBe(BEFORE); + }); + + it("starts a new filter without dates even with a temporary toolbar range", () => { + const { result } = setup(); + act(() => result.current.setDateRange({ from: new Date(2026, 9, 1) })); + expect(result.current.effectiveDateRange.createdAfter).toBeTruthy(); + + act(() => result.current.startCreateMode()); + + expect(result.current.parsedFilterData?.filters).not.toHaveProperty( + "created_after", + ); + expect(result.current.parsedFilterData?.filters).not.toHaveProperty( + "created_before", + ); + expect(result.current.parsedFilterData?.dateRangeContext).toBeUndefined(); + expect(result.current.dateRange).toBeNull(); + expect(result.current.effectiveDateRange).toEqual({}); + }); +}); diff --git a/frontend/contexts/knowledge-filter-context.tsx b/frontend/contexts/knowledge-filter-context.tsx index 3f9acc0c82..146301af69 100644 --- a/frontend/contexts/knowledge-filter-context.tsx +++ b/frontend/contexts/knowledge-filter-context.tsx @@ -5,6 +5,7 @@ import React, { type ReactNode, use, useCallback, + useMemo, useState, } from "react"; import { @@ -12,6 +13,16 @@ import { type FilterColor, type IconKey, } from "@/lib/filter-constants"; +import { + applySelection, + type DateRangeContext, + type DateRangeValue, + displayRangeFromSaved, + NO_DATE_RANGE, + parseSavedDateRange, + type SavedDateRange, + savedDateRangeOf, +} from "@/lib/saved-date-range"; interface KnowledgeFilter { id: string; @@ -30,7 +41,12 @@ export interface ParsedQueryData { document_types: string[]; owners: string[]; connector_types: string[]; + /** Saved inclusive UTC bounds on indexed_time; absent means any time. */ + created_after?: string; + created_before?: string; }; + /** Calendar days + timezone the saved bounds were picked in (display only). */ + dateRangeContext?: DateRangeContext; limit: number; scoreThreshold: number; color: FilterColor; @@ -56,6 +72,20 @@ interface KnowledgeFilterContextType { /** Filenames checked in the knowledge table; seeds data_sources on create. */ selectedSources: string[]; setSelectedSources: (sources: string[]) => void; + /** + * Effective date range for the file grid: the selected filter's saved range, + * unless the toolbar has overridden it. Overrides are temporary; they are + * saved only when the Knowledge Filter itself is saved. + */ + dateRange: DateRangeValue | null; + setDateRange: (range: DateRangeValue | null) => void; + /** Saved range of the selected filter (what the editor starts from). */ + savedDateRange: SavedDateRange; + /** + * UTC bounds used for retrieval. Unchanged saved ranges keep their exact + * saved bounds in any viewer timezone; a cleared range has neither bound. + */ + effectiveDateRange: SavedDateRange; } const KnowledgeFilterContext = createContext< @@ -90,14 +120,33 @@ export function KnowledgeFilterProvider({ const [createMode, setCreateMode] = useState(false); const [queryOverride, setQueryOverride] = useState(""); const [selectedSources, setSelectedSources] = useState([]); + // null = no toolbar override (use the saved range); NO_DATE_RANGE = "any time". + const [dateRangeOverride, setDateRangeOverride] = + useState(null); + + const savedDateRange = useMemo( + () => savedDateRangeOf(parsedFilterData), + [parsedFilterData], + ); + const effectiveDateRange = dateRangeOverride ?? savedDateRange; + const dateRange = useMemo( + () => displayRangeFromSaved(effectiveDateRange), + [effectiveDateRange], + ); + const setDateRange = (range: DateRangeValue | null) => { + setDateRangeOverride(applySelection(effectiveDateRange, range)); + }; const setSelectedFilter = (filter: KnowledgeFilter | null) => { setSelectedFilterState(filter); + // A temporary toolbar range never outlives the filter it was set against. + setDateRangeOverride(null); if (filter) { setCreateMode(false); try { const raw = JSON.parse(filter.query_data); + const saved = parseSavedDateRange(raw); // Normalize parsed data with defaults for missing fields // This handles filters created via API with incomplete queryData const parsed: ParsedQueryData = { @@ -107,7 +156,12 @@ export function KnowledgeFilterProvider({ document_types: raw.filters?.document_types ?? ["*"], owners: raw.filters?.owners ?? ["*"], connector_types: raw.filters?.connector_types ?? ["*"], + ...(saved.createdAfter && { created_after: saved.createdAfter }), + ...(saved.createdBefore && { + created_before: saved.createdBefore, + }), }, + ...(saved.context && { dateRangeContext: saved.context }), limit: raw.limit ?? 10, scoreThreshold: raw.scoreThreshold ?? 0, color: raw.color ?? "zinc", @@ -157,6 +211,7 @@ export function KnowledgeFilterProvider({ setPanelMode("filters"); setCreateMode(true); setSelectedFilterState(null); + setDateRangeOverride(null); setParsedFilterData({ query: "", filters: { @@ -164,6 +219,7 @@ export function KnowledgeFilterProvider({ document_types: ["*"], owners: ["*"], connector_types: ["*"], + // No dates: a temporary toolbar range is never saved implicitly. }, limit: 10, scoreThreshold: 0, @@ -202,6 +258,10 @@ export function KnowledgeFilterProvider({ setQueryOverride, selectedSources, setSelectedSources, + dateRange, + setDateRange, + savedDateRange, + effectiveDateRange, }; return ( diff --git a/frontend/lib/date-range-filter.test.ts b/frontend/lib/date-range-filter.test.ts new file mode 100644 index 0000000000..fda27d498b --- /dev/null +++ b/frontend/lib/date-range-filter.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from "vitest"; +import { isDateRangeFilterEnabled } from "./date-range-filter"; + +describe("isDateRangeFilterEnabled", () => { + it("is enabled for OSS", () => { + expect(isDateRangeFilterEnabled("oss")).toBe(true); + }); + + it.each([ + "saas", + "", + null, + undefined, + ])("is disabled for run mode %s", (runMode) => { + expect(isDateRangeFilterEnabled(runMode)).toBe(false); + }); +}); diff --git a/frontend/lib/date-range-filter.ts b/frontend/lib/date-range-filter.ts new file mode 100644 index 0000000000..4a24132f17 --- /dev/null +++ b/frontend/lib/date-range-filter.ts @@ -0,0 +1,29 @@ +/** Knowledge date-range filter gate (mirrors the ingest-preview run-mode gate). */ + +import { IBM_THEME_DEV } from "@/lib/brand"; + +/** + * Eligible run modes for showing the Knowledge date-range filter. + * SaaS omitted until product approval — add "saas" here to restore. + */ +const DATE_RANGE_FILTER_RUN_MODES = new Set(["oss"]); + +/** + * Client-side gate for showing the date-range filter. Pass `isCloudBrand` so the + * local IBM theme switch hides it the same way real `OPENRAG_RUN_MODE=saas` would. + * + * This gate controls date-picker visibility only. + * Saved date bounds still apply to search and chat when controls are hidden. + */ +export function isDateRangeFilterEnabled( + runMode: string | null | undefined, + options?: { isCloudBrand?: boolean }, +): boolean { + if (runMode == null || !DATE_RANGE_FILTER_RUN_MODES.has(runMode)) { + return false; + } + if (IBM_THEME_DEV && options?.isCloudBrand) { + return false; + } + return true; +} diff --git a/frontend/lib/filter-normalization.test.ts b/frontend/lib/filter-normalization.test.ts new file mode 100644 index 0000000000..79027f804d --- /dev/null +++ b/frontend/lib/filter-normalization.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "vitest"; +import { buildSearchPayloadFilters } from "./filter-normalization"; + +const AFTER = "2026-10-07T07:00:00.000Z"; +const BEFORE = "2026-10-11T06:59:59.999999Z"; + +describe("buildSearchPayloadFilters dates", () => { + it("keeps saved bounds alongside list filters", () => { + expect( + buildSearchPayloadFilters({ + document_types: ["application/pdf"], + owners: ["*"], + created_after: AFTER, + created_before: BEFORE, + }), + ).toEqual({ + document_types: ["application/pdf"], + created_after: AFTER, + created_before: BEFORE, + }); + }); + + it("treats a date-only filter as an active filter", () => { + expect(buildSearchPayloadFilters({ created_after: AFTER })).toEqual({ + created_after: AFTER, + }); + }); + + it("drops blank bounds and stays undefined for filters without dates", () => { + expect( + buildSearchPayloadFilters({ created_after: "", created_before: "" }), + ).toBeUndefined(); + expect(buildSearchPayloadFilters({ data_sources: ["*"] })).toBeUndefined(); + }); + + it("omits dates when the caller supplies its own effective bounds", () => { + expect( + buildSearchPayloadFilters( + { document_types: ["application/pdf"], created_after: AFTER }, + { includeDates: false }, + ), + ).toEqual({ document_types: ["application/pdf"] }); + }); +}); diff --git a/frontend/lib/filter-normalization.ts b/frontend/lib/filter-normalization.ts index 59427f288a..3c25f37334 100644 --- a/frontend/lib/filter-normalization.ts +++ b/frontend/lib/filter-normalization.ts @@ -3,6 +3,9 @@ export interface FilterInput { document_types?: string[]; owners?: string[]; connector_types?: string[]; + /** Inclusive UTC bounds on indexed_time (ISO 8601). */ + created_after?: string; + created_before?: string; } export interface NormalizedSelectedFilters { @@ -30,9 +33,22 @@ function normalizeSelectedFilters( }; } +function dateBound(value?: string): string | undefined { + return typeof value === "string" && value !== "" ? value : undefined; +} + +/** + * Build the `filters` object for search/chat/nudge requests. + * + * Saved date bounds are included by default. The Knowledge page passes + * `includeDates: false` because it sends its effective bounds separately, so a + * cleared temporary range can never fall back to the saved one. + */ export function buildSearchPayloadFilters( filters?: FilterInput, + options: { includeDates?: boolean } = {}, ): FilterInput | undefined { + const { includeDates = true } = options; const normalized = normalizeSelectedFilters(filters); const payloadFilters: FilterInput = {}; @@ -48,6 +64,12 @@ export function buildSearchPayloadFilters( if (normalized.connector_types.length > 0) { payloadFilters.connector_types = normalized.connector_types; } + if (includeDates) { + const createdAfter = dateBound(filters?.created_after); + const createdBefore = dateBound(filters?.created_before); + if (createdAfter) payloadFilters.created_after = createdAfter; + if (createdBefore) payloadFilters.created_before = createdBefore; + } return Object.keys(payloadFilters).length > 0 ? payloadFilters : undefined; } diff --git a/frontend/lib/knowledge-table-state.test.ts b/frontend/lib/knowledge-table-state.test.ts index ff187e2ab9..7b14eea1f6 100644 --- a/frontend/lib/knowledge-table-state.test.ts +++ b/frontend/lib/knowledge-table-state.test.ts @@ -171,6 +171,90 @@ describe("buildKnowledgeTableRows", () => { ]); }); + describe("date range", () => { + const yesterday = { + from: new Date(2026, 8, 7), + to: new Date(2026, 8, 7), + }; + + it("drops task files created outside the selected range", () => { + const rows = buildKnowledgeTableRows( + [], + [ + taskFile({ + filename: "today.pdf", + created_at: new Date(2026, 8, 8, 10).toISOString(), + }), + ], + false, + yesterday, + ); + expect(rows).toHaveLength(0); + }); + + it("keeps task files created inside the range, including its last day", () => { + const rows = buildKnowledgeTableRows( + [], + [ + taskFile({ + filename: "late.pdf", + created_at: new Date(2026, 8, 7, 23, 30).toISOString(), + }), + ], + false, + yesterday, + ); + expect(rows.map((r) => r.filename)).toEqual(["late.pdf"]); + }); + + it("applies an open-ended range", () => { + const created_at = new Date(2026, 8, 1, 10).toISOString(); + const files = [taskFile({ filename: "old.pdf", created_at })]; + expect( + buildKnowledgeTableRows([], files, false, { + from: new Date(2026, 8, 5), + }), + ).toHaveLength(0); + expect( + buildKnowledgeTableRows([], files, false, { to: new Date(2026, 8, 5) }), + ).toHaveLength(1); + }); + + it("exposes the task creation time as indexed_time on task-only rows", () => { + const created_at = new Date(2026, 8, 7, 10).toISOString(); + const rows = buildKnowledgeTableRows( + [], + [taskFile({ filename: "pending.pdf", created_at })], + ); + expect(rows[0].indexed_time).toBe(created_at); + }); + + it("keeps every task file when no range is set", () => { + const rows = buildKnowledgeTableRows([], [taskFile()], false, null); + expect(rows).toHaveLength(1); + }); + + it("keeps a task file whose timestamp is unparseable", () => { + const rows = buildKnowledgeTableRows( + [], + [taskFile({ created_at: "not-a-date" })], + false, + yesterday, + ); + expect(rows).toHaveLength(1); + }); + + it("never drops indexed backend rows, which the server already filtered", () => { + const rows = buildKnowledgeTableRows( + [searchFile({ filename: "indexed.pdf" })], + [], + false, + yesterday, + ); + expect(rows).toHaveLength(1); + }); + }); + it("labels a task file with no name or url as Untitled source", () => { const rows = buildKnowledgeTableRows( [], @@ -425,3 +509,42 @@ describe("buildActiveSourceOptions", () => { ).toEqual([]); }); }); + +describe("buildKnowledgeTableRows effective bounds", () => { + // Task rows take their indexed_time from the task's created_at. + const inFlight = (created_at: string) => + taskFile({ filename: "new.pdf", status: "processing", created_at }); + + it("filters in-flight rows by the UTC bounds, not the display days", () => { + // Bounds saved in Los Angeles; display days are deliberately wrong for the + // viewer to prove only the bounds decide. + const bounds = { + createdAfter: "2026-10-07T07:00:00.000Z", + createdBefore: "2026-10-11T06:59:59.999999Z", + }; + const rows = (time: string) => + buildKnowledgeTableRows( + [], + [inFlight(time)], + false, + { from: new Date(2000, 0, 1), to: new Date(2000, 0, 1) }, + bounds, + ).map((r) => r.filename); + + expect(rows("2026-10-07T07:00:00.000Z")).toEqual(["new.pdf"]); + expect(rows("2026-10-11T06:59:59.000Z")).toEqual(["new.pdf"]); + expect(rows("2026-10-07T06:59:59.000Z")).toEqual([]); + expect(rows("2026-10-11T07:00:00.000Z")).toEqual([]); + }); + + it("applies no date filter once the bounds are cleared", () => { + const rows = buildKnowledgeTableRows( + [], + [inFlight("2020-01-01T00:00:00.000Z")], + false, + null, + {}, + ); + expect(rows.map((r) => r.filename)).toEqual(["new.pdf"]); + }); +}); diff --git a/frontend/lib/knowledge-table-state.ts b/frontend/lib/knowledge-table-state.ts index 5b2592dfa4..cb2dd3f631 100644 --- a/frontend/lib/knowledge-table-state.ts +++ b/frontend/lib/knowledge-table-state.ts @@ -188,6 +188,8 @@ export function buildKnowledgeTableRows( searchData: SearchFile[], taskFiles: TaskFile[], hasActiveFilter = false, + dateRange?: { from?: Date; to?: Date } | null, + bounds?: { createdAfter?: string; createdBefore?: string }, ): SearchFile[] { const taskFilesAsFiles: SearchFile[] = taskFiles.map((taskFile) => { const normalizedFilename = @@ -197,6 +199,7 @@ export function buildKnowledgeTableRows( return { filename: normalizedFilename, + indexed_time: taskFile.created_at || undefined, mimetype: taskFile.mimetype, source_url: taskFile.source_url || "", size: taskFile.size, @@ -284,7 +287,55 @@ export function buildKnowledgeTableRows( return backendFiles; } - return [...backendFiles, ...filteredTaskFiles]; + // Effective UTC bounds win over the display days so in-flight rows agree with + // the server results even when the range was saved in another timezone. + const boundFrom = bounds?.createdAfter + ? new Date(bounds.createdAfter).getTime() + : undefined; + const boundTo = bounds?.createdBefore + ? new Date(bounds.createdBefore).getTime() + : undefined; + const rangeFrom = dateRange?.from; + const rangeTo = dateRange?.to; + const useBounds = bounds !== undefined; + const hasRange = useBounds + ? boundFrom !== undefined || boundTo !== undefined + : Boolean(rangeFrom ?? rangeTo); + const dateFilteredTaskFiles = hasRange + ? filteredTaskFiles.filter((f) => { + if (!f.indexed_time) return true; + const d = new Date(f.indexed_time); + if (Number.isNaN(d.getTime())) return true; + if (useBounds) { + if (boundFrom !== undefined && d.getTime() < boundFrom) return false; + if (boundTo !== undefined && d.getTime() > boundTo) return false; + return true; + } + if (rangeFrom) { + const start = new Date( + rangeFrom.getFullYear(), + rangeFrom.getMonth(), + rangeFrom.getDate(), + ); + if (d < start) return false; + } + if (rangeTo) { + const end = new Date( + rangeTo.getFullYear(), + rangeTo.getMonth(), + rangeTo.getDate(), + 23, + 59, + 59, + 999, + ); + if (d > end) return false; + } + return true; + }) + : filteredTaskFiles; + + return [...backendFiles, ...dateFilteredTaskFiles]; } export function buildActiveSourceOptions( diff --git a/frontend/lib/saved-date-range.test.ts b/frontend/lib/saved-date-range.test.ts new file mode 100644 index 0000000000..e49e367559 --- /dev/null +++ b/frontend/lib/saved-date-range.test.ts @@ -0,0 +1,113 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { + applySelection, + displayRangeFromSaved, + foreignTimeZone, + NO_DATE_RANGE, + parseSavedDateRange, + savedRangeFromSelection, +} from "./saved-date-range"; + +// October 7-10 picked in Los Angeles (PDT, UTC-7): full local days as UTC bounds. +const LA_RANGE = { + createdAfter: "2026-10-07T07:00:00.000Z", + createdBefore: "2026-10-11T06:59:59.999999Z", + context: { + from: "2026-10-07", + to: "2026-10-10", + timezone: "America/Los_Angeles", + }, +}; + +const originalTZ = process.env.TZ; +afterEach(() => { + if (originalTZ === undefined) delete process.env.TZ; + else process.env.TZ = originalTZ; +}); + +const ymd = (d: Date | undefined) => + d && [d.getFullYear(), d.getMonth() + 1, d.getDate()].join("-"); + +describe("parseSavedDateRange", () => { + it("returns an empty range for filters saved without dates", () => { + expect(parseSavedDateRange({ filters: { document_types: ["*"] } })).toEqual( + {}, + ); + expect(parseSavedDateRange(undefined)).toEqual({}); + }); + + it("reads bounds and context", () => { + expect( + parseSavedDateRange({ + filters: { + created_after: LA_RANGE.createdAfter, + created_before: LA_RANGE.createdBefore, + }, + date_range_context: LA_RANGE.context, + }), + ).toEqual(LA_RANGE); + }); + + it("ignores context when there are no bounds", () => { + expect( + parseSavedDateRange({ date_range_context: LA_RANGE.context }), + ).toEqual({}); + }); +}); + +describe("timezone handling", () => { + it("restores the saved calendar days, not the viewer's shifted days", () => { + process.env.TZ = "Europe/London"; + const shown = displayRangeFromSaved(LA_RANGE); + expect(ymd(shown?.from)).toBe("2026-10-7"); + expect(ymd(shown?.to)).toBe("2026-10-10"); + }); + + it("keeps the exact saved bounds when the same days are re-applied elsewhere", () => { + process.env.TZ = "Europe/London"; + const shown = displayRangeFromSaved(LA_RANGE); + expect(applySelection(LA_RANGE, shown)).toBe(LA_RANGE); + }); + + it("recomputes bounds in the viewer's timezone once the user edits the range", () => { + process.env.TZ = "Europe/London"; + const edited = applySelection(LA_RANGE, { + from: new Date(2026, 9, 7), + to: new Date(2026, 9, 9), + }); + // London is on BST (UTC+1) in October. + expect(edited.createdAfter).toBe("2026-10-06T23:00:00.000Z"); + expect(edited.createdBefore).toBe("2026-10-09T22:59:59.999999Z"); + expect(edited.context).toEqual({ + from: "2026-10-07", + to: "2026-10-09", + timezone: "Europe/London", + }); + }); + + it("surfaces the saved timezone only when it differs from the viewer's", () => { + process.env.TZ = "Europe/London"; + expect(foreignTimeZone(LA_RANGE)).toBe("America/Los_Angeles"); + process.env.TZ = "America/Los_Angeles"; + expect(foreignTimeZone(LA_RANGE)).toBeNull(); + expect(foreignTimeZone(NO_DATE_RANGE)).toBeNull(); + }); + + it("derives days from the bounds when a filter has no context", () => { + process.env.TZ = "UTC"; + const shown = displayRangeFromSaved({ + createdAfter: "2026-10-07T00:00:00.000Z", + createdBefore: "2026-10-10T23:59:59.999Z", + }); + expect(ymd(shown?.from)).toBe("2026-10-7"); + expect(ymd(shown?.to)).toBe("2026-10-10"); + }); +}); + +describe("clearing", () => { + it("clearing a range yields neither bound", () => { + expect(applySelection(LA_RANGE, null)).toEqual({}); + expect(applySelection(LA_RANGE, {})).toEqual({}); + expect(savedRangeFromSelection(null)).toEqual({}); + }); +}); diff --git a/frontend/lib/saved-date-range.ts b/frontend/lib/saved-date-range.ts new file mode 100644 index 0000000000..d25325e1be --- /dev/null +++ b/frontend/lib/saved-date-range.ts @@ -0,0 +1,209 @@ +/** + * Date range persisted on a Knowledge Filter. + * + * Two representations travel together and play different roles: + * - `createdAfter` / `createdBefore` are explicit UTC instants. They are the + * only values used for retrieval, so a shared filter never shifts its window. + * - `context` (calendar days + IANA timezone) is what the picker displays. + * + * An unchanged range keeps its saved bounds verbatim in any viewer timezone; + * bounds are only recomputed (in the viewer's timezone) when the user edits it. + */ + +import { isSameDay } from "date-fns"; +import { + toLocalDayEndISOString, + toLocalDayKey, + toLocalDayStartISOString, +} from "@/lib/time-utils"; + +export interface DateRangeContext { + from?: string; + to?: string; + timezone?: string; +} + +export interface SavedDateRange { + createdAfter?: string; + createdBefore?: string; + context?: DateRangeContext; +} + +export interface DateRangeValue { + from?: Date; + to?: Date; +} + +/** No bounds at all: "any time". Used as an explicit override that clears a saved range. */ +export const NO_DATE_RANGE: SavedDateRange = Object.freeze({}); + +const DAY_RE = /^(\d{4})-(\d{2})-(\d{2})$/; + +function nonEmptyString(value: unknown): string | undefined { + return typeof value === "string" && value !== "" ? value : undefined; +} + +export function parseCalendarDay(day: string): Date | null { + const match = DAY_RE.exec(day); + if (!match) return null; + const [y, m, d] = [Number(match[1]), Number(match[2]), Number(match[3])]; + const date = new Date(y, m - 1, d); + return date.getFullYear() === y && + date.getMonth() === m - 1 && + date.getDate() === d + ? date + : null; +} + +export function viewerTimeZone(): string { + return Intl.DateTimeFormat().resolvedOptions().timeZone; +} + +export function hasDateBounds(range: SavedDateRange | undefined): boolean { + return Boolean(range?.createdAfter || range?.createdBefore); +} + +/** Read the optional date fields out of parsed (already JSON-decoded) query data. */ +export function parseSavedDateRange(raw: unknown): SavedDateRange { + const data = (raw ?? {}) as { + filters?: { created_after?: unknown; created_before?: unknown }; + date_range_context?: { + from?: unknown; + to?: unknown; + timezone?: unknown; + } | null; + }; + const createdAfter = nonEmptyString(data.filters?.created_after); + const createdBefore = nonEmptyString(data.filters?.created_before); + if (!createdAfter && !createdBefore) return {}; + + const ctx = data.date_range_context; + const context: DateRangeContext = {}; + if (ctx && typeof ctx === "object") { + const from = nonEmptyString(ctx.from); + const to = nonEmptyString(ctx.to); + const timezone = nonEmptyString(ctx.timezone); + if (from) context.from = from; + if (to) context.to = to; + if (timezone) context.timezone = timezone; + } + return { + ...(createdAfter && { createdAfter }), + ...(createdBefore && { createdBefore }), + ...(Object.keys(context).length > 0 && { context }), + }; +} + +/** Calendar days to show in the picker for a saved range, or null when unrestricted. */ +export function displayRangeFromSaved( + range: SavedDateRange | undefined, +): DateRangeValue | null { + if (!range || !hasDateBounds(range)) return null; + + const from = range.context?.from + ? parseCalendarDay(range.context.from) + : null; + const to = range.context?.to ? parseCalendarDay(range.context.to) : null; + if (from || to) { + return { ...(from && { from }), ...(to && { to }) }; + } + + // No usable context (e.g. a filter created through the API): derive the + // days from the bounds in the viewer's timezone. + const after = range.createdAfter ? new Date(range.createdAfter) : null; + const before = range.createdBefore ? new Date(range.createdBefore) : null; + const derivedFrom = after && !Number.isNaN(after.getTime()) ? after : null; + const derivedTo = before && !Number.isNaN(before.getTime()) ? before : null; + if (!derivedFrom && !derivedTo) return null; + return { + ...(derivedFrom && { from: derivedFrom }), + ...(derivedTo && { to: derivedTo }), + }; +} + +/** Bounds + display context for a range the user just picked (viewer's timezone). */ +export function savedRangeFromSelection( + range: DateRangeValue | null | undefined, +): SavedDateRange { + if (!range?.from && !range?.to) return NO_DATE_RANGE; + return { + ...(range.from && { + createdAfter: toLocalDayStartISOString(range.from), + }), + ...(range.to && { createdBefore: toLocalDayEndISOString(range.to) }), + context: { + ...(range.from && { from: toLocalDayKey(range.from) }), + ...(range.to && { to: toLocalDayKey(range.to) }), + timezone: viewerTimeZone(), + }, + }; +} + +function dayKey(date: Date | undefined): string | undefined { + return date ? toLocalDayKey(date) : undefined; +} + +export function sameCalendarDays( + a: DateRangeValue | null | undefined, + b: DateRangeValue | null | undefined, +): boolean { + return dayKey(a?.from) === dayKey(b?.from) && dayKey(a?.to) === dayKey(b?.to); +} + +/** + * Next saved range after the user picks `next` while `current` is in effect. + * Re-picking the days already shown keeps `current` untouched, so applying an + * unchanged range in another timezone never rewrites its saved bounds. + */ +export function applySelection( + current: SavedDateRange, + next: DateRangeValue | null | undefined, +): SavedDateRange { + if (sameCalendarDays(displayRangeFromSaved(current), next)) return current; + return savedRangeFromSelection(next); +} + +/** Saved timezone worth surfacing: only when it differs from the viewer's. */ +export function foreignTimeZone( + range: SavedDateRange | undefined, +): string | null { + const tz = range?.context?.timezone; + return tz && hasDateBounds(range) && tz !== viewerTimeZone() ? tz : null; +} + +/** Saved range carried by parsed query data. */ +export function savedDateRangeOf( + parsed: { + filters: { created_after?: string; created_before?: string }; + dateRangeContext?: DateRangeContext; + } | null, +): SavedDateRange { + const createdAfter = parsed?.filters.created_after || undefined; + const createdBefore = parsed?.filters.created_before || undefined; + if (!createdAfter && !createdBefore) return NO_DATE_RANGE; + return { + ...(createdAfter && { createdAfter }), + ...(createdBefore && { createdBefore }), + ...(parsed?.dateRangeContext && { context: parsed.dateRangeContext }), + }; +} + +function formatLabelDate(date: Date): string { + return date.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + year: "numeric", + }); +} + +/** "A – B", a single day, or "From A" / "Until B" for a one-sided range. */ +export function describeRange(range: DateRangeValue | null): string | null { + if (range?.from && range.to) { + return isSameDay(range.from, range.to) + ? formatLabelDate(range.from) + : `${formatLabelDate(range.from)} – ${formatLabelDate(range.to)}`; + } + if (range?.from) return `From ${formatLabelDate(range.from)}`; + if (range?.to) return `Until ${formatLabelDate(range.to)}`; + return null; +} diff --git a/frontend/lib/time-utils.test.ts b/frontend/lib/time-utils.test.ts new file mode 100644 index 0000000000..fed12c4ac9 --- /dev/null +++ b/frontend/lib/time-utils.test.ts @@ -0,0 +1,106 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { toLocalDayEndISOString, toLocalDayStartISOString } from "./time-utils"; + +describe("toLocalDayStartISOString", () => { + it("sets time to local midnight (00:00:00.000)", () => { + const date = new Date("2024-03-15T15:30:45.123"); + const result = toLocalDayStartISOString(date); + const d = new Date(result); + // getHours/getMinutes/etc. return local time, matching what setHours set + expect(d.getHours()).toBe(0); + expect(d.getMinutes()).toBe(0); + expect(d.getSeconds()).toBe(0); + expect(d.getMilliseconds()).toBe(0); + }); + + it("does not mutate the input date", () => { + const date = new Date("2024-06-01T12:00:00"); + const original = date.getTime(); + toLocalDayStartISOString(date); + expect(date.getTime()).toBe(original); + }); + + it("returns a valid ISO 8601 string", () => { + const date = new Date("2024-03-15"); + const result = toLocalDayStartISOString(date); + expect(result).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/); + expect(new Date(result).getTime()).not.toBeNaN(); + }); +}); + +describe("toLocalDayEndISOString", () => { + it("sets time to local end-of-day (23:59:59.999)", () => { + const date = new Date("2024-03-15T10:00:00"); + const result = toLocalDayEndISOString(date); + const d = new Date(result); + expect(d.getHours()).toBe(23); + expect(d.getMinutes()).toBe(59); + expect(d.getSeconds()).toBe(59); + expect(d.getMilliseconds()).toBe(999); + }); + + it("does not mutate the input date", () => { + const date = new Date("2024-06-01T12:00:00"); + const original = date.getTime(); + toLocalDayEndISOString(date); + expect(date.getTime()).toBe(original); + }); + + it("returns a valid ISO 8601 string", () => { + const date = new Date("2024-03-15"); + const result = toLocalDayEndISOString(date); + expect(result).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.999999Z$/); + expect(new Date(result).getTime()).not.toBeNaN(); + }); + + it("end is later than start for the same input date", () => { + const date = new Date("2024-03-15"); + const start = toLocalDayStartISOString(date); + const end = toLocalDayEndISOString(date); + expect(new Date(end).getTime()).toBeGreaterThan(new Date(start).getTime()); + }); +}); + +describe("local-day filter bounds across timezones", () => { + const originalTz = process.env.TZ; + afterEach(() => { + process.env.TZ = originalTz; + }); + + // 2024-10-06T03:00Z is 8pm Oct 5 in Los Angeles, 11pm Oct 5 in New York, + // and 8:30am Oct 6 in Kolkata. + const uploadedAt = Date.parse("2024-10-06T03:00:00.000Z"); + + function dayContainsUpload(day: number) { + const picked = new Date(2024, 9, day); // what the calendar hands back + const start = Date.parse(toLocalDayStartISOString(picked)); + const end = Date.parse(toLocalDayEndISOString(picked)); + return start <= uploadedAt && uploadedAt <= end; + } + + it.each([ + { tz: "America/Los_Angeles", localDay: 5 }, + { tz: "America/New_York", localDay: 5 }, + { tz: "UTC", localDay: 6 }, + { tz: "Asia/Kolkata", localDay: 6 }, + ])("in $tz the upload lands on the day it is displayed under ($localDay)", ({ + tz, + localDay, + }) => { + process.env.TZ = tz; + expect( + new Date(uploadedAt).toLocaleDateString("en-US", { day: "numeric" }), + ).toBe(String(localDay)); + expect(dayContainsUpload(localDay)).toBe(true); + expect(dayContainsUpload(localDay === 5 ? 6 : 5)).toBe(false); + }); + + it("spans exactly one local day, including across a DST change", () => { + process.env.TZ = "America/New_York"; + // 2024-11-03 is the fall-back day: 25 hours long in New York. + const picked = new Date(2024, 10, 3); + const start = Date.parse(toLocalDayStartISOString(picked)); + const end = Date.parse(toLocalDayEndISOString(picked)); + expect(end - start + 1).toBe(25 * 60 * 60 * 1000); + }); +}); diff --git a/frontend/lib/time-utils.ts b/frontend/lib/time-utils.ts index ce7da67606..906a09a9b7 100644 --- a/frontend/lib/time-utils.ts +++ b/frontend/lib/time-utils.ts @@ -64,3 +64,28 @@ export function formatTaskTimestamp( return `${diffHours} hours ago · ${formatClockTime(date)}`; return `${Math.floor(diffHours / 24)} days ago · ${formatClockTime(date)}`; } + +// Day bounds are the viewer's local calendar day, sent as UTC instants. The +// "Date Ingested" column renders indexed_time in local time too, so a file +// listed under a day is the file that day's filter returns, in any timezone. +export function toLocalDayStartISOString(date: Date): string { + const d = new Date(date); + d.setHours(0, 0, 0, 0); + return d.toISOString(); +} + +export function toLocalDayEndISOString(date: Date): string { + const d = new Date(date); + d.setHours(23, 59, 59, 999); + // The backend stores microsecond-precision timestamps; extend the inclusive + // bound through the day's final microsecond. + return d.toISOString().replace(/\.999Z$/, ".999999Z"); +} + +// Matches the yyyy-MM-dd keys the upload histogram buckets by (viewer's time zone). +export function toLocalDayKey(date: Date): string { + const y = date.getFullYear(); + const m = String(date.getMonth() + 1).padStart(2, "0"); + const d = String(date.getDate()).padStart(2, "0"); + return `${y}-${m}-${d}`; +} diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 99b4b8308c..e92233e050 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -34,6 +34,7 @@ "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "cmdk": "^1.1.1", + "date-fns": "^3.6.0", "dotenv": "^17.4.2", "fuse.js": "^7.4.2", "lucide-react": "^1.22.0", @@ -42,6 +43,7 @@ "next-themes": "^0.4.6", "prom-client": "^15.1.3", "react": "^19.2.7", + "react-day-picker": "^8.10.2", "react-dom": "^19.2.7", "react-dropzone": "^15.0.0", "react-hook-form": "^7.80.0", @@ -5013,6 +5015,16 @@ "node": "^20.19.0 || ^22.12.0 || >=24.0.0" } }, + "node_modules/date-fns": { + "version": "3.6.0", + "resolved": "https://registry.npmjs.org/date-fns/-/date-fns-3.6.0.tgz", + "integrity": "sha512-fRHTG8g/Gif+kSh50gaGEdToemgfj74aRX3swtiouboip5JDLAyDE9F11nHMIcvOaXeOC6D7SpNhi7uFyB7Uww==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/kossnocorp" + } + }, "node_modules/debug": { "version": "4.3.7", "resolved": "https://registry.npmjs.org/debug/-/debug-4.3.7.tgz", @@ -8353,6 +8365,20 @@ "node": ">=0.10.0" } }, + "node_modules/react-day-picker": { + "version": "8.10.2", + "resolved": "https://registry.npmjs.org/react-day-picker/-/react-day-picker-8.10.2.tgz", + "integrity": "sha512-LK68OTbHB3oJNhl9cA0qVizzp3o26w61YSjAFkYi67N86iro32wx86kSNeFU/hq+gI8m1yzWhnomMLfZ041RzQ==", + "license": "MIT", + "funding": { + "type": "individual", + "url": "https://github.com/sponsors/gpbl" + }, + "peerDependencies": { + "date-fns": "^2.28.0 || ^3.0.0", + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/react-dom": { "version": "19.2.7", "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.7.tgz", diff --git a/frontend/package.json b/frontend/package.json index 646d5bd4ec..3b5c289826 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -47,6 +47,7 @@ "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "cmdk": "^1.1.1", + "date-fns": "^3.6.0", "dotenv": "^17.4.2", "fuse.js": "^7.4.2", "lucide-react": "^1.22.0", @@ -55,6 +56,7 @@ "next-themes": "^0.4.6", "prom-client": "^15.1.3", "react": "^19.2.7", + "react-day-picker": "^8.10.2", "react-dom": "^19.2.7", "react-dropzone": "^15.0.0", "react-hook-form": "^7.80.0", diff --git a/sdks/python/openrag_sdk/documents.py b/sdks/python/openrag_sdk/documents.py index d67ea46338..05d92ae6a4 100644 --- a/sdks/python/openrag_sdk/documents.py +++ b/sdks/python/openrag_sdk/documents.py @@ -196,6 +196,8 @@ async def list_files( owner: str | None = None, search: str | None = None, after_key: str | None = None, + created_after: str | None = None, + created_before: str | None = None, ) -> ListFilesResponse: """ List ingested files with cursor-based composite-aggregation pagination (v2). @@ -212,6 +214,10 @@ async def list_files( search: Substring/prefix match against filename. after_key: JSON-encoded composite cursor from a previous response's after_key field. Pass to fetch the next page. + created_after: ISO 8601 timestamp; return only files indexed on or + after this instant (e.g. "2024-01-01" or "2024-01-01T00:00:00Z"). + created_before: ISO 8601 timestamp; return only files indexed on or + before this instant. Returns: ListFilesResponse with files list, approximate total, and next @@ -233,6 +239,10 @@ async def list_files( params["search"] = search if after_key is not None: params["after_key"] = after_key + if created_after is not None: + params["created_after"] = created_after + if created_before is not None: + params["created_before"] = created_before response = await self._client._request( "GET", diff --git a/sdks/typescript/src/documents.ts b/sdks/typescript/src/documents.ts index acb9c16e51..8d89a0b206 100644 --- a/sdks/typescript/src/documents.ts +++ b/sdks/typescript/src/documents.ts @@ -165,6 +165,8 @@ export class DocumentsClient { if (options.owner !== undefined) params.set("owner", options.owner); if (options.search !== undefined) params.set("search", options.search); if (options.after_key !== undefined) params.set("after_key", options.after_key); + if (options.created_after !== undefined) params.set("created_after", options.created_after); + if (options.created_before !== undefined) params.set("created_before", options.created_before); const qs = params.toString(); const path = qs ? `/api/v2/files?${qs}` : "/api/v2/files"; diff --git a/sdks/typescript/src/types.ts b/sdks/typescript/src/types.ts index cc91f7c5bf..18810f8a73 100644 --- a/sdks/typescript/src/types.ts +++ b/sdks/typescript/src/types.ts @@ -147,6 +147,10 @@ export interface ListFilesOptions { owner?: string; search?: string; after_key?: string; + /** ISO 8601 timestamp; return only files indexed on or after this instant. */ + created_after?: string; + /** ISO 8601 timestamp; return only files indexed on or before this instant. */ + created_before?: string; } // Chat history types diff --git a/src/api/files.py b/src/api/files.py index c5ea921671..7b5b45a7cb 100644 --- a/src/api/files.py +++ b/src/api/files.py @@ -8,6 +8,9 @@ """ import json +import re +from datetime import UTC, datetime +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from fastapi import Depends, HTTPException, Query from fastapi.responses import JSONResponse @@ -18,6 +21,65 @@ logger = get_logger(__name__) +# fromisoformat also accepts basic forms (e.g. "20240101"), which OpenSearch's +# strict_date_optional_time mapping on indexed_time rejects. +_EXTENDED_ISO_RE = re.compile( + r"^\d{4}-\d{2}-\d{2}" + r"(T\d{2}:\d{2}(:\d{2}(\.\d{1,6})?)?)?" + r"(Z|[+-]\d{2}:\d{2})?$" +) + + +CREATED_AFTER_DESCRIPTION = ( + "Inclusive lower bound on indexed_time (ISO 8601). " + "A value with no time is UTC midnight at the start of that day. " + "Documents with no indexed_time are excluded." +) +CREATED_BEFORE_DESCRIPTION = ( + "Inclusive upper bound on indexed_time (ISO 8601). " + "A value with no time is UTC midnight at the start of that day, " + "so pass the next day to include the chosen day. " + "Documents with no indexed_time are excluded." +) + + +def _parse_iso_timestamp(value: str, param_name: str) -> datetime: + """Parse an extended ISO 8601 timestamp, raising HTTP 422 when it is not one.""" + detail = f"{param_name} must be a valid ISO 8601 timestamp" + if not _EXTENDED_ISO_RE.match(value): + raise HTTPException(status_code=422, detail=detail) + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError as err: + raise HTTPException(status_code=422, detail=detail) from err + # indexed_time is stored in UTC and OpenSearch reads offset-less values as UTC. + return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC) + + +def _validate_iso_timestamp(value: str | None, param_name: str) -> None: + """Raise HTTP 422 when *value* is present but is not a valid ISO 8601 timestamp.""" + if isinstance(value, str): + _parse_iso_timestamp(value, param_name) + + +def _validate_date_range(created_after: str | None, created_before: str | None) -> None: + """Raise HTTP 422 for a malformed bound or a range whose start is after its end.""" + after = ( + _parse_iso_timestamp(created_after, "created_after") + if isinstance(created_after, str) + else None + ) + before = ( + _parse_iso_timestamp(created_before, "created_before") + if isinstance(created_before, str) + else None + ) + if after and before and after > before: + raise HTTPException( + status_code=422, + detail="created_after must be earlier than or equal to created_before", + ) + def _parse_after_key(after_key: str | None) -> dict | None: """Parse a JSON-encoded composite cursor. @@ -52,10 +114,13 @@ async def list_files( search: str | None = Query(None, description="Search filename"), after_key: str | None = Query(None, description="Composite pagination cursor (JSON-encoded)"), data_sources: list[str] | None = Query(None, description="Filename whitelist (repeatable)"), + created_after: str | None = Query(None, description=CREATED_AFTER_DESCRIPTION), + created_before: str | None = Query(None, description=CREATED_BEFORE_DESCRIPTION), file_service=Depends(get_file_service_v2), user: User = Depends(get_current_user), ): """List ingested files with composite-aggregation pagination, filtering, and sorting.""" + _validate_date_range(created_after, created_before) parsed_after_key = _parse_after_key(after_key) try: @@ -72,6 +137,8 @@ async def list_files( search=search, after_key=parsed_after_key, data_sources=data_sources, + created_after=created_after, + created_before=created_before, ) return JSONResponse(result) except Exception as e: @@ -97,10 +164,13 @@ async def search_files( owner: list[str] | None = Query(None, description="Filter by owner (repeatable)"), after_key: str | None = Query(None, description="Composite pagination cursor (JSON-encoded)"), data_sources: list[str] | None = Query(None, description="Filename whitelist (repeatable)"), + created_after: str | None = Query(None, description=CREATED_AFTER_DESCRIPTION), + created_before: str | None = Query(None, description=CREATED_BEFORE_DESCRIPTION), file_service=Depends(get_file_service_v2), user: User = Depends(get_current_user), ): """Search files by name with fuzzy/partial matching (composite pagination).""" + _validate_date_range(created_after, created_before) parsed_after_key = _parse_after_key(after_key) try: @@ -115,6 +185,8 @@ async def search_files( owner=owner, after_key=parsed_after_key, data_sources=data_sources, + created_after=created_after, + created_before=created_before, ) return JSONResponse(result) except Exception as e: @@ -127,3 +199,46 @@ async def search_files( {"error": "Failed to search files"}, status_code=500, ) + + +async def upload_histogram( + start: str = Query(..., description="Window start (ISO 8601, inclusive)"), + end: str = Query(..., description="Window end (ISO 8601, inclusive)"), + tz: str = Query("UTC", description="IANA time zone that defines calendar days"), + connector_type: list[str] | None = Query( + None, description="Filter by connector type (repeatable)" + ), + mimetype: list[str] | None = Query(None, description="Filter by MIME type (repeatable)"), + owner: list[str] | None = Query(None, description="Filter by owner (repeatable)"), + data_sources: list[str] | None = Query(None, description="Filename whitelist (repeatable)"), + file_service=Depends(get_file_service_v2), + user: User = Depends(get_current_user), +): + """Count distinct uploaded files per calendar day, plus the all-time total.""" + if _parse_iso_timestamp(start, "start") > _parse_iso_timestamp(end, "end"): + raise HTTPException(status_code=422, detail="start must be earlier than or equal to end") + try: + ZoneInfo(tz) + except (ZoneInfoNotFoundError, ValueError) as err: + raise HTTPException(status_code=422, detail="tz must be an IANA time zone") from err + + try: + result = await file_service.upload_histogram( + user_id=user.user_id, + jwt_token=user.jwt_token, + start=start, + end=end, + time_zone=tz, + connector_type=connector_type, + mimetype=mimetype, + owner=owner, + data_sources=data_sources, + ) + return JSONResponse(result) + except Exception as e: + logger.error("Failed to build upload histogram", error=str(e)) + from utils.opensearch_utils import AUTH_ERROR_MESSAGE, is_opensearch_auth_error + + if is_opensearch_auth_error(e): + return JSONResponse({"error": AUTH_ERROR_MESSAGE}, status_code=401) + return JSONResponse({"error": "Failed to build upload histogram"}, status_code=500) diff --git a/src/api/knowledge_filter.py b/src/api/knowledge_filter.py index 1f9e1588d0..b2d66ccbd7 100644 --- a/src/api/knowledge_filter.py +++ b/src/api/knowledge_filter.py @@ -1,12 +1,15 @@ import json +import re import uuid from datetime import UTC, datetime from typing import Any +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError -from fastapi import Depends, Request +from fastapi import Depends, HTTPException, Request from fastapi.responses import JSONResponse from pydantic import BaseModel +from api.files import _validate_date_range from dependencies import ( get_knowledge_filter_service, get_monitor_service, @@ -40,10 +43,79 @@ def _error_response(result: dict) -> JSONResponse: return JSONResponse(result, status_code=500) +_CALENDAR_DAY_RE = re.compile(r"[0-9]{4}-[0-9]{2}-[0-9]{2}", re.ASCII) +_DATE_BOUND_KEYS = ("created_after", "created_before") +_DATE_CONTEXT_KEYS = ("from", "to", "timezone") + + +def _validate_calendar_day(value: Any, name: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{name} must be a YYYY-MM-DD string") + # Exact zero-padded form only: the range check compares these as strings. + if not _CALENDAR_DAY_RE.fullmatch(value): + raise ValueError(f"{name} must be a valid YYYY-MM-DD date") + try: + datetime.strptime(value, "%Y-%m-%d") + except ValueError as err: + raise ValueError(f"{name} must be a valid YYYY-MM-DD date") from err + return value + + +def _normalize_date_range(data: dict, filters: dict) -> tuple[dict, dict | None]: + """Validate the optional saved date bounds and their display context. + + Returns ``(bounds, context)``. ``bounds`` holds only the UTC bounds that are + set; ``context`` is the calendar days + IANA timezone used to restore the + picker, kept outside ``filters`` so it never reaches retrieval. + Raises ValueError for any malformed value. + """ + bounds: dict[str, str] = {} + for key in _DATE_BOUND_KEYS: + value = filters.get(key) + if value is None or value == "": + continue + if not isinstance(value, str): + raise ValueError(f"{key} must be an ISO 8601 string") + bounds[key] = value + try: + _validate_date_range(bounds.get("created_after"), bounds.get("created_before")) + except HTTPException as err: + raise ValueError(str(err.detail)) from err + + raw_context = data.get("date_range_context") + if raw_context is None: + return bounds, None + if not bounds: + # Context without bounds would describe a range that is not applied. + return bounds, None + if not isinstance(raw_context, dict): + raise ValueError("date_range_context must be an object") + context: dict[str, str] = {} + for key in ("from", "to"): + if raw_context.get(key) is not None: + context[key] = _validate_calendar_day(raw_context[key], f"date_range_context.{key}") + if "from" in context and "to" in context and context["from"] > context["to"]: + raise ValueError("date_range_context.from must be on or before date_range_context.to") + tz = raw_context.get("timezone") + if tz is not None: + try: + if not isinstance(tz, str) or not tz: + raise ValueError + ZoneInfo(tz) + except (ValueError, ZoneInfoNotFoundError, OSError) as err: + raise ValueError("date_range_context.timezone must be a valid IANA timezone") from err + context["timezone"] = tz + return bounds, context or None + + def normalize_query_data(query_data: str | dict) -> str: """ Normalize query_data to ensure all required fields exist with defaults. This prevents frontend crashes when API-created filters have incomplete data. + + Optional ``created_after``/``created_before`` bounds are validated and kept + inside ``filters``; filters saved without dates stay unrestricted. + Raises ValueError for malformed date data. """ if isinstance(query_data, str): try: @@ -60,6 +132,8 @@ def normalize_query_data(query_data: str | dict) -> str: "owners": filters.get("owners", ["*"]), "connector_types": filters.get("connector_types", ["*"]), } + bounds, date_context = _normalize_date_range(data, filters) + normalized_filters.update(bounds) normalized = { "query": data.get("query", ""), @@ -69,6 +143,8 @@ def normalize_query_data(query_data: str | dict) -> str: "color": data.get("color", "zinc"), "icon": data.get("icon", "filter"), } + if date_context: + normalized["date_range_context"] = date_context return json.dumps(normalized) @@ -113,9 +189,9 @@ async def create_knowledge_filter( try: normalized_query_data = normalize_query_data(body.queryData) - except Exception as e: - logger.error(f"Failed to normalize query_data: {e}") - return JSONResponse({"error": f"Invalid queryData format: {str(e)}"}, status_code=400) + except Exception: + logger.error("Failed to normalize query_data", exc_info=True) + return JSONResponse({"error": "Invalid queryData format"}, status_code=400) jwt_token = user.jwt_token @@ -195,6 +271,13 @@ async def update_knowledge_filter( existing_filter = existing_result["filter"] + query_data = body.queryData if body.queryData is not None else existing_filter["query_data"] + try: + normalized_query_data = normalize_query_data(query_data) + except Exception: + logger.error("Failed to normalize query_data", exc_info=True) + return JSONResponse({"error": "Invalid queryData format"}, status_code=400) + delete_result = await knowledge_filter_service.delete_knowledge_filter( filter_id, user_id=user.user_id, jwt_token=jwt_token ) @@ -203,13 +286,6 @@ async def update_knowledge_filter( {"error": "Failed to delete existing knowledge filter"}, status_code=500 ) - query_data = body.queryData if body.queryData is not None else existing_filter["query_data"] - try: - normalized_query_data = normalize_query_data(query_data) - except Exception as e: - logger.error(f"Failed to normalize query_data: {e}") - return JSONResponse({"error": f"Invalid queryData format: {str(e)}"}, status_code=400) - updated_filter = { "id": filter_id, "name": body.name if body.name is not None else existing_filter["name"], diff --git a/src/api/search.py b/src/api/search.py index 4ebdea7642..62a26dca05 100644 --- a/src/api/search.py +++ b/src/api/search.py @@ -1,27 +1,29 @@ -from typing import Any, Dict +from typing import Any -from fastapi import Depends -from pydantic import BaseModel, Field +from fastapi import Depends, HTTPException from fastapi.responses import JSONResponse -from utils.logging_config import get_logger -from utils.opensearch_utils import OpenSearchDiskSpaceError, DISK_SPACE_ERROR_MESSAGE +from pydantic import BaseModel, Field +from api.files import _validate_date_range from dependencies import ( get_search_service, get_session_manager, - get_current_user, require_permission, ) from session_manager import User +from utils.logging_config import get_logger +from utils.opensearch_utils import DISK_SPACE_ERROR_MESSAGE, OpenSearchDiskSpaceError logger = get_logger(__name__) class SearchBody(BaseModel): query: str - filters: Dict[str, Any] = Field(default_factory=dict) + filters: dict[str, Any] = Field(default_factory=dict) limit: int = 10 scoreThreshold: float = Field(default=0, alias="scoreThreshold") + created_after: str | None = None + created_before: str | None = None model_config = {"populate_by_name": True} @@ -33,6 +35,18 @@ async def search( user: User = Depends(require_permission("search:use")), ): """Search for documents""" + merged_filters: dict[str, Any] = dict(body.filters) + if body.created_after: + merged_filters["created_after"] = body.created_after + if body.created_before: + merged_filters["created_before"] = body.created_before + for bound in ("created_after", "created_before"): + value = merged_filters.get(bound) + if value is not None and not isinstance(value, str): + raise HTTPException(status_code=422, detail=f"{bound} must be a string") + # Validate after merging: the bounds can also arrive inside `filters`. + _validate_date_range(merged_filters.get("created_after"), merged_filters.get("created_before")) + try: jwt_token = user.jwt_token @@ -50,7 +64,7 @@ async def search( body.query, user_id=user.user_id, jwt_token=jwt_token, - filters=body.filters, + filters=merged_filters if merged_filters else None, limit=body.limit, score_threshold=body.scoreThreshold, ) @@ -59,10 +73,7 @@ async def search( return JSONResponse({"error": DISK_SPACE_ERROR_MESSAGE}, status_code=507) except Exception as e: error_msg = str(e) - if ( - "AuthenticationException" in error_msg - or "access denied" in error_msg.lower() - ): + if "AuthenticationException" in error_msg or "access denied" in error_msg.lower(): return JSONResponse({"error": error_msg}, status_code=403) else: return JSONResponse({"error": error_msg}, status_code=500) diff --git a/src/api/v1/_filter_resolution.py b/src/api/v1/_filter_resolution.py index 417e306cac..c4a7ad5c45 100644 --- a/src/api/v1/_filter_resolution.py +++ b/src/api/v1/_filter_resolution.py @@ -15,13 +15,14 @@ from fastapi import HTTPException _FILTER_DIMENSIONS = ("data_sources", "document_types", "owners", "connector_types") +_DATE_BOUNDS = ("created_after", "created_before") -def _strip_wildcards(filters: dict[str, Any] | None) -> dict[str, list[str]]: - """Keep only filter dimensions that contain concrete values.""" +def _strip_wildcards(filters: dict[str, Any] | None) -> dict[str, Any]: + """Keep concrete filter dimensions and any saved date bounds.""" if not filters: return {} - cleaned: dict[str, list[str]] = {} + cleaned: dict[str, Any] = {} for key in _FILTER_DIMENSIONS: values = filters.get(key) if not values or not isinstance(values, list): @@ -29,6 +30,10 @@ def _strip_wildcards(filters: dict[str, Any] | None) -> dict[str, list[str]]: if "*" in values: continue cleaned[key] = values + for key in _DATE_BOUNDS: + value = filters.get(key) + if isinstance(value, str) and value: + cleaned[key] = value return cleaned diff --git a/src/api/v2/files.py b/src/api/v2/files.py index cec3b70d78..d7c9338250 100644 --- a/src/api/v2/files.py +++ b/src/api/v2/files.py @@ -9,6 +9,7 @@ from fastapi import Depends, Query +from api.files import CREATED_AFTER_DESCRIPTION, CREATED_BEFORE_DESCRIPTION from api.files import list_files as _internal_list_files from api.files import search_files as _internal_search_files from dependencies import get_file_service_v2, require_api_key_permission @@ -30,6 +31,8 @@ async def list_files( search: str | None = Query(None, description="Search filename"), after_key: str | None = Query(None, description="Composite pagination cursor (JSON-encoded)"), data_sources: list[str] | None = Query(None, description="Filename whitelist (repeatable)"), + created_after: str | None = Query(None, description=CREATED_AFTER_DESCRIPTION), + created_before: str | None = Query(None, description=CREATED_BEFORE_DESCRIPTION), file_service=Depends(get_file_service_v2), user: User = Depends(require_api_key_permission("knowledge:read:own")), ): @@ -45,6 +48,8 @@ async def list_files( search=search, after_key=after_key, data_sources=data_sources, + created_after=created_after, + created_before=created_before, file_service=file_service, user=user, ) @@ -61,6 +66,8 @@ async def search_files( owner: list[str] | None = Query(None, description="Filter by owner (repeatable)"), after_key: str | None = Query(None, description="Composite pagination cursor (JSON-encoded)"), data_sources: list[str] | None = Query(None, description="Filename whitelist (repeatable)"), + created_after: str | None = Query(None, description=CREATED_AFTER_DESCRIPTION), + created_before: str | None = Query(None, description=CREATED_BEFORE_DESCRIPTION), file_service=Depends(get_file_service_v2), user: User = Depends(require_api_key_permission("knowledge:read:own")), ): @@ -74,6 +81,8 @@ async def search_files( owner=owner, after_key=after_key, data_sources=data_sources, + created_after=created_after, + created_before=created_before, file_service=file_service, user=user, ) diff --git a/src/app/routes/internal.py b/src/app/routes/internal.py index efe5167d14..bf3f05c407 100644 --- a/src/app/routes/internal.py +++ b/src/app/routes/internal.py @@ -161,6 +161,9 @@ def register_internal_routes(app: FastAPI): # File listing/search endpoints app.add_api_route("/files", files.list_files, methods=["GET"], tags=["internal"]) app.add_api_route("/files/search", files.search_files, methods=["GET"], tags=["internal"]) + app.add_api_route( + "/files/upload-histogram", files.upload_histogram, methods=["GET"], tags=["internal"] + ) # Knowledge Filter endpoints app.add_api_route( diff --git a/src/services/chat_service.py b/src/services/chat_service.py index 7b139cc48e..68d28e230c 100644 --- a/src/services/chat_service.py +++ b/src/services/chat_service.py @@ -6,6 +6,7 @@ from config.settings import LANGFLOW_CHAT_FLOW_ID, LANGFLOW_URL, NUDGES_FLOW_ID, clients from utils.langflow_utils import fence_untrusted_text from utils.logging_config import get_logger +from utils.opensearch_queries import build_filter_clauses logger = get_logger(__name__) @@ -171,26 +172,7 @@ async def langflow_chat( # Build the complete filter expression like the search service does filter_expression: dict[str, Any] = {} if filters: - filter_clauses = [] - # Map frontend filter names to backend field names - field_mapping = { - "data_sources": "filename", - "document_types": "mimetype", - "owners": "owner", - "connector_types": "connector_type", - } - - for filter_key, values in filters.items(): - if values is not None and isinstance(values, list) and len(values) > 0: - # Map frontend key to backend field name - field_name = field_mapping.get(filter_key, filter_key) - - if len(values) == 1: - # Single value filter - filter_clauses.append({"term": {field_name: values[0]}}) - else: - # Multiple values filter - filter_clauses.append({"terms": {field_name: values}}) + filter_clauses = build_filter_clauses(filters) if filter_clauses: filter_expression["filter"] = filter_clauses @@ -298,29 +280,9 @@ async def langflow_nudges_chat( # Build the complete filter expression like the chat service does filter_expression: dict[str, Any] = {} has_user_filters = False - filter_clauses = [] + filter_clauses = build_filter_clauses(filters) if filters: - # Map frontend filter names to backend field names - field_mapping = { - "data_sources": "filename", - "document_types": "mimetype", - "owners": "owner", - "connector_types": "connector_type", - } - - for filter_key, values in filters.items(): - if values is not None and isinstance(values, list) and len(values) > 0: - # Map frontend key to backend field name - field_name = field_mapping.get(filter_key, filter_key) - - if len(values) == 1: - # Single value filter - filter_clauses.append({"term": {field_name: values[0]}}) - else: - # Multiple values filter - filter_clauses.append({"terms": {field_name: values}}) - if filter_clauses: has_user_filters = True diff --git a/src/services/document_index_writer.py b/src/services/document_index_writer.py index 0a25a991c3..476c5bf5d2 100644 --- a/src/services/document_index_writer.py +++ b/src/services/document_index_writer.py @@ -67,6 +67,7 @@ class DocumentIndexContext: parser: str | None = None chunk_size: int | None = None chunk_overlap: int | None = None + indexed_time: str | None = None @dataclass @@ -131,7 +132,7 @@ async def index_chunks( dimensions=dimensions, ) - now = datetime.datetime.now(datetime.UTC).isoformat() + now = context.indexed_time or datetime.datetime.now(datetime.UTC).isoformat() bulk_body: list[dict[str, Any]] = [] for chunk in chunks: if len(chunk.vector) != dimensions: diff --git a/src/services/file_service_v2.py b/src/services/file_service_v2.py index ce41326d8c..4401d7db98 100644 --- a/src/services/file_service_v2.py +++ b/src/services/file_service_v2.py @@ -21,6 +21,7 @@ from config.settings import get_index_name from utils.logging_config import get_logger +from utils.opensearch_queries import build_date_range_query logger = get_logger(__name__) @@ -55,6 +56,8 @@ async def list_files( search: str | None = None, after_key: dict | None = None, data_sources: list[str] | None = None, + created_after: str | None = None, + created_before: str | None = None, ) -> dict[str, Any]: """ List files with server-side pagination via composite aggregation. @@ -71,7 +74,14 @@ async def list_files( opensearch_client = self.session_manager.get_user_opensearch_client(user_id, jwt_token) query = self._build_filter_query( - user_id, connector_type, mimetype, owner, search, data_sources + user_id, + connector_type, + mimetype, + owner, + search, + data_sources, + created_after, + created_before, ) total, is_approximate = await self._get_file_count(opensearch_client, query) @@ -141,6 +151,66 @@ async def list_files( "after_key": next_after_key, } + async def upload_histogram( + self, + user_id: str, + jwt_token: str = None, + start: str = "", + end: str = "", + time_zone: str = "UTC", + connector_type: list[str] | None = None, + mimetype: list[str] | None = None, + owner: list[str] | None = None, + data_sources: list[str] | None = None, + ) -> dict[str, Any]: + """Distinct files per calendar day (in *time_zone*) within [start, end], plus the all-time total.""" + opensearch_client = self.session_manager.get_user_opensearch_client(user_id, jwt_token) + query = self._build_filter_query( + user_id, connector_type, mimetype, owner, None, data_sources + ) + # The window is a filter sub-agg rather than part of the query so the + # same request also yields the all-time total. + body = { + "size": 0, + "query": query, + "aggs": { + "total": {"cardinality": {"field": "filename", "precision_threshold": 3000}}, + "window": { + "filter": build_date_range_query("indexed_time", gte=start, lte=end), + "aggs": { + "days": { + "date_histogram": { + "field": "indexed_time", + "calendar_interval": "day", + "time_zone": time_zone, + "format": "yyyy-MM-dd", + "min_doc_count": 1, + }, + "aggs": { + "files": { + "cardinality": { + "field": "filename", + "precision_threshold": 3000, + } + } + }, + } + }, + }, + }, + } + result = await opensearch_client.search(index=get_index_name(), body=body) + aggs = result.get("aggregations", {}) + buckets = aggs.get("window", {}).get("days", {}).get("buckets", []) + return { + "total": aggs.get("total", {}).get("value", 0), + "days": { + b["key_as_string"]: b.get("files", {}).get("value", 0) + for b in buckets + if b.get("key_as_string") + }, + } + async def _list_files_by_chunk_count( self, opensearch_client: Any, @@ -207,6 +277,8 @@ async def search_files( owner: list[str] | None = None, after_key: dict | None = None, data_sources: list[str] | None = None, + created_after: str | None = None, + created_before: str | None = None, ) -> dict[str, Any]: """Search files by name with fuzzy/prefix matching.""" return await self.list_files( @@ -222,6 +294,8 @@ async def search_files( search=query, after_key=after_key, data_sources=data_sources, + created_after=created_after, + created_before=created_before, ) def _build_filter_query( @@ -232,6 +306,8 @@ def _build_filter_query( owner: list[str] | None = None, search: str | None = None, data_sources: list[str] | None = None, + created_after: str | None = None, + created_before: str | None = None, ) -> dict[str, Any]: def _effective_values(values: str | list[str] | None) -> list[str]: if values is None: @@ -279,6 +355,12 @@ def _effective_values(values: str | list[str] | None) -> list[str]: ) filter_clauses.append(clause) + date_range_clause = build_date_range_query( + "indexed_time", gte=created_after, lte=created_before + ) + if date_range_clause is not None: + filter_clauses.append(date_range_clause) + if search: must.append( { diff --git a/src/services/indexed_time_backfill.py b/src/services/indexed_time_backfill.py new file mode 100644 index 0000000000..2c9ea2b779 --- /dev/null +++ b/src/services/indexed_time_backfill.py @@ -0,0 +1,172 @@ +"""One-shot repair: give every chunk of an ingest run the run's earliest indexed_time. + +Before per-run timestamps, each Langflow ingest-callback batch stamped its own +``indexed_time``, so one file could carry several values and the Date Uploaded +sort (grouped by ``(indexed_time, filename)``) listed it as several rows. + +Scope: this makes ``indexed_time`` uniform within an ``ingest_run_id``. It does +not merge different runs that share a filename, and chunks without an +``ingest_run_id`` are not touched. + +Eligibility: a split run is deferred while its newest chunk (``max_t``) is +within ``LANGFLOW_INGEST_CALLBACK_TTL_SECONDS`` of now. This is an inactivity +precaution based on the configured token lifetime, not the run's actual token +expiry: expiry stops new callback requests being accepted, but a request +accepted just before expiry can still finish writing afterwards, so the cutoff +deliberately keys off the last observed write rather than ``min_t + TTL``. It +lowers the chance of rewriting a run that is still being ingested; it is not a +guarantee that no writes are in flight (that would need completion or +in-flight-write tracking, which is out of scope here). + +Deferred runs, like any failed or partial pass, are retried on the next boot; +the migration is only marked done after a complete clean pass. +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any + +from config.settings import LANGFLOW_INGEST_CALLBACK_TTL_SECONDS, clients, get_index_name +from db.migrations_runtime import _already_done, _mark_done +from utils.logging_config import get_logger + +logger = get_logger(__name__) + +MIGRATION_NAME = "indexed_time_per_ingest_run_v1" +_PAGE_SIZE = 500 + +_FIX_SCRIPT = """ +def t = params.times[ctx._source.ingest_run_id]; +if (t == null || t.equals(ctx._source.indexed_time)) { ctx.op = 'noop'; } +else { ctx._source.indexed_time = t; } +""" + + +def _scan_body(after_key: dict[str, Any] | None) -> dict[str, Any]: + composite: dict[str, Any] = { + "size": _PAGE_SIZE, + "sources": [{"run": {"terms": {"field": "ingest_run_id"}}}], + } + if after_key: + composite["after"] = after_key + return { + "size": 0, + "query": {"exists": {"field": "ingest_run_id"}}, + "aggs": { + "runs": { + "composite": composite, + "aggs": { + "min_t": {"min": {"field": "indexed_time"}}, + "max_t": {"max": {"field": "indexed_time"}}, + "earliest": { + "top_hits": { + "size": 1, + "_source": ["indexed_time"], + "sort": [{"indexed_time": {"order": "asc"}}], + } + }, + }, + } + }, + } + + +def _select_runs(buckets: list[dict[str, Any]], cutoff_ms: float) -> tuple[dict[str, str], int]: + """Return ({run_id: earliest indexed_time string}, deferred_count).""" + times: dict[str, str] = {} + deferred = 0 + for bucket in buckets: + min_t = bucket.get("min_t", {}).get("value") + max_t = bucket.get("max_t", {}).get("value") + if min_t is None or min_t == max_t: + continue + if max_t >= cutoff_ms: # recent write activity: see module docstring + deferred += 1 + continue + hits = bucket.get("earliest", {}).get("hits", {}).get("hits", []) + target = hits[0].get("_source", {}).get("indexed_time") if hits else None + if target: + times[str(bucket["key"]["run"])] = target + return times, deferred + + +async def _run_backfill(client: Any, index_name: str) -> dict[str, int] | None: + """Scan and repair. Returns counts, or None when the pass was incomplete.""" + stats = {"fixed": 0, "deferred": 0, "updated": 0, "conflicts": 0, "failures": 0} + if not await client.indices.exists(index=index_name): + return stats + + cutoff_ms = (time.time() - LANGFLOW_INGEST_CALLBACK_TTL_SECONDS) * 1000 + after_key: dict[str, Any] | None = None + while True: + response = await client.search(index=index_name, body=_scan_body(after_key)) + if response.get("timed_out") or response.get("_shards", {}).get("failed", 0): + logger.warning("Indexed-time backfill scan was partial", index_name=index_name) + return None + agg = response.get("aggregations", {}).get("runs", {}) + times, deferred = _select_runs(agg.get("buckets", []), cutoff_ms) + stats["deferred"] += deferred + + if times: + result = await client.update_by_query( + index=index_name, + body={ + "query": {"terms": {"ingest_run_id": list(times)}}, + "script": { + "lang": "painless", + "source": _FIX_SCRIPT, + "params": {"times": times}, + }, + }, + conflicts="proceed", + wait_for_completion=True, + request_timeout=300, + ) + if result.get("timed_out"): + logger.warning("Indexed-time backfill update timed out", index_name=index_name) + return None + stats["fixed"] += len(times) + stats["updated"] += int(result.get("updated", 0)) + stats["conflicts"] += int(result.get("version_conflicts", 0)) + stats["failures"] += len(result.get("failures", [])) + + after_key = agg.get("after_key") + if not after_key: + break + + if stats["updated"] > 0: + await client.indices.refresh(index=index_name) + return stats + + +async def run_indexed_time_backfill() -> None: + """Run the repair once; never raises (except cancellation) and never blocks startup.""" + try: + import db.engine as db_engine + + session_factory = db_engine.SessionLocal + client = clients.opensearch + if session_factory is None or client is None: + return + + async with session_factory() as session: + if await _already_done(session, MIGRATION_NAME): + return + + stats = await _run_backfill(client, get_index_name()) + if stats is None or stats["failures"] or stats["conflicts"] or stats["deferred"]: + logger.warning( + "Indexed-time backfill incomplete; will retry on next boot", + **(stats or {}), + ) + return + + await _mark_done(session, MIGRATION_NAME, notes=str(stats)) + await session.commit() + logger.info("Indexed-time backfill complete", **stats) + except asyncio.CancelledError: + raise + except Exception as e: + logger.error("Indexed-time backfill failed; will retry on next boot", error=str(e)) diff --git a/src/services/langflow_file_service.py b/src/services/langflow_file_service.py index d980a0879a..8d435cac6e 100644 --- a/src/services/langflow_file_service.py +++ b/src/services/langflow_file_service.py @@ -1,4 +1,5 @@ import asyncio +import datetime import io import json import time @@ -302,6 +303,7 @@ def _configure_ingest_callback( chunk_size=chunk_size, chunk_overlap=chunk_overlap, connector_file_id=connector_file_id, + indexed_time=datetime.datetime.now(datetime.UTC).isoformat(), ) token = self.ingest_token_service.create_token(context) logger.info( diff --git a/src/services/langflow_ingest_token_service.py b/src/services/langflow_ingest_token_service.py index c6502e22d4..2cb22a4d64 100644 --- a/src/services/langflow_ingest_token_service.py +++ b/src/services/langflow_ingest_token_service.py @@ -2,6 +2,8 @@ from __future__ import annotations +import dataclasses +import datetime import time import uuid from typing import Any @@ -83,6 +85,11 @@ def __init__(self, secret: str | None = None, ttl_seconds: int | None = None): def create_token(self, context: DocumentIndexContext) -> str: now = int(time.time()) + if context.indexed_time is None: + context = dataclasses.replace( + context, + indexed_time=datetime.datetime.fromtimestamp(now, datetime.UTC).isoformat(), + ) jti = str(uuid.uuid4()) payload = { "aud": self.audience, @@ -169,6 +176,7 @@ def _context_to_payload(context: DocumentIndexContext) -> dict[str, Any]: "chunk_size": context.chunk_size, "chunk_overlap": context.chunk_overlap, "connector_file_id": context.connector_file_id, + "indexed_time": context.indexed_time, } @staticmethod @@ -206,4 +214,5 @@ def _payload_to_context(payload: dict[str, Any]) -> DocumentIndexContext: chunk_size=payload.get("chunk_size"), chunk_overlap=payload.get("chunk_overlap"), connector_file_id=payload.get("connector_file_id"), + indexed_time=payload.get("indexed_time"), ) diff --git a/src/services/search_service.py b/src/services/search_service.py index 56b64789b7..83985c1058 100644 --- a/src/services/search_service.py +++ b/src/services/search_service.py @@ -25,6 +25,7 @@ get_embedding_space_id, ) from utils.logging_config import get_logger +from utils.opensearch_queries import build_filter_clauses logger = get_logger(__name__) @@ -376,30 +377,9 @@ async def _search(self, query: str, embedding_model: str) -> dict[str, Any]: if not is_wildcard_match_all: # Build filter clauses first so we can use them in model detection - filter_clauses: list[dict[str, Any]] = [] - if filters: - # Map frontend filter names to backend field names - field_mapping = { - "data_sources": "filename", - "document_types": "mimetype", - "owners": "owner", - "connector_types": "connector_type", - } - - for filter_key, values in filters.items(): - if values is not None and isinstance(values, list): - # Map frontend key to backend field name - field_name = field_mapping.get(filter_key, filter_key) - - if len(values) == 0: - # Empty array means "match nothing" - use impossible filter - filter_clauses.append({"term": {field_name: "__IMPOSSIBLE_VALUE__"}}) - elif len(values) == 1: - # Single value filter - filter_clauses.append({"term": {field_name: values[0]}}) - else: - # Multiple values filter - filter_clauses.append({"terms": {field_name: values}}) + filter_clauses: list[dict[str, Any]] = build_filter_clauses( + filters, empty_list_matches_nothing=True + ) try: seen_spaces: set[str] = set() @@ -564,30 +544,7 @@ async def embed_with_space( ) else: # Wildcard query - no embedding needed - filter_clauses = [] - if filters: - # Map frontend filter names to backend field names - field_mapping = { - "data_sources": "filename", - "document_types": "mimetype", - "owners": "owner", - "connector_types": "connector_type", - } - - for filter_key, values in filters.items(): - if values is not None and isinstance(values, list): - # Map frontend key to backend field name - field_name = field_mapping.get(filter_key, filter_key) - - if len(values) == 0: - # Empty array means "match nothing" - use impossible filter - filter_clauses.append({"term": {field_name: "__IMPOSSIBLE_VALUE__"}}) - elif len(values) == 1: - # Single value filter - filter_clauses.append({"term": {field_name: values[0]}}) - else: - # Multiple values filter - filter_clauses.append({"terms": {field_name: values}}) + filter_clauses = build_filter_clauses(filters, empty_list_matches_nothing=True) # Build query body if is_wildcard_match_all: @@ -602,17 +559,17 @@ async def embed_with_space( knn_queries = [] for field_identity, embedding_vector in query_embeddings.items(): field_name = get_embedding_field_name(field_identity) - knn_queries.append( - { - "knn": { - field_name: { - "vector": embedding_vector, - "k": 50, - "num_candidates": 1000, - } - } - } - ) + knn_clause: dict[str, Any] = { + "vector": embedding_vector, + "k": 50, + "num_candidates": 1000, + } + # Filter inside the kNN clause so neighbors are chosen among + # in-range documents; the outer bool filter alone is a + # post-filter over the global top-k. + if filter_clauses: + knn_clause["filter"] = {"bool": {"filter": filter_clauses}} + knn_queries.append({"knn": {field_name: knn_clause}}) # KNN clauses naturally apply only to documents carrying their # vector field. Keep all other documents eligible for keyword @@ -630,7 +587,7 @@ async def embed_with_space( # Hybrid search (semantic + keyword) when embeddings are available; # keyword-only fallback when none succeeded. When falling back, bump # the multi_match boost so keyword scoring isn't artificially damped. - should_clauses = [] + should_clauses: list[dict[str, Any]] = [] if knn_queries: should_clauses.append( { @@ -711,6 +668,7 @@ async def embed_with_space( "allowed_users", "allowed_groups", "allowed_principal_labels", + "indexed_time", ], "size": limit, } @@ -872,6 +830,7 @@ async def embed_with_space( "allowed_users": source.get("allowed_users", []), "allowed_groups": source.get("allowed_groups", []), "allowed_principal_labels": source.get("allowed_principal_labels", []), + "indexed_time": source.get("indexed_time"), } ) diff --git a/src/services/startup_orchestrator.py b/src/services/startup_orchestrator.py index 779c9c017f..1dcd92a870 100644 --- a/src/services/startup_orchestrator.py +++ b/src/services/startup_orchestrator.py @@ -223,3 +223,7 @@ async def startup_tasks(services) -> None: # Update MCP server URLs and required global-var headers after flows and variables exist await _update_mcp_server_urls(services["langflow_mcp_service"]) + + from services.indexed_time_backfill import run_indexed_time_backfill + + await run_indexed_time_backfill() diff --git a/src/utils/opensearch_queries.py b/src/utils/opensearch_queries.py index d16dfb7b5b..c3d1932b99 100644 --- a/src/utils/opensearch_queries.py +++ b/src/utils/opensearch_queries.py @@ -116,3 +116,76 @@ def build_replace_filename_query(filename: str, owner: str) -> dict: ] } } + + +def build_date_range_query( + field: str, gte: str | None = None, lte: str | None = None +) -> dict | None: + """ + Build a date range query clause for OpenSearch. + + Args: + field: The date field name to filter on (e.g. 'created_time') + gte: Lower bound ISO-8601 string (inclusive) + lte: Upper bound ISO-8601 string (inclusive) + + Returns: + A dict containing {"range": {field: {...}}} or None if both bounds are None. + """ + range_bounds: dict[str, str] = {} + if gte is not None: + range_bounds["gte"] = gte + if lte is not None: + range_bounds["lte"] = lte + if not range_bounds: + return None + return {"range": {field: range_bounds}} + + +FILTER_FIELD_MAPPING = { + "data_sources": "filename", + "document_types": "mimetype", + "owners": "owner", + "connector_types": "connector_type", +} + + +def build_filter_clauses( + filters: dict | None, *, empty_list_matches_nothing: bool = False +) -> list[dict]: + """ + Convert request filters into OpenSearch filter clauses. + + List-valued filters become term/terms clauses on the mapped field. + ``created_after``/``created_before`` become one inclusive ``indexed_time`` + range clause. + + Args: + filters: Frontend filter dict (may be None/empty). + empty_list_matches_nothing: When True an empty list matches no documents + (search behavior); when False it is ignored (chat/nudges behavior). + """ + clauses: list[dict] = [] + if not filters: + return clauses + + for filter_key, values in filters.items(): + if values is None or not isinstance(values, list): + continue + field_name = FILTER_FIELD_MAPPING.get(filter_key, filter_key) + if len(values) == 0: + if empty_list_matches_nothing: + clauses.append({"term": {field_name: "__IMPOSSIBLE_VALUE__"}}) + elif len(values) == 1: + clauses.append({"term": {field_name: values[0]}}) + else: + clauses.append({"terms": {field_name: values}}) + + date_clause = build_date_range_query( + "indexed_time", + gte=filters.get("created_after") or None, + lte=filters.get("created_before") or None, + ) + if date_clause is not None: + clauses.append(date_clause) + return clauses diff --git a/tests/unit/api/test_files_search_date_range.py b/tests/unit/api/test_files_search_date_range.py new file mode 100644 index 0000000000..103e6d514b --- /dev/null +++ b/tests/unit/api/test_files_search_date_range.py @@ -0,0 +1,101 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from api import files as internal_files +from api.v2 import files as v2_files +from services.file_service_v2 import FileServiceV2 + +USER = SimpleNamespace(user_id="u1", jwt_token="jwt") +AFTER = "2024-01-01T00:00:00Z" +BEFORE = "2024-01-31T23:59:59Z" + + +def _search_kwargs(**overrides): + kwargs = { + "q": "report", + "page": 1, + "page_size": 25, + "connector_type": None, + "mimetype": None, + "owner": None, + "after_key": None, + "data_sources": None, + "created_after": AFTER, + "created_before": BEFORE, + "user": USER, + } + kwargs.update(overrides) + return kwargs + + +def _service(): + service = MagicMock() + service.search_files = AsyncMock(return_value={"files": []}) + return service + + +@pytest.mark.asyncio +async def test_internal_search_forwards_date_bounds(): + service = _service() + response = await internal_files.search_files(**_search_kwargs(file_service=service)) + assert response.status_code == 200 + kwargs = service.search_files.await_args.kwargs + assert kwargs["created_after"] == AFTER + assert kwargs["created_before"] == BEFORE + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides", + [ + {"created_after": BEFORE, "created_before": AFTER}, + {"created_after": "20240101"}, + ], +) +async def test_internal_search_rejects_bad_bounds(overrides): + service = _service() + with pytest.raises(HTTPException) as exc: + await internal_files.search_files(**_search_kwargs(file_service=service, **overrides)) + assert exc.value.status_code == 422 + service.search_files.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_search_forwards_date_bounds(): + service = _service() + await v2_files.search_files(**_search_kwargs(file_service=service)) + kwargs = service.search_files.await_args.kwargs + assert kwargs["created_after"] == AFTER + assert kwargs["created_before"] == BEFORE + + +@pytest.mark.asyncio +async def test_service_search_passes_bounds_to_list_files(): + service = FileServiceV2(MagicMock()) + service.list_files = AsyncMock(return_value={"files": []}) + await service.search_files( + user_id="u1", query="report", created_after=AFTER, created_before=BEFORE + ) + kwargs = service.list_files.await_args.kwargs + assert kwargs["created_after"] == AFTER + assert kwargs["created_before"] == BEFORE + + +@pytest.mark.parametrize( + "handler", + [ + internal_files.list_files, + internal_files.search_files, + v2_files.list_files, + v2_files.search_files, + ], +) +def test_bounds_are_documented_as_inclusive(handler): + defaults = handler.__defaults__ or () + descriptions = [getattr(d, "description", "") or "" for d in defaults] + assert internal_files.CREATED_AFTER_DESCRIPTION in descriptions + assert internal_files.CREATED_BEFORE_DESCRIPTION in descriptions + assert "Inclusive" in internal_files.CREATED_BEFORE_DESCRIPTION diff --git a/tests/unit/api/test_files_timestamp_validation.py b/tests/unit/api/test_files_timestamp_validation.py new file mode 100644 index 0000000000..684e943129 --- /dev/null +++ b/tests/unit/api/test_files_timestamp_validation.py @@ -0,0 +1,165 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from api.files import _validate_date_range, _validate_iso_timestamp +from api.search import SearchBody +from api.search import search as search_endpoint +from api.v2.files import list_files as v2_list_files + + +@pytest.mark.parametrize( + "value", + [ + None, + "2024-01-01", + "2024-01-01T00:00", + "2024-01-01T00:00:00", + "2024-01-01T00:00:00.000Z", + "2024-01-31T23:59:59.999+05:30", + ], +) +def test_accepts_extended_iso_forms(value): + _validate_iso_timestamp(value, "created_after") + + +@pytest.mark.parametrize( + "value", + [ + "20240101", + "20240101T000000", + "2024-01-01T000000", + "2024-13-01", + "not-a-date", + "", + ], +) +def test_rejects_basic_and_invalid_forms(value): + with pytest.raises(HTTPException) as exc: + _validate_iso_timestamp(value, "created_after") + assert exc.value.status_code == 422 + assert "created_after" in exc.value.detail + + +@pytest.mark.parametrize( + ("created_after", "created_before"), + [ + (None, None), + ("2024-01-01", None), + (None, "2024-01-31"), + ("2024-01-01", "2024-01-31"), + ("2024-01-01T00:00:00Z", "2024-01-01T00:00:00Z"), + # 09:00 at +05:30 is 03:30 UTC, which is before 04:00 UTC. + ("2024-01-01T09:00:00+05:30", "2024-01-01T04:00:00Z"), + # An offset-less bound is read as UTC, so it compares against an aware one. + ("2024-01-01T00:00:00", "2024-01-01T00:00:00Z"), + ], +) +def test_accepts_ordered_ranges(created_after, created_before): + _validate_date_range(created_after, created_before) + + +@pytest.mark.parametrize( + ("created_after", "created_before"), + [ + ("2024-02-01", "2024-01-01"), + ("2024-01-01T00:00:01Z", "2024-01-01T00:00:00Z"), + # 05:00 UTC is after 09:00 at +05:30 (03:30 UTC). + ("2024-01-01T05:00:00Z", "2024-01-01T09:00:00+05:30"), + ], +) +def test_rejects_inverted_ranges(created_after, created_before): + with pytest.raises(HTTPException) as exc: + _validate_date_range(created_after, created_before) + assert exc.value.status_code == 422 + assert "created_after must be earlier" in exc.value.detail + + +@pytest.mark.parametrize( + ("created_after", "created_before", "bad_param"), + [ + ("20240101", None, "created_after"), + (None, "nope", "created_before"), + ], +) +def test_range_rejects_malformed_bound(created_after, created_before, bad_param): + with pytest.raises(HTTPException) as exc: + _validate_date_range(created_after, created_before) + assert exc.value.status_code == 422 + assert bad_param in exc.value.detail + + +def _call_v2_list_files(file_service, **overrides): + params = { + "page": 1, + "page_size": 25, + "sort_by": "filename", + "sort_order": "asc", + "connector_type": None, + "mimetype": None, + "owner": None, + "search": None, + "after_key": None, + "data_sources": None, + "created_after": None, + "created_before": None, + **overrides, + } + user = SimpleNamespace(user_id="user-1", jwt_token="jwt") + return v2_list_files(**params, file_service=file_service, user=user) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides", + [ + {"created_after": "not-a-date"}, + {"created_before": "20240101"}, + {"created_after": "2024-02-01", "created_before": "2024-01-01"}, + ], +) +async def test_v2_list_files_rejects_invalid_dates_before_querying(overrides): + file_service = SimpleNamespace(list_files=AsyncMock()) + with pytest.raises(HTTPException) as exc: + await _call_v2_list_files(file_service, **overrides) + assert exc.value.status_code == 422 + file_service.list_files.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_list_files_passes_valid_range_through(): + file_service = SimpleNamespace(list_files=AsyncMock(return_value={"files": []})) + await _call_v2_list_files(file_service, created_after="2024-01-01", created_before="2024-01-31") + kwargs = file_service.list_files.await_args.kwargs + assert kwargs["created_after"] == "2024-01-01" + assert kwargs["created_before"] == "2024-01-31" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"query": "q", "created_after": "2024-02-01", "created_before": "2024-01-01"}, + # Bounds smuggled through `filters` are validated too. + {"query": "q", "filters": {"created_after": "garbage"}}, + { + "query": "q", + "filters": {"created_after": "2024-02-01"}, + "created_before": "2024-01-01", + }, + ], +) +async def test_search_rejects_invalid_date_bounds(body): + search_service = SimpleNamespace(search=AsyncMock()) + user = SimpleNamespace(user_id="user-1", jwt_token="jwt") + with pytest.raises(HTTPException) as exc: + await search_endpoint( + SearchBody(**body), + search_service=search_service, + session_manager=None, + user=user, + ) + assert exc.value.status_code == 422 + search_service.search.assert_not_awaited() diff --git a/tests/unit/api/test_files_upload_histogram.py b/tests/unit/api/test_files_upload_histogram.py new file mode 100644 index 0000000000..1386c4c368 --- /dev/null +++ b/tests/unit/api/test_files_upload_histogram.py @@ -0,0 +1,110 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from api.files import upload_histogram +from services.file_service_v2 import FileServiceV2 + +USER = SimpleNamespace(user_id="u1", jwt_token="jwt") + + +def _call(file_service, **overrides): + params = { + "start": "2024-05-01T00:00:00Z", + "end": "2024-06-20T23:59:59Z", + "tz": "America/New_York", + "connector_type": None, + "mimetype": None, + "owner": None, + "data_sources": None, + "file_service": file_service, + "user": USER, + } + params.update(overrides) + return upload_histogram(**params) + + +@pytest.mark.asyncio +async def test_passes_window_timezone_and_filters_to_service(): + service = MagicMock() + service.upload_histogram = AsyncMock(return_value={"total": 2, "days": {"2024-06-01": 2}}) + + response = await _call(service, owner=["u1"]) + + assert response.status_code == 200 + kwargs = service.upload_histogram.await_args.kwargs + assert kwargs["time_zone"] == "America/New_York" + assert kwargs["start"] == "2024-05-01T00:00:00Z" + assert kwargs["owner"] == ["u1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides", + [ + {"start": "20240501"}, + {"end": "not-a-date"}, + {"start": "2024-07-01T00:00:00Z"}, + {"tz": "Mars/Olympus_Mons"}, + ], +) +async def test_rejects_bad_input(overrides): + service = MagicMock() + service.upload_histogram = AsyncMock() + with pytest.raises(HTTPException) as exc: + await _call(service, **overrides) + assert exc.value.status_code == 422 + service.upload_histogram.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_returns_500_when_service_fails(): + service = MagicMock() + service.upload_histogram = AsyncMock(side_effect=RuntimeError("down")) + response = await _call(service) + assert response.status_code == 500 + + +@pytest.mark.asyncio +async def test_service_buckets_by_local_day_and_keeps_all_time_total(): + client = MagicMock() + client.search = AsyncMock( + return_value={ + "aggregations": { + "total": {"value": 7}, + "window": { + "days": { + "buckets": [ + {"key_as_string": "2024-06-01", "files": {"value": 2}}, + {"key_as_string": "2024-06-03", "files": {"value": 5}}, + ] + } + }, + } + } + ) + session_manager = MagicMock() + session_manager.get_user_opensearch_client.return_value = client + + result = await FileServiceV2(session_manager).upload_histogram( + user_id="u1", + start="2024-05-01T00:00:00Z", + end="2024-06-20T23:59:59Z", + time_zone="Asia/Kolkata", + connector_type=["local"], + ) + + assert result == {"total": 7, "days": {"2024-06-01": 2, "2024-06-03": 5}} + body = client.search.await_args.kwargs["body"] + # The date window narrows only the histogram, never the all-time total. + assert "range" not in str(body["query"]) + assert {"term": {"connector_type": "local"}} in body["query"]["bool"]["filter"] + window = body["aggs"]["window"] + assert window["filter"] == { + "range": {"indexed_time": {"gte": "2024-05-01T00:00:00Z", "lte": "2024-06-20T23:59:59Z"}} + } + histogram = window["aggs"]["days"]["date_histogram"] + assert histogram["time_zone"] == "Asia/Kolkata" + assert histogram["calendar_interval"] == "day" diff --git a/tests/unit/api/test_knowledge_filter_date_range.py b/tests/unit/api/test_knowledge_filter_date_range.py new file mode 100644 index 0000000000..5ac43f1125 --- /dev/null +++ b/tests/unit/api/test_knowledge_filter_date_range.py @@ -0,0 +1,136 @@ +import json + +import pytest + +from api.knowledge_filter import normalize_query_data +from api.v1._filter_resolution import _strip_wildcards + +AFTER = "2026-10-07T07:00:00Z" +BEFORE = "2026-10-11T06:59:59.999Z" +CONTEXT = {"from": "2026-10-07", "to": "2026-10-10", "timezone": "America/Los_Angeles"} + + +def _norm(data): + return json.loads(normalize_query_data(data)) + + +def test_filter_without_dates_stays_unrestricted(): + out = _norm({"filters": {"document_types": ["application/pdf"]}}) + assert "created_after" not in out["filters"] + assert "created_before" not in out["filters"] + assert "date_range_context" not in out + + +def test_dates_and_context_round_trip_with_context_outside_filters(): + out = _norm( + { + "filters": {"created_after": AFTER, "created_before": BEFORE}, + "date_range_context": CONTEXT, + } + ) + assert out["filters"]["created_after"] == AFTER + assert out["filters"]["created_before"] == BEFORE + assert out["date_range_context"] == CONTEXT + assert "date_range_context" not in out["filters"] + # Idempotent through a second save + assert _norm(out) == out + + +def test_context_dropped_when_no_bounds(): + assert "date_range_context" not in _norm({"date_range_context": CONTEXT}) + + +@pytest.mark.parametrize( + "filters", + [ + {"created_after": 123}, + {"created_before": ["2026-10-07"]}, + {"created_after": "not-a-date"}, + {"created_after": BEFORE, "created_before": AFTER}, + ], +) +def test_invalid_bounds_rejected(filters): + with pytest.raises(ValueError): + normalize_query_data({"filters": filters}) + + +@pytest.mark.parametrize( + "context", + [ + "x", + {"from": "2026-13-01"}, + {"from": 5}, + {"from": "2026-10-10", "to": "2026-10-07"}, + {"timezone": "Not/AZone"}, + {"timezone": 7}, + ], +) +def test_invalid_context_rejected(context): + with pytest.raises(ValueError): + normalize_query_data({"filters": {"created_after": AFTER}, "date_range_context": context}) + + +def test_v1_resolution_keeps_date_bounds(): + out = _strip_wildcards( + { + "document_types": ["application/pdf"], + "owners": ["*"], + "created_after": AFTER, + "created_before": BEFORE, + } + ) + assert out == { + "document_types": ["application/pdf"], + "created_after": AFTER, + "created_before": BEFORE, + } + assert _strip_wildcards({"created_after": 5}) == {} + + +def test_unpadded_calendar_days_rejected(): + # "2026-2-9" would string-compare as later than "2026-10-10". + for context in ( + {"from": "2026-10-10", "to": "2026-2-9"}, + {"from": "2026-2-9"}, + {"to": "2026-10-1"}, + ): + with pytest.raises(ValueError): + normalize_query_data( + {"filters": {"created_after": AFTER}, "date_range_context": context} + ) + + +@pytest.mark.asyncio +async def test_rejected_update_neither_deletes_nor_recreates(): + from unittest.mock import AsyncMock, MagicMock + + from api.knowledge_filter import UpdateFilterBody, update_knowledge_filter + + service = MagicMock() + service.get_knowledge_filter = AsyncMock( + return_value={ + "success": True, + "filter": { + "id": "f1", + "name": "n", + "description": "", + "query_data": "{}", + "owner": "u", + }, + } + ) + service.delete_knowledge_filter = AsyncMock() + service.create_knowledge_filter = AsyncMock() + user = MagicMock(user_id="u", jwt_token="t") + + response = await update_knowledge_filter( + "f1", + UpdateFilterBody(queryData={"filters": {"created_after": 123}}), + knowledge_filter_service=service, + session_manager=MagicMock(), + user=user, + ) + + assert response.status_code == 400 + service.delete_knowledge_filter.assert_not_called() + service.create_knowledge_filter.assert_not_called() diff --git a/tests/unit/services/test_chat_service_date_filters.py b/tests/unit/services/test_chat_service_date_filters.py new file mode 100644 index 0000000000..f2e5e2523e --- /dev/null +++ b/tests/unit/services/test_chat_service_date_filters.py @@ -0,0 +1,83 @@ +"""ChatService forwards saved date ranges in the Langflow query-filter header.""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import agent +from auth_context import set_search_filters +from services.chat_service import ChatService + +HEADER = "X-LANGFLOW-GLOBAL-VAR-OPENRAG-QUERY-FILTER" +RANGE = { + "range": {"indexed_time": {"gte": "2026-10-07T07:00:00Z", "lte": "2026-10-11T06:59:59.999Z"}} +} +PDF = {"term": {"mimetype": "application/pdf"}} +SAMPLE_EXCLUSION = {"bool": {"must_not": [{"term": {"is_sample_data": "true"}}]}} +DATES = {"created_after": "2026-10-07T07:00:00Z", "created_before": "2026-10-11T06:59:59.999Z"} + + +@pytest.fixture +def captured(monkeypatch): + monkeypatch.setattr("services.chat_service.LANGFLOW_URL", "http://lf") + monkeypatch.setattr("services.chat_service.LANGFLOW_CHAT_FLOW_ID", "chat") + monkeypatch.setattr("services.chat_service.NUDGES_FLOW_ID", "nudges") + monkeypatch.setattr( + "services.chat_service.clients.ensure_langflow_client", + AsyncMock(return_value=MagicMock()), + ) + monkeypatch.setattr("utils.langflow_headers.add_provider_credentials_to_headers", AsyncMock()) + monkeypatch.setattr( + "config.settings.get_openrag_config", + lambda: SimpleNamespace( + knowledge=SimpleNamespace( + embedding_model="m", + embedding_provider=None, + chunk_size=1000, + chunk_overlap=200, + index_name="documents", + ) + ), + ) + monkeypatch.setattr( + "services.langflow_ingest_token_service.LangflowIngestTokenService.create_token", + lambda self, context: "tok", + ) + seen = {} + + async def fake_chat(client, flow_id, prompt, user_id, **kwargs): + seen["headers"] = kwargs["extra_headers"] + return "ok", "rid", [] + + monkeypatch.setattr(agent, "async_langflow_chat", fake_chat) + return seen + + +def _filter_expression(seen): + return json.loads(seen["headers"][HEADER]) + + +@pytest.mark.asyncio +async def test_chat_header_carries_pdf_and_date_range(captured): + set_search_filters({"document_types": ["application/pdf"], **DATES}) + try: + await ChatService().langflow_chat(prompt="hi", user_id="u") + finally: + set_search_filters(None) + assert _filter_expression(captured)["filter"] == [PDF, RANGE] + + +@pytest.mark.asyncio +async def test_nudges_date_only_filter_sends_range_without_sample_exclusion(captured): + await ChatService().langflow_nudges_chat(user_id="u", filters=dict(DATES)) + clauses = _filter_expression(captured)["filter"] + assert clauses == [RANGE] + assert SAMPLE_EXCLUSION not in clauses + + +@pytest.mark.asyncio +async def test_nudges_unrestricted_request_still_excludes_sample_data(captured): + await ChatService().langflow_nudges_chat(user_id="u", filters=None) + assert _filter_expression(captured)["filter"] == [SAMPLE_EXCLUSION] diff --git a/tests/unit/test_document_index_writer_ownership.py b/tests/unit/test_document_index_writer_ownership.py index 5682998dc8..837948af42 100644 --- a/tests/unit/test_document_index_writer_ownership.py +++ b/tests/unit/test_document_index_writer_ownership.py @@ -119,3 +119,30 @@ async def test_new_chunks_persist_provider_qualified_embedding_space() -> None: assert document["embedding_provider"] == "azure" assert document["embedding_space_id"] == "azure:test-model" assert document["chunk_embedding_azure_test_model"] == [0.1, 0.2, 0.3] + + +@pytest.mark.asyncio +async def test_batches_of_one_ingest_run_share_the_context_indexed_time() -> None: + opensearch = InMemoryOpenSearch() + writer = DocumentIndexWriter(opensearch_client=opensearch) + context = make_context("user-a") + context.indexed_time = "2026-01-01T00:00:00+00:00" + + first = DocumentIndexChunk(chunk_id="doc_0", text="a", vector=[0.1, 0.2, 0.3]) + second = DocumentIndexChunk(chunk_id="doc_1", text="b", vector=[0.1, 0.2, 0.3]) + await writer.index_chunks(context, [first]) + await writer.index_chunks(context, [second]) + + assert {d["indexed_time"] for d in opensearch.documents.values()} == { + "2026-01-01T00:00:00+00:00" + } + + +@pytest.mark.asyncio +async def test_indexed_time_defaults_to_now_without_a_context_value() -> None: + opensearch = InMemoryOpenSearch() + writer = DocumentIndexWriter(opensearch_client=opensearch) + + await writer.index_chunks(make_context("user-a"), [make_chunk()]) + + assert next(iter(opensearch.documents.values()))["indexed_time"] diff --git a/tests/unit/test_filter_clause_builder.py b/tests/unit/test_filter_clause_builder.py new file mode 100644 index 0000000000..6343bdf964 --- /dev/null +++ b/tests/unit/test_filter_clause_builder.py @@ -0,0 +1,45 @@ +from utils.opensearch_queries import build_filter_clauses + +RANGE = { + "range": {"indexed_time": {"gte": "2026-10-07T00:00:00Z", "lte": "2026-10-10T23:59:59.999Z"}} +} +DATES = {"created_after": "2026-10-07T00:00:00Z", "created_before": "2026-10-10T23:59:59.999Z"} + + +def test_no_filters_returns_empty(): + assert build_filter_clauses(None) == [] + assert build_filter_clauses({}) == [] + + +def test_single_and_multiple_values(): + clauses = build_filter_clauses({"document_types": ["application/pdf"], "owners": ["a", "b"]}) + assert clauses == [ + {"term": {"mimetype": "application/pdf"}}, + {"terms": {"owner": ["a", "b"]}}, + ] + + +def test_empty_list_ignored_by_default(): + assert build_filter_clauses({"owners": []}) == [] + + +def test_empty_list_matches_nothing_for_search(): + assert build_filter_clauses({"owners": []}, empty_list_matches_nothing=True) == [ + {"term": {"owner": "__IMPOSSIBLE_VALUE__"}} + ] + + +def test_date_only_filter_builds_range(): + assert build_filter_clauses(dict(DATES)) == [RANGE] + + +def test_dates_combined_with_lists_and_not_treated_as_terms(): + clauses = build_filter_clauses({"document_types": ["application/pdf"], **DATES}) + assert clauses == [{"term": {"mimetype": "application/pdf"}}, RANGE] + + +def test_single_bound_and_blank_bounds(): + assert build_filter_clauses({"created_after": "2026-10-07T00:00:00Z"}) == [ + {"range": {"indexed_time": {"gte": "2026-10-07T00:00:00Z"}}} + ] + assert build_filter_clauses({"created_after": "", "created_before": None}) == [] diff --git a/tests/unit/test_indexed_time_backfill.py b/tests/unit/test_indexed_time_backfill.py new file mode 100644 index 0000000000..4bbe664fc4 --- /dev/null +++ b/tests/unit/test_indexed_time_backfill.py @@ -0,0 +1,304 @@ +import time +from unittest.mock import AsyncMock, MagicMock + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlmodel import SQLModel + +import db.engine as db_engine +import db.models # noqa: F401 (registers tables) +from config.settings import LANGFLOW_INGEST_CALLBACK_TTL_SECONDS +from db.migrations_runtime import _already_done +from services import indexed_time_backfill as backfill + +EARLIEST = "2026-10-15T14:00:10.123456+00:00" +NOW_MS = time.time() * 1000 +OLD_MS = NOW_MS - (LANGFLOW_INGEST_CALLBACK_TTL_SECONDS + 3600) * 1000 +RECENT_MS = NOW_MS - 1000 + + +def _bucket(run, min_t, max_t, earliest=EARLIEST): + return { + "key": {"run": run}, + "min_t": {"value": min_t}, + "max_t": {"value": max_t}, + "earliest": {"hits": {"hits": [{"_source": {"indexed_time": earliest}}]}}, + } + + +def _page(buckets, after_key=None, **extra): + agg = {"buckets": buckets} + if after_key: + agg["after_key"] = after_key + return {"aggregations": {"runs": agg}, "_shards": {"failed": 0}, **extra} + + +def _split_run(run="run-1", max_t=None): + return _bucket(run, OLD_MS - 5000, OLD_MS if max_t is None else max_t) + + +def _client(*pages, update=None): + client = MagicMock() + client.indices.exists = AsyncMock(return_value=True) + client.indices.refresh = AsyncMock() + client.search = AsyncMock(side_effect=list(pages)) + client.update_by_query = AsyncMock(return_value=update or {"updated": 3}) + return client + + +@pytest_asyncio.fixture +async def session_factory(monkeypatch): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(SQLModel.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + monkeypatch.setattr(db_engine, "SessionLocal", factory) + yield factory + await engine.dispose() + + +@pytest.fixture +def use_client(monkeypatch): + def _use(client): + monkeypatch.setattr(backfill.clients, "opensearch", client, raising=False) + monkeypatch.setattr(backfill, "get_index_name", lambda: "documents") + return client + + return _use + + +async def _is_done(factory): + async with factory() as session: + return await _already_done(session, backfill.MIGRATION_NAME) + + +@pytest.mark.asyncio +async def test_split_old_run_is_updated_with_exact_earliest_string(session_factory, use_client): + client = use_client(_client(_page([_split_run()]))) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_awaited_once() + kwargs = client.update_by_query.await_args.kwargs + assert kwargs["body"]["script"]["params"]["times"] == {"run-1": EARLIEST} + assert kwargs["body"]["query"] == {"terms": {"ingest_run_id": ["run-1"]}} + assert kwargs["conflicts"] == "proceed" + assert kwargs["wait_for_completion"] is True + assert kwargs["request_timeout"] == 300 + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_uniform_run_is_not_selected(session_factory, use_client): + client = use_client(_client(_page([_bucket("run-1", OLD_MS, OLD_MS)]))) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_not_awaited() + client.indices.refresh.assert_not_awaited() + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_split_run_within_ttl_window_is_deferred(session_factory, use_client): + ttl_ms = LANGFLOW_INGEST_CALLBACK_TTL_SECONDS * 1000 + run = _bucket("run-1", NOW_MS - ttl_ms // 2, RECENT_MS) + client = use_client(_client(_page([run]))) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_not_awaited() + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_run_past_ttl_since_start_but_recent_max_is_still_deferred( + session_factory, use_client +): + # min_t + TTL has passed (the token is expired), but the run was written to + # recently, so the conservative inactivity guard still defers it. + ttl_ms = LANGFLOW_INGEST_CALLBACK_TTL_SECONDS * 1000 + run = _bucket("run-1", NOW_MS - 2 * ttl_ms, RECENT_MS) + assert run["min_t"]["value"] + ttl_ms < NOW_MS + client = use_client(_client(_page([run]))) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_not_awaited() + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_old_split_run_outside_ttl_window_is_selected(session_factory, use_client): + ttl_ms = LANGFLOW_INGEST_CALLBACK_TTL_SECONDS * 1000 + run = _bucket("run-1", NOW_MS - 3 * ttl_ms, NOW_MS - 2 * ttl_ms) + client = use_client(_client(_page([run]))) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_awaited_once() + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_any_deferred_run_blocks_completion_and_retries_next_boot( + session_factory, use_client +): + deferred = _split_run("run-recent", max_t=RECENT_MS) + client = use_client(_client(_page([_split_run("run-old"), deferred]))) + + await backfill.run_indexed_time_backfill() + + # The eligible run is still repaired, but the deferred one keeps the + # migration open. + assert client.update_by_query.await_args.kwargs["body"]["script"]["params"]["times"] == { + "run-old": EARLIEST + } + assert not await _is_done(session_factory) + + # Next boot: the deferred run has aged out, so it is repaired and the + # migration completes. + aged = _split_run("run-recent") + client.search.side_effect = [_page([aged])] + client.update_by_query.reset_mock() + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_awaited_once() + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_follows_after_key_until_absent(session_factory, use_client): + client = use_client( + _client( + _page([_split_run("run-1")], after_key={"run": "run-1"}), + _page([_split_run("run-2")]), + ) + ) + + await backfill.run_indexed_time_backfill() + + assert client.search.await_count == 2 + assert client.search.await_args_list[1].kwargs["body"]["aggs"]["runs"]["composite"][ + "after" + ] == {"run": "run-1"} + assert client.update_by_query.await_count == 2 + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_update_failures_are_not_marked_done(session_factory, use_client): + use_client(_client(_page([_split_run()]), update={"updated": 1, "failures": [{"x": 1}]})) + + await backfill.run_indexed_time_backfill() + + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_version_conflicts_are_not_marked_done(session_factory, use_client): + use_client(_client(_page([_split_run()]), update={"updated": 1, "version_conflicts": 2})) + + await backfill.run_indexed_time_backfill() + + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_update_timeout_is_not_marked_done(session_factory, use_client): + use_client(_client(_page([_split_run()]), update={"updated": 1, "timed_out": True})) + + await backfill.run_indexed_time_backfill() + + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "partial", [{"timed_out": True}, {"_shards": {"failed": 1}}], ids=["timed_out", "shards"] +) +async def test_partial_scan_is_not_marked_done(session_factory, use_client, partial): + client = use_client(_client({**_page([_split_run()]), **partial})) + + await backfill.run_indexed_time_backfill() + + client.update_by_query.assert_not_awaited() + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_missing_index_is_a_clean_pass(session_factory, use_client): + client = use_client(_client()) + client.indices.exists.return_value = False + + await backfill.run_indexed_time_backfill() + + client.search.assert_not_awaited() + assert await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_search_error_does_not_escape_or_mark_done(session_factory, use_client): + client = use_client(_client()) + client.search.side_effect = RuntimeError("boom") + + await backfill.run_indexed_time_backfill() + + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_cancellation_is_re_raised(session_factory, use_client): + import asyncio + + client = use_client(_client()) + client.search.side_effect = asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await backfill.run_indexed_time_backfill() + assert not await _is_done(session_factory) + + +@pytest.mark.asyncio +async def test_already_done_skips_opensearch(session_factory, use_client): + client = use_client(_client(_page([_split_run()]))) + await backfill.run_indexed_time_backfill() + client.search.reset_mock() + client.update_by_query.reset_mock() + + await backfill.run_indexed_time_backfill() + + client.indices.exists.assert_awaited_once() + client.search.assert_not_awaited() + client.update_by_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_refresh_only_when_documents_were_updated(session_factory, use_client): + client = use_client(_client(_page([_split_run()]), update={"updated": 0})) + await backfill.run_indexed_time_backfill() + client.indices.refresh.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_refresh_runs_once_after_updates(session_factory, use_client): + client = use_client( + _client( + _page([_split_run("run-1")], after_key={"run": "run-1"}), + _page([_split_run("run-2")]), + ) + ) + await backfill.run_indexed_time_backfill() + client.indices.refresh.assert_awaited_once_with(index="documents") + + +@pytest.mark.asyncio +async def test_returns_early_without_session_factory_or_client(monkeypatch, use_client): + client = use_client(_client(_page([_split_run()]))) + monkeypatch.setattr(db_engine, "SessionLocal", None) + + await backfill.run_indexed_time_backfill() + + client.indices.exists.assert_not_awaited() diff --git a/tests/unit/test_langflow_ingest_callback.py b/tests/unit/test_langflow_ingest_callback.py index f077675d57..b97caa4ec0 100644 --- a/tests/unit/test_langflow_ingest_callback.py +++ b/tests/unit/test_langflow_ingest_callback.py @@ -443,3 +443,58 @@ def test_ingest_token_prefers_jwt_signing_key_over_session_secret(monkeypatch): session_service = LangflowIngestTokenService(secret=session_secret, ttl_seconds=60) with pytest.raises(ValueError, match="Invalid Langflow ingest token"): session_service.validate_token(token) + + +def test_ingest_token_round_trips_indexed_time(): + from services.langflow_ingest_token_service import LangflowIngestTokenService + + payload = LangflowIngestTokenService._context_to_payload( + DocumentIndexContext( + document_id="d", + mimetype="application/pdf", + embedding_model="m", + indexed_time="2026-01-01T00:00:00+00:00", + ) + ) + context = LangflowIngestTokenService._payload_to_context(payload) + + assert context.indexed_time == "2026-01-01T00:00:00+00:00" + + +def _token_context(**overrides): + return DocumentIndexContext( + document_id="d", mimetype="application/pdf", embedding_model="m", **overrides + ) + + +def test_create_token_defaults_indexed_time_to_issue_time(): + import datetime + + import jwt + + service = LangflowIngestTokenService(secret="test-secret" * 4, ttl_seconds=60) + original = _token_context() + + token = service.create_token(original) + + iat = jwt.decode(token, options={"verify_signature": False})["iat"] + context, _ = service.validate_token(token) + assert context.indexed_time == datetime.datetime.fromtimestamp(iat, datetime.UTC).isoformat() + + +def test_create_token_preserves_explicit_indexed_time(): + service = LangflowIngestTokenService(secret="test-secret" * 4, ttl_seconds=60) + + token = service.create_token(_token_context(indexed_time="2026-01-01T00:00:00+00:00")) + + context, _ = service.validate_token(token) + assert context.indexed_time == "2026-01-01T00:00:00+00:00" + + +def test_create_token_does_not_mutate_callers_context(): + service = LangflowIngestTokenService(secret="test-secret" * 4, ttl_seconds=60) + original = _token_context() + + service.create_token(original) + + assert original.indexed_time is None diff --git a/tests/unit/test_opensearch_embedding_spaces.py b/tests/unit/test_opensearch_embedding_spaces.py index 2dd184895d..aff1a4231e 100644 --- a/tests/unit/test_opensearch_embedding_spaces.py +++ b/tests/unit/test_opensearch_embedding_spaces.py @@ -1,4 +1,5 @@ import importlib +import json import sys from types import ModuleType, SimpleNamespace @@ -392,3 +393,104 @@ def _bulk_ingest(**kwargs): assert embedding.embedded == ["11 月 27 日 ( 木 )"] assert indexed == ["11 月 27 日 ( 木 )"] + + +_RANGE = {"range": {"indexed_time": {"gte": "2026-10-07T00:00:00Z", "lte": "2026-10-10T23:59:59Z"}}} +_PDF = {"term": {"mimetype": "application/pdf"}} + + +@pytest.mark.parametrize( + ("filter_obj", "expected"), + [ + ({"filter": [_RANGE]}, [_RANGE]), + ({"filter": [_PDF, _RANGE]}, [_PDF, _RANGE]), + ({"filter": [{"range": {"indexed_time": {}}}]}, []), + ({"filter": [{"range": {"indexed_time": {"gte": None}}}]}, []), + ({"filter": [{"range": "bad"}]}, []), + ( + {"filter": [{"range": {"a": {"gte": 1}, "b": {"gte": 2}}}]}, + [], + ), + ], +) +def test_filter_parser_keeps_valid_range_clauses( + monkeypatch: pytest.MonkeyPatch, filter_obj, expected +) -> None: + module = _load_opensearch_module(monkeypatch) + component = module.OpenSearchVectorStoreComponentMultimodalMultiEmbedding() + + assert component._coerce_filter_clauses(filter_obj) == expected + assert component._coerce_filter_clauses(json.dumps(filter_obj)) == expected + + +@pytest.mark.parametrize("clauses", [[_RANGE], [_PDF, _RANGE]], ids=["date-only", "pdf-plus-date"]) +def test_search_applies_date_filter_outside_and_inside_each_knn( + monkeypatch: pytest.MonkeyPatch, clauses +) -> None: + module = _load_opensearch_module(monkeypatch) + aggregation = { + "aggregations": { + "embedding_spaces": { + "buckets": [{"key": {"space_id": "azure:text-embedding-3-small"}, "doc_count": 2}] + }, + "legacy_embedding_models": {"buckets": []}, + } + } + fields = { + "chunk_embedding_azure_text_embedding_3_small": {"type": "knn_vector", "dimension": 2} + } + client = _SearchClient(aggregation, fields) + + class Child: + def embed_query(self, query): + return [0.1, 0.2] + + class Root(Child): + deployment = "selected-model" + + def for_model(self, route): + return Child() + + component = _component(module, client, Root()) + component.filter_expression = json.dumps({"filter": clauses}) + component.search("needle") + + bool_query = client.query["query"]["bool"] + assert bool_query["filter"] == clauses + knn_queries = bool_query["should"][0]["dis_max"]["queries"] + assert knn_queries + for query in knn_queries: + (knn,) = query["knn"].values() + assert knn["filter"] == {"bool": {"filter": clauses}} + + +def test_search_without_filter_leaves_knn_unfiltered(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_opensearch_module(monkeypatch) + aggregation = { + "aggregations": { + "embedding_spaces": { + "buckets": [{"key": {"space_id": "azure:text-embedding-3-small"}, "doc_count": 2}] + }, + "legacy_embedding_models": {"buckets": []}, + } + } + fields = { + "chunk_embedding_azure_text_embedding_3_small": {"type": "knn_vector", "dimension": 2} + } + client = _SearchClient(aggregation, fields) + + class Child: + def embed_query(self, query): + return [0.1, 0.2] + + class Root(Child): + deployment = "selected-model" + + def for_model(self, route): + return Child() + + _component(module, client, Root()).search("needle") + + for query in client.query["query"]["bool"]["should"][0]["dis_max"]["queries"]: + (knn,) = query["knn"].values() + assert "filter" not in knn diff --git a/tests/unit/test_search_api_date_bounds.py b/tests/unit/test_search_api_date_bounds.py new file mode 100644 index 0000000000..23f56c631c --- /dev/null +++ b/tests/unit/test_search_api_date_bounds.py @@ -0,0 +1,13 @@ +import pytest +from fastapi import HTTPException + +from api.search import SearchBody, search + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bound", ["created_after", "created_before"]) +async def test_non_string_bound_in_filters_is_rejected(bound): + body = SearchBody(query="*", filters={bound: 123}) + with pytest.raises(HTTPException) as exc: + await search(body, search_service=None, session_manager=None, user=None) + assert exc.value.status_code == 422 diff --git a/tests/unit/test_search_service_score_threshold.py b/tests/unit/test_search_service_score_threshold.py index f65eea6009..475e9ab980 100644 --- a/tests/unit/test_search_service_score_threshold.py +++ b/tests/unit/test_search_service_score_threshold.py @@ -155,3 +155,21 @@ async def test_num_candidates_retry_keeps_the_threshold(monkeypatch): assert all("num_candidates" not in params for params in _knn_params(retry["query"])) assert retry["query"]["function_score"]["min_score"] == 1.0 assert "min_score" not in retry + + +@pytest.mark.asyncio +async def test_filters_are_pushed_into_each_knn_clause(monkeypatch): + """Filters must constrain kNN neighbor selection, not only post-filter the top-k.""" + os_client = _OpenSearch() + service = _make_service(os_client, monkeypatch) + _set_context(0) + set_search_filters({"created_after": "2024-01-01T00:00:00Z"}) + + await service.search_tool("MCP") + + params = _knn_params(os_client.bodies[-1]["query"]) + assert params + for knn in params: + assert knn["filter"] == { + "bool": {"filter": [{"range": {"indexed_time": {"gte": "2024-01-01T00:00:00Z"}}}]} + }