Skip to content

Commit 9a883be

Browse files
eavanvalkenburgeavanvalkenburg
committed
Python: samples and misc (#12123)
### 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. --> Final update with last pieces and the samples ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [ ] The code builds clean without any errors or warnings - [ ] 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 - [ ] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄
1 parent 48494a3 commit 9a883be

10 files changed

Lines changed: 235 additions & 269 deletions

File tree

‎python/samples/concepts/chat_history/store_chat_history_in_cosmosdb.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from samples.concepts.setup.chat_completion_services import Services, get_chat_completion_service_and_request_settings
88
from semantic_kernel import Kernel
99
from semantic_kernel.connectors.ai import FunctionChoiceBehavior
10-
from semantic_kernel.connectors.memory.azure_cosmos_db.azure_cosmos_db_no_sql_store import AzureCosmosDBNoSQLStore
10+
from semantic_kernel.connectors.memory.azure_cosmos_db import AzureCosmosDBNoSQLStore
1111
from semantic_kernel.contents import ChatHistory, ChatMessageContent
1212
from semantic_kernel.core_plugins.math_plugin import MathPlugin
1313
from semantic_kernel.core_plugins.time_plugin import TimePlugin
@@ -40,8 +40,8 @@
4040
@dataclass
4141
class ChatHistoryModel:
4242
session_id: Annotated[str, VectorStoreRecordKeyField]
43-
user_id: Annotated[str, VectorStoreRecordDataField(is_filterable=True)]
44-
messages: Annotated[list[dict[str, str]], VectorStoreRecordDataField(is_filterable=True)]
43+
user_id: Annotated[str, VectorStoreRecordDataField(is_indexed=True)]
44+
messages: Annotated[list[dict[str, str]], VectorStoreRecordDataField(is_indexed=True)]
4545

4646

4747
# 2. We then create a class that extends the ChatHistory class

‎python/samples/concepts/memory/complex_memory.py‎

Lines changed: 75 additions & 180 deletions
Large diffs are not rendered by default.

‎python/samples/concepts/memory/memory_with_pandas.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,15 @@
66
import pandas as pd
77

88
from semantic_kernel import Kernel
9-
from semantic_kernel.connectors.ai.open_ai import (
10-
OpenAIEmbeddingPromptExecutionSettings,
11-
OpenAITextEmbedding,
12-
)
9+
from semantic_kernel.connectors.ai.open_ai import OpenAIEmbeddingPromptExecutionSettings, OpenAITextEmbedding
1310
from semantic_kernel.connectors.memory.azure_ai_search import AzureAISearchCollection
1411
from semantic_kernel.data import (
1512
VectorStoreRecordDataField,
1613
VectorStoreRecordDefinition,
1714
VectorStoreRecordKeyField,
1815
VectorStoreRecordVectorField,
1916
)
20-
from semantic_kernel.data.vector_search import add_vector_to_records
17+
from semantic_kernel.data.vector_storage import add_vector_to_records
2118

2219
model_fields = VectorStoreRecordDefinition(
2320
container_mode=True,

‎python/samples/concepts/memory/simple_memory.py‎

Lines changed: 22 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -8,30 +8,21 @@
88

99
from samples.concepts.memory.utils import print_record
1010
from samples.concepts.resources.utils import Colors, print_with_color
11-
from semantic_kernel import Kernel
12-
from semantic_kernel.connectors.ai.open_ai import (
13-
OpenAIEmbeddingPromptExecutionSettings,
14-
OpenAITextEmbedding,
15-
)
16-
from semantic_kernel.connectors.memory.in_memory import InMemoryVectorCollection
11+
from semantic_kernel.connectors.ai.open_ai import OpenAITextEmbedding
12+
from semantic_kernel.connectors.memory import InMemoryCollection
1713
from semantic_kernel.data import (
18-
VectorSearchFilter,
1914
VectorSearchOptions,
2015
VectorStoreRecordDataField,
2116
VectorStoreRecordKeyField,
2217
VectorStoreRecordVectorField,
2318
vectorstoremodel,
2419
)
25-
from semantic_kernel.data.const import DISTANCE_FUNCTION_DIRECTION_HELPER, DistanceFunction, IndexKind
26-
from semantic_kernel.data.vector_search import add_vector_to_records
2720

2821
# This is the most basic example of a vector store and collection
2922
# For a more complex example, using different collection types, see "complex_memory.py"
3023
# This sample uses openai text embeddings, so make sure to have your environment variables set up
3124
# it needs openai api key and embedding model id
32-
kernel = Kernel()
3325
embedder = OpenAITextEmbedding(service_id="embedding")
34-
kernel.add_service(embedder)
3526

3627
# Next, you need to define your data structure
3728
# In this case, we are using a dataclass to define our data structure
@@ -41,36 +32,22 @@
4132
# This has been done in constants here for simplicity, but you can also define them in the model itself
4233
# Next we create three records using that model
4334

44-
DISTANCE_FUNCTION = DistanceFunction.COSINE_SIMILARITY
45-
# The in memory collection does not actually use a index, so this variable is not relevant, here for completeness
46-
INDEX_KIND = IndexKind.IVF_FLAT
47-
4835

49-
@vectorstoremodel
36+
@vectorstoremodel(collection_name="test")
5037
@dataclass
5138
class DataModel:
39+
content: Annotated[str, VectorStoreRecordDataField()]
40+
id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
5241
vector: Annotated[
53-
list[float] | None,
54-
VectorStoreRecordVectorField(
55-
embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings()},
56-
index_kind=INDEX_KIND,
57-
dimensions=1536,
58-
distance_function=DISTANCE_FUNCTION,
59-
property_type="float",
60-
),
42+
str | list[float] | None,
43+
VectorStoreRecordVectorField(dimensions=1536, property_type="float"),
6144
] = None
62-
id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
63-
content: Annotated[
64-
str,
65-
VectorStoreRecordDataField(
66-
has_embedding=True,
67-
embedding_property_name="vector",
68-
property_type="str",
69-
is_full_text_searchable=True,
70-
),
71-
] = "content1"
72-
title: Annotated[str, VectorStoreRecordDataField(property_type="str", is_full_text_searchable=True)] = "title"
73-
tag: Annotated[str, VectorStoreRecordDataField(property_type="str", is_filterable=True)] = "tag"
45+
title: Annotated[str, VectorStoreRecordDataField(property_type="str", is_full_text_indexed=True)] = "title"
46+
tag: Annotated[str, VectorStoreRecordDataField(property_type="str", is_indexed=True)] = "tag"
47+
48+
def __post_init__(self):
49+
if self.vector is None:
50+
self.vector = self.content
7451

7552

7653
records = [
@@ -102,9 +79,9 @@ async def main():
10279
# we also use the async with to open and close the connection
10380
# for the in memory collection, this is just a no-op
10481
# but for other collections, like Azure AI Search, this will open and close the connection
105-
async with InMemoryVectorCollection[str, DataModel](
106-
collection_name="test",
82+
async with InMemoryCollection[str, DataModel](
10783
data_model_type=DataModel,
84+
embedding_generator=embedder,
10885
) as record_collection:
10986
# Create the collection after wiping it
11087
print_with_color("Creating test collection!", Colors.CGREY)
@@ -113,15 +90,13 @@ async def main():
11390

11491
# First add vectors to the records
11592
print_with_color("Adding records!", Colors.CBLUE)
116-
records_with_embedding = await add_vector_to_records(kernel, records, data_model_type=DataModel)
117-
# Next upsert them to the store.
118-
keys = await record_collection.upsert_batch(records_with_embedding)
93+
keys = await record_collection.upsert(records)
11994
print(f" Upserted {keys=}")
12095
print("-" * 30)
12196

12297
# Now we can get the records back
12398
print_with_color("Getting records!", Colors.CBLUE)
124-
results = await record_collection.get_batch([records[0].id, records[1].id, records[2].id])
99+
results = await record_collection.get([records[0].id, records[1].id, records[2].id])
125100
if results and isinstance(results, Sequence):
126101
[print_record(record=result) for result in results]
127102
else:
@@ -134,20 +109,17 @@ async def main():
134109
# The other options are optional, but can be useful
135110
# The filter option is used to filter the results based on the tag field
136111
options = VectorSearchOptions(
137-
vector_field_name="vector",
138-
include_vectors=True,
139-
filter=VectorSearchFilter.equal_to("tag", "general"),
112+
vector_property_name="vector",
113+
filter=lambda x: x.tag == "general",
140114
)
141115
query = "python"
142116
print_with_color(f"Searching for '{query}', with filter 'tag == general'", Colors.CBLUE)
143117
print_with_color(
144-
f"Using vectorized search, for {DISTANCE_FUNCTION.value}, "
145-
f"the {'higher' if DISTANCE_FUNCTION_DIRECTION_HELPER[DISTANCE_FUNCTION](1, 0) else 'lower'} the score the better" # noqa: E501
146-
f"",
118+
"Using vectorized search, the lower the score the better",
147119
Colors.CBLUE,
148120
)
149-
search_results = await record_collection.vectorized_search(
150-
vector=(await embedder.generate_raw_embeddings([query]))[0],
121+
search_results = await record_collection.search(
122+
values=query,
151123
options=options,
152124
)
153125
if search_results.total_count == 0:

‎python/samples/concepts/rag/rag_with_vector_collection.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
import asyncio
3+
from dataclasses import dataclass
4+
from typing import Annotated
5+
6+
from semantic_kernel import Kernel
7+
from semantic_kernel.connectors.ai import FunctionChoiceBehavior
8+
from semantic_kernel.connectors.ai.open_ai import (
9+
OpenAIChatCompletion,
10+
OpenAIChatPromptExecutionSettings,
11+
OpenAITextEmbedding,
12+
)
13+
from semantic_kernel.connectors.memory import InMemoryCollection
14+
from semantic_kernel.data import (
15+
VectorStoreRecordDataField,
16+
VectorStoreRecordKeyField,
17+
VectorStoreRecordVectorField,
18+
vectorstoremodel,
19+
)
20+
from semantic_kernel.functions import KernelArguments
21+
22+
"""
23+
This sample shows a really easy way to have RAG with a vector store.
24+
It creates a simple datamodel, and then creates a collection with that datamodel.
25+
Then we create a function that can search the collection.
26+
Finally, in two different ways we call the function to search the collection.
27+
"""
28+
29+
30+
# Define a data model for the collection
31+
# This model will be used to store the information in the collection
32+
@vectorstoremodel(collection_name="budget")
33+
@dataclass
34+
class BudgetItem:
35+
id: Annotated[str, VectorStoreRecordKeyField]
36+
text: Annotated[str, VectorStoreRecordDataField]
37+
embedding: Annotated[
38+
list[float] | str | None,
39+
VectorStoreRecordVectorField(dimensions=1536, embedding_generator=OpenAITextEmbedding()),
40+
] = None
41+
42+
def __post_init__(self):
43+
if self.embedding is None:
44+
self.embedding = self.text
45+
46+
47+
async def main():
48+
kernel = Kernel()
49+
50+
kernel.add_service(OpenAIChatCompletion())
51+
52+
async with InMemoryCollection(data_model_type=BudgetItem) as collection:
53+
# Add information to the collection
54+
await collection.upsert(
55+
[
56+
BudgetItem(id="info1", text="My budget for 2022 is $50,000"),
57+
BudgetItem(id="info1", text="My budget for 2023 is $75,000"),
58+
BudgetItem(id="info1", text="My budget for 2024 is $100,000"),
59+
BudgetItem(id="info2", text="My budget for 2025 is $150,000"),
60+
],
61+
)
62+
# Create a function to search the collection
63+
# note the string_mapper, this is used to map the result of the search to a string
64+
kernel.add_function(
65+
"memory",
66+
collection.create_search_function(
67+
function_name="recall",
68+
description="Recalls the budget information.",
69+
string_mapper=lambda x: x.record.text,
70+
),
71+
)
72+
73+
# Call the search function directly from from a template.
74+
result = await kernel.invoke_prompt(
75+
function_name="budget",
76+
plugin_name="BudgetPlugin",
77+
prompt="{{memory.recall 'budget by year'}} What is my budget for 2024?",
78+
)
79+
print("Called from template")
80+
print(result)
81+
print("======================")
82+
# Let the LLM choose the function to call
83+
result = await kernel.invoke_prompt(
84+
function_name="budget",
85+
plugin_name="BudgetPlugin",
86+
prompt="What is my budget for 2024?",
87+
arguments=KernelArguments(
88+
settings=OpenAIChatPromptExecutionSettings(
89+
function_choice_behavior=FunctionChoiceBehavior.Auto(),
90+
),
91+
),
92+
)
93+
print("Called from LLM")
94+
print(result)
95+
96+
97+
"""
98+
Output:
99+
100+
Called from template
101+
Your budget for 2024 is $100,000.
102+
======================
103+
Called from LLM
104+
Your budget for 2024 is $100,000.
105+
106+
"""
107+
108+
if __name__ == "__main__":
109+
asyncio.run(main())

‎python/tests/samples/test_concepts.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@
4646
from samples.concepts.prompt_templates.configuring_prompts import main as configuring_prompts
4747
from samples.concepts.prompt_templates.load_yaml_prompt import main as load_yaml_prompt
4848
from samples.concepts.prompt_templates.template_language import main as template_language
49-
from samples.concepts.rag.rag_with_text_memory_plugin import main as rag_with_text_memory_plugin
49+
from samples.concepts.rag.rag_with_vector_collection import main as rag_with_text_memory_plugin
5050
from samples.concepts.service_selector.custom_service_selector import main as custom_service_selector
5151
from samples.concepts.text_completion.text_completion import main as text_completion
5252
from samples.getting_started_with_agents.chat_completion.step01_chat_completion_agent_simple import (

‎python/tests/unit/connectors/memory/mongodb_atlas/test_mongodb_atlas_collection.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@
77
from pymongo.results import UpdateResult
88
from pytest import mark, raises
99

10-
from semantic_kernel.connectors.memory.mongodb_atlas.const import DEFAULT_DB_NAME, DEFAULT_SEARCH_INDEX_NAME
11-
from semantic_kernel.connectors.memory.mongodb_atlas.mongodb_atlas_collection import MongoDBAtlasCollection
10+
from semantic_kernel.connectors.memory.mongodb import DEFAULT_DB_NAME, DEFAULT_SEARCH_INDEX_NAME, MongoDBAtlasCollection
1211
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreInitializationException
1312

1413

‎python/tests/unit/connectors/memory/mongodb_atlas/test_mongodb_atlas_store.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@
33

44
from pymongo import AsyncMongoClient
55

6-
from semantic_kernel.connectors.memory.mongodb_atlas.mongodb_atlas_collection import MongoDBAtlasCollection
7-
from semantic_kernel.connectors.memory.mongodb_atlas.mongodb_atlas_store import MongoDBAtlasStore
6+
from semantic_kernel.connectors.memory.mongodb import MongoDBAtlasCollection, MongoDBAtlasStore
87

98

109
def test_mongodb_atlas_store_initialization(mongodb_atlas_unit_test_env):

‎python/tests/unit/connectors/memory/weaviate/test_weaviate_collection.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from weaviate.collections.classes.config_vectorizers import VectorDistances
99
from weaviate.collections.classes.data import DataObject
1010

11-
from semantic_kernel.connectors.memory.weaviate.weaviate_collection import WeaviateCollection
11+
from semantic_kernel.connectors.memory.weaviate import WeaviateCollection
1212
from semantic_kernel.exceptions import (
1313
ServiceInvalidExecutionSettingsError,
1414
VectorStoreInitializationException,
@@ -17,7 +17,7 @@
1717

1818

1919
@patch(
20-
"semantic_kernel.connectors.memory.weaviate.weaviate_collection.use_async_with_weaviate_cloud",
20+
"semantic_kernel.connectors.memory.weaviate.use_async_with_weaviate_cloud",
2121
return_value=AsyncMock(spec=WeaviateAsyncClient),
2222
)
2323
def test_weaviate_collection_init_with_weaviate_cloud(
@@ -47,7 +47,7 @@ def test_weaviate_collection_init_with_weaviate_cloud(
4747

4848

4949
@patch(
50-
"semantic_kernel.connectors.memory.weaviate.weaviate_collection.use_async_with_local",
50+
"semantic_kernel.connectors.memory.weaviate.use_async_with_local",
5151
return_value=AsyncMock(spec=WeaviateAsyncClient),
5252
)
5353
def test_weaviate_collection_init_with_local(
@@ -73,7 +73,7 @@ def test_weaviate_collection_init_with_local(
7373

7474

7575
@patch(
76-
"semantic_kernel.connectors.memory.weaviate.weaviate_collection.use_async_with_embedded",
76+
"semantic_kernel.connectors.memory.weaviate.use_async_with_embedded",
7777
return_value=AsyncMock(spec=WeaviateAsyncClient),
7878
)
7979
def test_weaviate_collection_init_with_embedded(
@@ -153,7 +153,7 @@ def test_weaviate_collection_init_with_custom_client(
153153

154154

155155
@patch(
156-
"semantic_kernel.connectors.memory.weaviate.weaviate_collection.use_async_with_local",
156+
"semantic_kernel.connectors.memory.weaviate.use_async_with_local",
157157
side_effect=Exception,
158158
)
159159
def test_weaviate_collection_init_fail_to_create_client(
@@ -175,7 +175,7 @@ def test_weaviate_collection_init_fail_to_create_client(
175175

176176

177177
@patch(
178-
"semantic_kernel.connectors.memory.weaviate.weaviate_collection.use_async_with_weaviate_cloud",
178+
"semantic_kernel.connectors.memory.weaviate.use_async_with_weaviate_cloud",
179179
return_value=AsyncMock(spec=WeaviateAsyncClient),
180180
)
181181
def test_weaviate_collection_init_with_lower_case_collection_name(
@@ -424,6 +424,6 @@ async def test_weaviate_collection_deserialize_data(
424424
)
425425

426426
with patch.object(collection, "_inner_get", return_value=[weaviate_data_object]) as mock_inner_get:
427-
await collection.get(data.id)
427+
await collection.get(key=data.id)
428428

429-
mock_inner_get.assert_called_once_with([data.id], include_vectors=True)
429+
mock_inner_get.assert_called_once_with([data.id], include_vectors=False, options=None)

0 commit comments

Comments
 (0)