Python: Azure Cosmos DB NoSQL Vector Store & Collection implementation

[TaoChenOSU]

Motivation and Context

We are implementing the Azure Cosmos DB NoSQL vector store and vector collection.

Description

Azure Cosmos DB NoSQL vector store & collection implementation.

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 😄

[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
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/services
mistral_ai_chat_completion.py	118	7	94%	118–121, 307–310
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%	113
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	27	3	89%	38–39, 44
open_ai_chat_completion_base.py	127	5	96%	71, 121, 141, 177, 287
open_ai_handler.py	98	6	94%	133, 141–142, 154, 163–164
open_ai_text_completion_base.py	80	2	98%	56, 161
semantic_kernel/connectors/ai/open_ai/settings
azure_open_ai_settings.py	23	4	83%	104–107
semantic_kernel/connectors/memory/azure_ai_search
azure_ai_search_collection.py	134	33	75%	171, 173, 246–284, 294–304, 308, 312, 317–320
azure_ai_search_store.py	42	2	95%	130–131
utils.py	66	2	97%	127, 129
semantic_kernel/connectors/memory/azure_cosmos_db
azure_cosmos_db_no_sql_base.py	46	9	80%	75, 78–79, 85, 95–99
azure_cosmos_db_no_sql_collection.py	101	25	75%	98–103, 114–119, 123–130, 154, 187–200, 211–212
azure_cosmos_db_no_sql_store.py	30	6	80%	82–87
utils.py	58	4	93%	139, 153–156
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	95	2	98%	262–263
qdrant_store.py	48	2	96%	139–140
semantic_kernel/connectors/memory/redis
redis_collection.py	163	4	98%	148, 153–154, 324
redis_store.py	42	2	95%	108–109
utils.py	45	11	76%	145–146, 164, 166, 173–188
semantic_kernel/connectors/memory/weaviate
utils.py	61	4	93%	85–90, 253
weaviate_collection.py	130	32	75%	149–158, 162–182, 186–191, 275–280, 286
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
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_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	2	88%	41, 51
vector_store_record_collection.py	248	20	92%	65, 431, 491–495, 503–507, 547–550, 557–560
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	68	39	43%	43–52, 56–58, 64–74, 78–85, 89–92, 96–100, 105, 109–114
process_end_step.py	19	2	89%	37, 41
process_function_target_builder.py	25	3	88%	37–40
process_step_builder.py	105	24	77%	44, 89, 103, 110–123, 135–142, 151, 160–169, 178, 192, 209
process_step_edge_builder.py	35	3	91%	43, 58, 68
process_types.py	25	1	96%	35
semantic_kernel/processes/kernel_process
kernel_process_step_context.py	17	1	94%	37
semantic_kernel/processes/local_runtime
local_kernel_process.py	20	2	90%	23, 30
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	178	114	36%	61, 72, 81–169, 173, 177, 181–182, 187–249, 253–270, 274–277, 281–284, 288–297, 303–306, 310–312
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	14549	883	94%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
2845	4 💤	0 ❌	0 🔥	1m 7s ⏱️

[eavanvalkenburg]

couple of small comments

[eavanvalkenburg]

I've introduced a concept called managed_client in VectorStore and VectorStoreRecordCollection, this together with a context manager can be used to make sure we cleanup what we create, but not what someone else created, so have a look at Azure AI Search for the mechanics!

[eavanvalkenburg]

Azure AI Search actually uses two clients (one for the service, one for an index) so it introduces a additional one, that can be applied to the database creation

[eavanvalkenburg]

a vector field can also have Int as the datatype, which would map toint8: https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/vector-search#vector-indexing-policies