wehub-resource-sync
132 lines
5.2 KiB
Python
132 lines
5.2 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
import asyncio
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass, field
|
|
from typing import Annotated
|
|
from uuid import uuid4
|
|
|
|
from samples.concepts.memory.utils import print_record
|
|
from samples.concepts.resources.utils import Colors, print_with_color
|
|
from semantic_kernel.connectors.ai.open_ai import OpenAITextEmbedding
|
|
from semantic_kernel.connectors.in_memory import InMemoryCollection
|
|
from semantic_kernel.data.vector import VectorStoreField, vectorstoremodel
|
|
|
|
# This is the most basic example of a vector store and collection
|
|
# For a more complex example, using different collection types, see "complex_memory.py"
|
|
# This sample uses openai text embeddings, so make sure to have your environment variables set up
|
|
# it needs openai api key and embedding model id
|
|
embedder = OpenAITextEmbedding(service_id="embedding")
|
|
|
|
# Next, you need to define your data structure
|
|
# In this case, we are using a dataclass to define our data structure
|
|
# you can also use a pydantic model, or a vanilla python class, see "data_models.py" for more examples
|
|
# Inside the model we define which fields we want to use, and which fields are vectors
|
|
# and for vector fields we define what kind of index we want to use, and what distance function we want to use
|
|
# This has been done in constants here for simplicity, but you can also define them in the model itself
|
|
# Next we create three records using that model
|
|
|
|
|
|
@vectorstoremodel(collection_name="test")
|
|
@dataclass
|
|
class DataModel:
|
|
content: Annotated[str, VectorStoreField("data")]
|
|
id: Annotated[str, VectorStoreField("key")] = field(default_factory=lambda: str(uuid4()))
|
|
vector: Annotated[
|
|
list[float] | str | None,
|
|
VectorStoreField("vector", dimensions=1536),
|
|
] = None
|
|
title: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] = "title"
|
|
tag: Annotated[str, VectorStoreField("data", is_indexed=True)] = "tag"
|
|
|
|
def __post_init__(self):
|
|
if self.vector is None:
|
|
self.vector = self.content
|
|
|
|
|
|
records = [
|
|
DataModel(
|
|
content="Semantic Kernel is awesome",
|
|
id="e6103c03-487f-4d7d-9c23-4723651c17f4",
|
|
title="Overview",
|
|
tag="general",
|
|
),
|
|
DataModel(
|
|
content="Semantic Kernel is available in dotnet, python and Java.",
|
|
id="09caec77-f7e1-466a-bcec-f1d51c5b15be",
|
|
title="Semantic Kernel Languages",
|
|
tag="general",
|
|
),
|
|
DataModel(
|
|
content="```python\nfrom semantic_kernel import Kernel\nkernel = Kernel()\n```",
|
|
id="d5c9913a-e015-4944-b960-5d4a84bca002",
|
|
title="Code sample",
|
|
tag="code",
|
|
),
|
|
]
|
|
|
|
|
|
async def main():
|
|
print("-" * 30)
|
|
# Create the collection here
|
|
# by using the generic we make sure that IDE's understand what you need to pass in and get back
|
|
# we also use the async with to open and close the connection
|
|
# for the in memory collection, this is just a no-op
|
|
# but for other collections, like Azure AI Search, this will open and close the connection
|
|
async with InMemoryCollection[str, DataModel](
|
|
record_type=DataModel,
|
|
embedding_generator=embedder,
|
|
) as record_collection:
|
|
# Create the collection after wiping it
|
|
print_with_color("Creating test collection!", Colors.CGREY)
|
|
await record_collection.ensure_collection_exists()
|
|
|
|
# First add vectors to the records
|
|
print_with_color("Adding records!", Colors.CBLUE)
|
|
keys = await record_collection.upsert(records)
|
|
print(f" Upserted {keys=}")
|
|
print("-" * 30)
|
|
|
|
# Now we can get the records back
|
|
print_with_color("Getting records!", Colors.CBLUE)
|
|
results = await record_collection.get([records[0].id, records[1].id, records[2].id])
|
|
if results and isinstance(results, Sequence):
|
|
[print_record(record=result) for result in results]
|
|
else:
|
|
print("Nothing found...")
|
|
print("-" * 30)
|
|
|
|
# Now we can search for records
|
|
# First we define the options
|
|
# The most important option is the vector_field_name, which is the name of the field that contains the vector
|
|
# The other options are optional, but can be useful
|
|
# The filter option is used to filter the results based on the tag field
|
|
options = {
|
|
"vector_property_name": "vector",
|
|
"filter": lambda x: x.tag == "general",
|
|
}
|
|
query = "python"
|
|
print_with_color(f"Searching for '{query}', with filter 'tag == general'", Colors.CBLUE)
|
|
print_with_color(
|
|
"Using vectorized search, the lower the score the better",
|
|
Colors.CBLUE,
|
|
)
|
|
search_results = await record_collection.search(
|
|
values=query,
|
|
**options,
|
|
)
|
|
if search_results.total_count == 0:
|
|
print("\nNothing found...\n")
|
|
else:
|
|
[print_record(result) async for result in search_results.results]
|
|
print("-" * 30)
|
|
|
|
# lets cleanup!
|
|
print_with_color("Deleting collection!", Colors.CBLUE)
|
|
await record_collection.ensure_collection_deleted()
|
|
print_with_color("Done!", Colors.CGREY)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|