semantic-kernel 1.8.3__zip

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/redis/utils.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import contextlib
+import json
+from datetime import datetime
+from typing import Any
+
+import numpy as np
+from redis.asyncio.client import Redis
+from redis.commands.search.document import Document
+from redis.commands.search.field import Field as RedisField
+from redis.commands.search.field import NumericField, TagField, TextField, VectorField
+
+from semantic_kernel.connectors.memory.azure_ai_search.const import DISTANCE_FUNCTION_MAP
+from semantic_kernel.connectors.memory.redis.const import TYPE_MAPPER_VECTOR, RedisCollectionTypes
+from semantic_kernel.data.vector_store_model_definition import VectorStoreRecordDefinition
+from semantic_kernel.data.vector_store_record_fields import (
+    VectorStoreRecordDataField,
+    VectorStoreRecordKeyField,
+    VectorStoreRecordVectorField,
+)
+from semantic_kernel.memory.memory_record import MemoryRecord
+
+
+def get_redis_key(collection_name: str, record_id: str) -> str:  # pragma: no cover
+    """Returns the Redis key for an element called record_id within collection_name.
+
+    Args:
+        collection_name (str): Name for a collection of embeddings
+        record_id (str): ID associated with a memory record
+
+    Returns:
+        str: Redis key in the format collection_name:id
+    """
+    return f"{collection_name}:{record_id}"
+
+
+def split_redis_key(redis_key: str) -> tuple[str, str]:  # pragma: no cover
+    """Split a Redis key into its collection name and record ID.
+
+    Args:
+        redis_key (str): Redis key
+
+    Returns:
+        tuple[str, str]: Tuple of the collection name and ID
+    """
+    collection, record_id = redis_key.split(":")
+    return collection, record_id
+
+
+def serialize_record_to_redis(record: MemoryRecord, vector_type: np.dtype) -> dict[str, Any]:  # pragma: no cover
+    """Serialize a MemoryRecord to Redis fields."""
+    all_metadata = {
+        "is_reference": record._is_reference,
+        "external_source_name": record._external_source_name or "",
+        "id": record._id or "",
+        "description": record._description or "",
+        "text": record._text or "",
+        "additional_metadata": record._additional_metadata or "",
+    }
+
+    return {
+        "key": record._key or "",
+        "timestamp": record._timestamp.isoformat() if record._timestamp else "",
+        "metadata": json.dumps(all_metadata),
+        "embedding": (record._embedding.astype(vector_type).tobytes() if record._embedding is not None else ""),
+    }
+
+
+def deserialize_redis_to_record(
+    fields: dict[str, Any], vector_type: np.dtype, with_embedding: bool
+) -> MemoryRecord:  # pragma: no cover
+    """Deserialize Redis fields to a MemoryRecord."""
+    metadata = json.loads(fields[b"metadata"])
+    record = MemoryRecord(
+        id=metadata["id"],
+        is_reference=metadata["is_reference"] is True,
+        description=metadata["description"],
+        external_source_name=metadata["external_source_name"],
+        text=metadata["text"],
+        additional_metadata=metadata["additional_metadata"],
+        embedding=None,
+    )
+
+    if fields[b"timestamp"] != b"":
+        record._timestamp = datetime.fromisoformat(fields[b"timestamp"].decode())
+
+    if with_embedding:
+        # Extract using the vector type, then convert to regular Python float type
+        record._embedding = np.frombuffer(fields[b"embedding"], dtype=vector_type).astype(float)
+
+    return record
+
+
+def deserialize_document_to_record(
+    database: Redis, doc: Document, vector_type: np.dtype, with_embedding: bool
+) -> MemoryRecord:  # pragma: no cover
+    """Deserialize document to a MemoryRecord."""
+    # Document's ID refers to the Redis key
+    redis_key = doc["id"]
+    _, id_str = split_redis_key(redis_key)
+
+    metadata = json.loads(doc["metadata"])
+    record = MemoryRecord(
+        id=id_str,
+        is_reference=metadata["is_reference"] is True,
+        description=metadata["description"],
+        external_source_name=metadata["external_source_name"],
+        text=metadata["text"],
+        additional_metadata=metadata["additional_metadata"],
+        embedding=None,
+    )
+
+    if doc["timestamp"] != "":
+        record._timestamp = datetime.fromisoformat(doc["timestamp"])
+
+    if with_embedding:
+        # Some bytes are lost when retrieving a document, fetch raw embedding
+        eb = database.hget(redis_key, "embedding")
+        record._embedding = np.frombuffer(eb, dtype=vector_type).astype(float)
+
+    return record
+
+
+class RedisWrapper(Redis):
+    """Wrapper to make sure the connection is closed when the object is deleted."""
+
+    def __del__(self) -> None:
+        """Close connection, done when the object is deleted, used when SK creates a client."""
+        with contextlib.suppress(Exception):
+            asyncio.get_running_loop().create_task(self.close())
+
+
+def data_model_definition_to_redis_fields(
+    data_model_definition: VectorStoreRecordDefinition, collection_type: RedisCollectionTypes
+) -> list[RedisField]:
+    """Create a list of fields for Redis from a data_model_definition."""
+    fields: list[RedisField] = []
+    for name, field in data_model_definition.fields.items():
+        if isinstance(field, VectorStoreRecordKeyField):
+            continue
+        if collection_type == RedisCollectionTypes.HASHSET:
+            fields.append(_field_to_redis_field_hashset(name, field))
+        elif collection_type == RedisCollectionTypes.JSON:
+            fields.append(_field_to_redis_field_json(name, field))
+    return fields
+
+
+def _field_to_redis_field_hashset(
+    name: str, field: VectorStoreRecordVectorField | VectorStoreRecordDataField
+) -> RedisField:
+    if isinstance(field, VectorStoreRecordVectorField):
+        return VectorField(
+            name=name,
+            algorithm=field.index_kind.value.upper() if field.index_kind else "HNSW",
+            attributes={
+                "type": TYPE_MAPPER_VECTOR[field.property_type or "default"],
+                "dim": field.dimensions,
+                "distance_metric": DISTANCE_FUNCTION_MAP[field.distance_function or "default"],
+            },
+        )
+    if field.property_type in ["int", "float"]:
+        return NumericField(name=name)
+    if field.is_full_text_searchable:
+        return TextField(name=name)
+    return TagField(name=name)
+
+
+def _field_to_redis_field_json(
+    name: str, field: VectorStoreRecordVectorField | VectorStoreRecordDataField
+) -> RedisField:
+    if isinstance(field, VectorStoreRecordVectorField):
+        return VectorField(
+            name=f"$.{name}",
+            algorithm=field.index_kind.value.upper() if field.index_kind else "HNSW",
+            attributes={
+                "type": TYPE_MAPPER_VECTOR[field.property_type or "default"],
+                "dim": field.dimensions,
+                "distance_metric": DISTANCE_FUNCTION_MAP[field.distance_function or "default"],
+            },
+            as_name=name,
+        )
+    if field.property_type in ["int", "float"]:
+        return NumericField(name=f"$.{name}", as_name=name)
+    if field.is_full_text_searchable:
+        return TextField(name=f"$.{name}", as_name=name)
+    return TagField(name=f"$.{name}", as_name=name)

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/usearch/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from semantic_kernel.connectors.memory.usearch.usearch_memory_store import (
+    USearchMemoryStore,
+)
+
+__all__ = ["USearchMemoryStore"]