Python: Introducing search for CosmosDB NoSQL Collections

[eavanvalkenburg]

Motivation and Context

Adds the search functions for Cosmos DB NoSQL.

Closes #6835

Description

Adds the search pieces and turns out the upsert cannot work cross partition key with execute batch, so refactored to using gather(upsert).

Contribution Checklist

• The code builds clean without any errors or warnings
• The PR follows the SK Contribution Guidelines and the pre-submission formatting script raises no violations
• All unit tests pass, and I have added new tests where possible
• I didn't break anyone 😄

[TaoChenOSU]

Should we add an integration test case that test inserting multiple records across partitions?

[markwallace-microsoft]

Python Test Coverage Report

File	Stmts	Miss	Cover	Missing
semantic_kernel
kernel.py	199	47	76%	148, 159, 163, 313–316, 423, 437–480
semantic_kernel/agents/group_chat
agent_chat.py	124	2	98%	78, 171
agent_group_chat.py	100	2	98%	151, 201
broadcast_queue.py	72	1	99%	35
semantic_kernel/agents/open_ai
assistant_content_generation.py	141	9	94%	97–98, 329–337, 379, 381
azure_assistant_agent.py	107	2	98%	284, 304
open_ai_assistant_agent.py	105	2	98%	252, 272
open_ai_assistant_base.py	467	8	98%	260, 338–339, 747, 868, 871, 945, 1007
semantic_kernel/connectors/ai
audio_to_text_client_base.py	9	1	89%	51
chat_completion_client_base.py	116	2	98%	382, 392
completion_usage.py	8	1	88%	17
text_to_audio_client_base.py	8	1	88%	52
semantic_kernel/connectors/ai/anthropic/services
anthropic_chat_completion.py	176	5	97%	147, 165, 169, 223, 419
semantic_kernel/connectors/ai/azure_ai_inference/services
azure_ai_inference_chat_completion.py	119	7	94%	120, 146–149, 159, 180, 202
azure_ai_inference_text_embedding.py	41	1	98%	86
semantic_kernel/connectors/ai/bedrock/services
bedrock_chat_completion.py	136	14	90%	117, 138, 163, 167–170, 228, 246–265, 324
bedrock_text_completion.py	57	2	96%	95, 118
bedrock_text_embedding.py	45	1	98%	94
semantic_kernel/connectors/ai/bedrock/services/model_provider
bedrock_ai21_labs.py	13	1	92%	67
bedrock_anthropic_claude.py	12	1	92%	54
bedrock_cohere.py	20	1	95%	75
utils.py	80	20	75%	68, 71, 102, 106–115, 132–150, 171–174
semantic_kernel/connectors/ai/embeddings
embedding_generator_base.py	8	1	88%	50
semantic_kernel/connectors/ai/google/google_ai/services
google_ai_chat_completion.py	119	4	97%	126, 152, 175, 177
google_ai_text_completion.py	63	2	97%	98, 121
utils.py	65	3	95%	139, 159–164
semantic_kernel/connectors/ai/google/vertex_ai/services
utils.py	66	3	95%	140, 160–165
vertex_ai_chat_completion.py	119	4	97%	121, 147, 170, 172
vertex_ai_text_completion.py	62	2	97%	95, 116
semantic_kernel/connectors/ai/hugging_face/services
hf_text_completion.py	60	3	95%	103, 112, 127
hf_text_embedding.py	32	5	84%	79–83
semantic_kernel/connectors/ai/mistral_ai/prompt_execution_settings
mistral_ai_prompt_execution_settings.py	33	2	94%	59–62
semantic_kernel/connectors/ai/mistral_ai/services
mistral_ai_chat_completion.py	122	36	70%	119–122, 147–150, 180–184, 197–205, 221–229, 241–254, 260, 269–273, 318–321
mistral_ai_text_embedding.py	39	1	97%	108
semantic_kernel/connectors/ai/ollama/services
ollama_chat_completion.py	107	11	90%	114, 139, 143–144, 154, 186, 223, 233–234, 256, 283
ollama_text_completion.py	57	3	95%	93, 103, 130
utils.py	46	25	46%	29, 44–52, 64–86, 98–102, 119–122
semantic_kernel/connectors/ai/onnx
utils.py	53	3	94%	50–51, 112
semantic_kernel/connectors/ai/onnx/services
onnx_gen_ai_chat_completion.py	72	7	90%	67–68, 98, 122, 167, 173, 179
onnx_gen_ai_completion_base.py	58	21	64%	59–71, 79–90
onnx_gen_ai_text_completion.py	46	5	89%	54–55, 87, 117, 133
semantic_kernel/connectors/ai/open_ai/prompt_execution_settings
open_ai_prompt_execution_settings.py	95	1	99%	117
open_ai_text_to_audio_execution_settings.py	17	1	94%	29
open_ai_text_to_image_execution_settings.py	36	1	97%	60
semantic_kernel/connectors/ai/open_ai/services
azure_chat_completion.py	103	3	97%	140, 149, 152
open_ai_audio_to_text_base.py	28	3	89%	37–38, 43
open_ai_chat_completion_base.py	127	5	96%	71, 121, 141, 177, 287
open_ai_handler.py	104	8	92%	134, 142–143, 159–160, 172, 181–182
open_ai_text_completion_base.py	80	2	98%	56, 161
open_ai_text_to_audio_base.py	25	3	88%	35–36, 41
semantic_kernel/connectors/ai/open_ai/settings
azure_open_ai_settings.py	24	4	83%	111–114
semantic_kernel/connectors/memory/azure_ai_search
azure_ai_search_collection.py	134	33	75%	170, 172, 245–283, 293–303, 307, 311, 316–319
azure_ai_search_store.py	42	2	95%	130–131
utils.py	73	7	90%	125, 127, 129, 150–151, 154–155
semantic_kernel/connectors/memory/azure_cosmos_db
azure_cosmos_db_no_sql_base.py	49	9	82%	90, 93–94, 102, 112–116
azure_cosmos_db_no_sql_collection.py	164	59	64%	98–99, 126–131, 145–150, 154–161, 172–186, 192–198, 204–208, 231–242, 246, 250, 274, 303, 307–312
azure_cosmos_db_no_sql_store.py	33	6	82%	91–96
utils.py	60	7	88%	41, 53, 67, 155, 169–172
semantic_kernel/connectors/memory/in_memory
in_memory_collection.py	142	16	89%	65, 117, 119, 139, 149, 174, 189, 201, 219–222, 226, 228, 230–231
semantic_kernel/connectors/memory/qdrant
qdrant_collection.py	119	4	97%	187, 191, 320–321
qdrant_store.py	48	2	96%	139–140
semantic_kernel/connectors/memory/redis
redis_collection.py	242	47	81%	173, 178–179, 190–200, 208–228, 231–238, 253–258, 262, 338, 375–382, 460, 482–489
redis_store.py	42	2	95%	108–109
utils.py	71	32	55%	153–154, 172, 174, 181–196, 203–233
semantic_kernel/connectors/memory/weaviate
utils.py	80	22	72%	92–97, 194, 243, 262, 269–272, 281–292
weaviate_collection.py	193	74	62%	172–184, 188–201, 205–213, 224–241, 248–255, 264–284, 293–309, 313, 316–320, 364, 368–375, 390, 409, 425, 434–435, 441, 446
weaviate_store.py	59	17	71%	110–118, 122–127, 132–137, 142–143
semantic_kernel/connectors/openapi_plugin
openapi_manager.py	58	2	97%	110–111
openapi_parser.py	88	2	98%	71, 128
openapi_runner.py	105	2	98%	181–182
semantic_kernel/connectors/openapi_plugin/models
rest_api_operation.py	129	1	99%	242
semantic_kernel/contents
audio_content.py	18	1	94%	53
binary_content.py	92	2	98%	170–171
function_call_content.py	100	2	98%	185, 213
streaming_chat_message_content.py	68	1	99%	210
streaming_content_mixin.py	39	2	95%	37, 64
semantic_kernel/core_plugins/sessions_python_tool
sessions_python_plugin.py	134	8	94%	69, 82–91, 99
sessions_python_settings.py	39	4	90%	84–87
semantic_kernel/data
search_filter.py	25	1	96%	7
semantic_kernel/data/record_definition
vector_store_model_definition.py	68	17	75%	66, 80–89, 101–106
vector_store_record_utils.py	28	2	93%	55, 57
semantic_kernel/data/text_search
text_search.py	72	4	94%	125, 165, 205, 293
utils.py	33	7	79%	23, 54–60, 69–70
vector_store_text_search.py	76	17	78%	167–174, 180–187, 192
semantic_kernel/data/vector_search
vector_search_filter.py	20	1	95%	6
vector_text_search.py	16	1	94%	45
vectorizable_text_search.py	15	1	93%	50
vectorized_search.py	15	1	93%	45
semantic_kernel/data/vector_storage
vector_store.py	16	1	94%	51
vector_store_record_collection.py	253	21	92%	431, 491–495, 503–507, 548–552, 559–563
semantic_kernel/functions
kernel_function_decorator.py	98	1	99%	102
kernel_function_from_method.py	96	1	99%	153
kernel_function_from_prompt.py	154	7	95%	165–166, 180, 201, 219, 239, 322
kernel_function_log_messages.py	36	6	83%	37–43
kernel_plugin.py	199	5	97%	468, 471, 500, 521, 546
semantic_kernel/planners
plan.py	234	45	81%	54, 163–165, 197, 214–227, 264, 269, 277–278, 288–291, 308, 313, 329, 332–337, 355, 360, 363, 365, 372, 386–388, 393–397
semantic_kernel/planners/function_calling_stepwise_planner
function_calling_stepwise_planner.py	116	4	97%	145, 189–190, 198
semantic_kernel/planners/sequential_planner
sequential_planner.py	64	6	91%	71, 75, 109, 125, 134–135
sequential_planner_extensions.py	50	9	82%	31–32, 56, 110–124
sequential_planner_parser.py	77	12	84%	66–74, 93, 117–120
semantic_kernel/processes
process_builder.py	71	41	42%	44–53, 57–59, 65–75, 79–88, 92–98, 102–106, 111, 115–120
process_end_step.py	19	2	89%	37, 41
process_event.py	15	3	80%	21, 26, 31
process_function_target_builder.py	27	3	89%	43–46
process_step_builder.py	109	25	77%	45, 94, 108, 115–128, 140–147, 156, 165–175, 180, 194, 211
process_step_edge_builder.py	35	3	91%	43, 58, 68
process_types.py	25	1	96%	35
semantic_kernel/processes/dapr_runtime
dapr_actor_registration.py	26	17	35%	25–31, 37–42, 48–53
dapr_kernel_process.py	22	2	91%	26, 38
dapr_kernel_process_context.py	34	2	94%	25, 27
dapr_process_info.py	31	4	87%	25, 30, 40, 47
dapr_step_info.py	31	2	94%	30, 39
semantic_kernel/processes/dapr_runtime/actors
event_buffer_actor.py	52	17	67%	40–43, 65–68, 81–85, 89–92
external_event_buffer_actor.py	55	12	78%	52–55, 77–80, 103–106
message_buffer_actor.py	49	12	76%	39–42, 57–60, 74–77
process_actor.py	242	113	53%	69–73, 78, 81, 87, 90, 94, 100, 103, 119–122, 127, 140, 156–160, 165, 175, 180, 189–197, 201–203, 211–219, 226, 242–303, 308, 319, 333–341, 344–351, 355–383, 414–427
step_actor.py	250	148	41%	77–81, 86, 89, 93–94, 106–108, 152–155, 169–171, 176–250, 254–350, 354, 358–361, 365–381, 385–404, 408–426, 430–433, 437, 441–444
semantic_kernel/processes/kernel_process
kernel_process_step_context.py	20	2	90%	38, 41
semantic_kernel/processes/local_runtime
local_kernel_process.py	22	2	91%	24, 35
local_kernel_process_context.py	32	2	94%	66–67
local_process.py	134	52	61%	92, 102, 120–130, 163–190, 194–199, 203, 207–213, 217–227, 231–232
local_step.py	163	100	39%	61, 72, 81–169, 173, 177, 181–182, 187–249, 253–256, 260–263, 267–276, 282–285, 289–291
semantic_kernel/prompt_template
kernel_prompt_template.py	78	7	91%	144–151
semantic_kernel/schema
kernel_json_schema_builder.py	131	9	93%	54, 93, 189, 197, 208, 216, 231, 235–236
semantic_kernel/services
ai_service_client_base.py	22	1	95%	64
semantic_kernel/template_engine/blocks
code_block.py	77	1	99%	119
named_arg_block.py	43	1	98%	98
semantic_kernel/utils/authentication
entra_id_authentication.py	15	2	87%	26, 38
semantic_kernel/utils/telemetry
user_agent.py	16	2	88%	18–19
semantic_kernel/utils/telemetry/model_diagnostics
decorators.py	171	4	98%	364–367
TOTAL	15930	1431	91%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
2906	4 💤	0 ❌	0 🔥	1m 13s ⏱️