Skip to content

Add multiple new search types to AzureCosmosDBNoSqlVectorSearch#20871

Open
allenkim0129 wants to merge 4 commits intorun-llama:mainfrom
allenkim0129:main
Open

Add multiple new search types to AzureCosmosDBNoSqlVectorSearch#20871
allenkim0129 wants to merge 4 commits intorun-llama:mainfrom
allenkim0129:main

Conversation

@allenkim0129
Copy link

Description

This PR significantly extends llama-index-vector-stores-azurecosmosnosql with multiple new search types, a fully parameterised query engine, comprehensive test coverage, and an updated README.

What changed

base.py — new search engine

  • Replaced the old _query method (single hard-coded vector search) with a new _search_query method that supports 6 search types via a search_type kwarg
  • All query components (TOP/OFFSET-LIMIT, SELECT, WHERE, ORDER BY) are built from composable helpers using a ParamMapping helper to produce fully parameterised CosmosDB SQL queries (no string interpolation of user values)
  • Backward-compatible: the legacy pre_filter={"where_clause": ..., "limit_offset_clause": ...} dict is still accepted and mapped to the new where / offset_limit kwargs
  • Added _generate_projection_fields, _generate_limit_clause, _generate_order_by_clause, _generate_order_by_component_with_full_text_rank_filter, _construct_search_query, _execute_search_query, _validate_search_args, _is_vector_search_type, _is_full_text_search_type, _is_vector_search_with_threshold helpers

utils.py — new helper module (extracted from base.py)

  • Constants — all field/parameter name literals in one place
  • ParamMapping — builds @param-style parameterised queries; handles projections, vector distance, inline vector literals for ORDER BY RANK contexts
  • AzureCosmosDBNoSqlVectorSearchType — enum for all supported search types

Search types added

search_type Description
vector Standard nearest-neighbour via VectorDistance, ORDER BY
vector_score_threshold Vector search filtered by a minimum similarity score
full_text_search FullTextContains WHERE predicate, no ranking
full_text_ranking ORDER BY RANK FullTextScore(...), single or multi-field RRF
hybrid ORDER BY RANK RRF(FullTextScore, VectorDistance)
hybrid_score_threshold Hybrid RRF + client-side score threshold filter
weighted_hybrid_search Hybrid RRF + client-side per-component weight re-ranking (weight= is not supported in CosmosDB SQL syntax, so weights are applied after results are returned)

Query options added

  • where — raw CosmosDB SQL WHERE predicate
  • offset_limitOFFSET x LIMIT y pagination (replaces TOP for hybrid types)
  • projection_mapping — project specific document fields as node metadata
  • return_with_vectors — include raw embeddings in results
  • full_text_rank_filter — list of {search_field, search_text} dicts for full-text ranking components
  • threshold — minimum similarity score for *_score_threshold types
  • weights — per-component float list for weighted_hybrid_search (client-side RRF re-ranking)

Tests

  • test_unit.py — 93 unit tests covering all helpers without any DB access:
    • TestParamMapping, TestGenerateProjectionFields, TestGenerateLimitClause, TestGenerateOrderByComponentWithFullTextRankFilter, TestGenerateOrderByClause, TestConstructSearchQuery, TestIsHelpers, TestSearchQueryPreFilterCompat
  • test_azurecosmosnosql.py — 15 integration tests against a live CosmosDB account:
    • test_add_and_delete, test_query, test_query_vector, test_query_vector_with_k, test_query_vector_with_projection_mapping, test_query_vector_with_offset_limit, test_query_vector_with_where_filter, test_query_full_text_search, test_query_full_text_ranking, test_query_full_text_ranking_multiple_fields, test_query_hybrid_with_score_threshold, test_query_hybrid_with_weights, test_query_weighted_hybrid_search, test_cosmos_client_with_host_and_key, test_cosmos_client_with_connection_string

README

  • Rewrote from scratch to cover: installation, quick start, all 7 search types with code examples, all query options, full text search container setup, weighted hybrid search with weight notation

New Package?

  • Yes
  • No

Version Bump?

Did I bump the version in the pyproject.toml?

  • Yes
  • No

Note for reviewer: the changes are additive and backward-compatible (old pre_filter callers are unaffected). Please advise if a minor or patch bump is required.

Type of Change

  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

How Has This Been Tested?

  • I added new unit tests to cover this change

Unit tests (no DB): pytest tests/test_unit.py93 passed

Integration tests (live CosmosDB):

export AZURE_COSMOSDB_URI="https://<account>.documents.azure.com:443/"
export AZURE_COSMOSDB_KEY="<key>"
pytest tests/test_azurecosmosnosql.py

15 passed

Suggested Checklist

  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • I have added Google Colab support for the newly added notebooks
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I ran uv run make format; uv run make lint to appease the lint gods

Updated to use 'id', instead of 'ref_doc_id'.
Also fixed tests.
- VECTOR_SCORE_THRESHOLD
- FULL_TEXT_SEARCH
- FULL_TEXT_RANKING
- HYBRID
- HYBRID_SCORE_THRESHOLD
- WEIGHTED_HYBRID_SEARCH
…RF re-ranking

- Add WEIGHTED_HYBRID_SEARCH enum value to AzureCosmosDBNoSqlVectorSearchType
- Accept 'weights' param in _search_query, _construct_search_query, _execute_search_query
- Apply position-based weighted RRF re-ranking client-side (CosmosDB SQL does not
  support weight= inside ORDER BY RANK RRF)
- Update _is_full_text_search_type and _is_vector_search_type to include new type
- Add unit tests: ORDER BY clause, construct query, is_* helpers
- Add integration test: test_query_weighted_hybrid_search
- Update README with weighted hybrid search section and usage examples
@dosubot dosubot bot added the size:XXL This PR changes 1000+ lines, ignoring generated files. label Mar 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:XXL This PR changes 1000+ lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant