Skip to content

Commit 1b22fd7

Browse files
Python: Fix Cosmos hybrid search field validation (#14650)
### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> Cosmos DB NoSQL hybrid search previously inserted the requested text-field name directly into the generated query. This change limits hybrid text search to fields declared by the collection definition and safely renders the resolved storage property. ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> - Resolve `additional_property_name` against declared data-field logical and storage names. - Reject unknown field names before executing a Cosmos query. - Render the resolved property using Cosmos bracket notation with JSON string encoding. - Add regression coverage for logical names, storage aliases, and injection-shaped unknown names. Validation: - `uv run --frozen pytest tests\unit\connectors\memory\azure_cosmos_db\test_azure_cosmos_db_no_sql_collection.py -q` — passed (30 tests). - `uv run --frozen ruff check semantic_kernel\connectors\azure_cosmos_db.py tests\unit\connectors\memory\azure_cosmos_db\test_azure_cosmos_db_no_sql_collection.py` — passed. - `uv run --frozen ruff format --check semantic_kernel\connectors\azure_cosmos_db.py tests\unit\connectors\memory\azure_cosmos_db\test_azure_cosmos_db_no_sql_collection.py` — passed. - `uv run --frozen mypy semantic_kernel\connectors\azure_cosmos_db.py --config-file mypy.ini` — passed. - `uv run --frozen python -m compileall -q semantic_kernel\connectors\azure_cosmos_db.py tests\unit\connectors\memory\azure_cosmos_db\test_azure_cosmos_db_no_sql_collection.py` — passed. - `uv run --project python --frozen pre-commit run --config python\.pre-commit-config.yaml --files python\semantic_kernel\connectors\azure_cosmos_db.py python\tests\unit\connectors\memory\azure_cosmos_db\test_azure_cosmos_db_no_sql_collection.py` — passed (all applicable hooks). - `git diff --check` — passed. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 6f1ae42 commit 1b22fd7

2 files changed

Lines changed: 101 additions & 3 deletions

File tree

‎python/semantic_kernel/connectors/azure_cosmos_db.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import ast
44
import asyncio
5+
import json
56
import re
67
import sys
78
from collections.abc import Sequence
@@ -833,11 +834,24 @@ def _substitute(match: "re.Match[str]", mapping: dict[str, str] = rename) -> str
833834
elif search_type == SearchType.KEYWORD_HYBRID:
834835
# Hybrid search: requires both a vector and keywords
835836
params.append({"name": "@keywords", "value": values})
836-
text_field = options.additional_property_name
837-
if not text_field:
837+
requested_text_field_name = options.additional_property_name
838+
if not requested_text_field_name:
838839
raise VectorStoreModelException("Hybrid search requires 'keyword_field_name' in options.")
840+
text_field = next(
841+
(
842+
field
843+
for field in self.definition.data_fields
844+
if requested_text_field_name in {field.name, field.storage_name}
845+
),
846+
None,
847+
)
848+
if not text_field:
849+
raise VectorStoreModelException(
850+
f"Data field '{requested_text_field_name}' not found in the data model definition."
851+
)
852+
text_field_reference = f"c[{json.dumps(text_field.storage_name or text_field.name)}]"
839853
distance_clause = (
840-
f"RRF(VectorDistance(c.{vector_field_name}, @vector), FullTextScore(c.{text_field}, @keywords))"
854+
f"RRF(VectorDistance(c.{vector_field_name}, @vector), FullTextScore({text_field_reference}, @keywords))"
841855
)
842856
else:
843857
raise VectorStoreModelException(f"Search type '{search_type}' is not supported.")

‎python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,90 @@ async def _empty_results():
252252
assert {"name": "@filter_p1", "value": "beta"} in parameters
253253

254254

255+
@pytest.mark.parametrize("additional_property_name", ["content", "stored_content"])
256+
async def test_azure_cosmos_db_no_sql_inner_search_resolves_hybrid_text_field(
257+
azure_cosmos_db_no_sql_unit_test_env,
258+
record_type,
259+
definition,
260+
collection_name: str,
261+
additional_property_name: str,
262+
) -> None:
263+
"""Test hybrid search resolves logical and storage field names before building the query."""
264+
from semantic_kernel.data.vector import SearchType, VectorSearchOptions
265+
266+
definition.data_fields[0].storage_name = "stored_content"
267+
vector_collection = CosmosNoSqlCollection(
268+
record_type=record_type,
269+
definition=definition,
270+
collection_name=collection_name,
271+
)
272+
273+
async def _empty_results():
274+
for item in []:
275+
yield item
276+
277+
mock_container_proxy = MagicMock()
278+
mock_container_proxy.query_items = MagicMock(return_value=_empty_results())
279+
options = VectorSearchOptions(
280+
vector_property_name="vector",
281+
additional_property_name=additional_property_name,
282+
top=3,
283+
)
284+
285+
with patch.object(
286+
CosmosNoSqlCollection,
287+
"_get_container_proxy",
288+
new=AsyncMock(return_value=mock_container_proxy),
289+
):
290+
await vector_collection._inner_search(
291+
search_type=SearchType.KEYWORD_HYBRID,
292+
options=options,
293+
values="alpha",
294+
vector=[0.1, 0.2, 0.3, 0.4, 0.5],
295+
)
296+
297+
query, kwargs = mock_container_proxy.query_items.call_args
298+
assert 'FullTextScore(c["stored_content"], @keywords)' in query[0]
299+
assert {"name": "@keywords", "value": "alpha"} in kwargs["parameters"]
300+
301+
302+
async def test_azure_cosmos_db_no_sql_inner_search_rejects_unknown_hybrid_text_field(
303+
azure_cosmos_db_no_sql_unit_test_env,
304+
record_type,
305+
collection_name: str,
306+
) -> None:
307+
"""Test hybrid search rejects injection-shaped text fields before executing a query."""
308+
from semantic_kernel.data.vector import SearchType, VectorSearchOptions
309+
310+
vector_collection = CosmosNoSqlCollection(
311+
record_type=record_type,
312+
collection_name=collection_name,
313+
)
314+
mock_container_proxy = MagicMock()
315+
options = VectorSearchOptions(
316+
vector_property_name="vector",
317+
additional_property_name="content), @keywords)) FROM c WHERE true --",
318+
top=3,
319+
)
320+
321+
with (
322+
patch.object(
323+
CosmosNoSqlCollection,
324+
"_get_container_proxy",
325+
new=AsyncMock(return_value=mock_container_proxy),
326+
),
327+
pytest.raises(VectorStoreModelException, match="not found in the data model definition"),
328+
):
329+
await vector_collection._inner_search(
330+
search_type=SearchType.KEYWORD_HYBRID,
331+
options=options,
332+
values="alpha",
333+
vector=[0.1, 0.2, 0.3, 0.4, 0.5],
334+
)
335+
336+
mock_container_proxy.query_items.assert_not_called()
337+
338+
255339
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_URL"]], indirect=True)
256340
def test_azure_cosmos_db_no_sql_collection_init_no_url(
257341
azure_cosmos_db_no_sql_unit_test_env,

0 commit comments

Comments
 (0)