unique_toolkit 1.38.3__py3-none-any.whl → 1.39.0__py3-none-any.whl

unique_toolkit/services/chat_service.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Sequence
+from typing import Any, Sequence, overload
 
 import unique_sdk
 from openai.types.chat import ChatCompletionToolChoiceOptionParam
@@ -13,6 +13,7 @@ from openai.types.responses import (
     response_create_params,
 )
 from openai.types.shared_params import Metadata, Reasoning
+from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from unique_toolkit._common.utils.files import is_file_content, is_image_content
@@ -93,6 +94,13 @@ from unique_toolkit.language_model.schemas import (
     LanguageModelToolDescription,
     ResponsesLanguageModelStreamResponse,
 )
+from unique_toolkit.short_term_memory.functions import (
+    create_memory,
+    create_memory_async,
+    find_latest_memory,
+    find_latest_memory_async,
+)
+from unique_toolkit.short_term_memory.schemas import ShortTermMemory
 
 logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
 
@@ -1654,3 +1662,403 @@ class ChatService(ChatServiceDeprecated):
             if is_image_content(filename=c.key):
                 images.append(c)
         return images, files
+
+    # Short Term Memories
+    ############################################################################
+
+    def create_chat_memory_by_id(
+        self, *, chat_id: str, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for a specific chat synchronously.
+
+        Args:
+            chat_id (str): The chat ID
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        # Convert BaseModel to JSON string if needed
+        if isinstance(value, BaseModel):
+            value = value.model_dump_json()
+
+        return create_memory(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            value=value,
+            chat_id=chat_id,
+        )
+
+    async def create_chat_memory_by_id_async(
+        self, *, chat_id: str, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for a specific chat asynchronously.
+
+        Args:
+            chat_id (str): The chat ID
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        # Convert BaseModel to JSON string if needed
+        if isinstance(value, BaseModel):
+            value = value.model_dump_json()
+
+        return await create_memory_async(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            value=value,
+            chat_id=chat_id,
+        )
+
+    def create_message_memory_by_id(
+        self, *, message_id: str, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for a specific message synchronously.
+
+        Args:
+            message_id (str): The message ID
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        # Convert BaseModel to JSON string if needed
+        if isinstance(value, BaseModel):
+            value = value.model_dump_json()
+
+        return create_memory(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            value=value,
+            message_id=message_id,
+        )
+
+    async def create_message_memory_by_id_async(
+        self, *, message_id: str, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for a specific message asynchronously.
+
+        Args:
+            message_id (str): The message ID
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        # Convert BaseModel to JSON string if needed
+        if isinstance(value, BaseModel):
+            value = value.model_dump_json()
+
+        return await create_memory_async(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            value=value,
+            message_id=message_id,
+        )
+
+    def find_chat_memory_by_id(self, *, chat_id: str, key: str) -> ShortTermMemory:
+        """Finds the latest short-term memory for a specific chat synchronously.
+
+        Args:
+            chat_id (str): The chat ID
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return find_latest_memory(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            chat_id=chat_id,
+        )
+
+    async def find_chat_memory_by_id_async(
+        self, *, chat_id: str, key: str
+    ) -> ShortTermMemory:
+        """Finds the latest short-term memory for a specific chat asynchronously.
+
+        Args:
+            chat_id (str): The chat ID
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return await find_latest_memory_async(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            chat_id=chat_id,
+        )
+
+    def find_message_memory_by_id(
+        self, *, message_id: str, key: str
+    ) -> ShortTermMemory:
+        """Finds the latest short-term memory for a specific message synchronously.
+
+        Args:
+            message_id (str): The message ID
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return find_latest_memory(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            message_id=message_id,
+        )
+
+    async def find_message_memory_by_id_async(
+        self, *, message_id: str, key: str
+    ) -> ShortTermMemory:
+        """Finds the latest short-term memory for a specific message asynchronously.
+
+        Args:
+            message_id (str): The message ID
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return await find_latest_memory_async(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            key=key,
+            message_id=message_id,
+        )
+
+    # Convenience methods using current chat/message IDs
+    ############################################################################
+
+    def create_chat_memory(
+        self, *, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for the current chat synchronously.
+
+        Args:
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        return self.create_chat_memory_by_id(
+            chat_id=self._chat_id,
+            key=key,
+            value=value,
+        )
+
+    async def create_chat_memory_async(
+        self, *, key: str, value: str | dict | BaseModel
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for the current chat asynchronously.
+
+        Args:
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        return await self.create_chat_memory_by_id_async(
+            chat_id=self._chat_id,
+            key=key,
+            value=value,
+        )
+
+    @overload
+    def create_message_memory(
+        self,
+        *,
+        key: str,
+        value: str | dict | BaseModel,
+    ) -> ShortTermMemory: ...
+
+    @overload
+    def create_message_memory(
+        self, *, key: str, value: str | dict | BaseModel, message_id: str
+    ) -> ShortTermMemory: ...
+
+    def create_message_memory(
+        self, *, key: str, value: str | dict | BaseModel, message_id: str | None = None
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for the current assistant message synchronously.
+
+        Args:
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        return self.create_message_memory_by_id(
+            key=key,
+            value=value,
+            message_id=message_id or self._assistant_message_id,
+        )
+
+    @overload
+    async def create_message_memory_async(
+        self,
+        *,
+        key: str,
+        value: str | dict | BaseModel,
+    ) -> ShortTermMemory: ...
+
+    @overload
+    async def create_message_memory_async(
+        self, *, key: str, value: str | dict | BaseModel, message_id: str
+    ) -> ShortTermMemory: ...
+
+    async def create_message_memory_async(
+        self, *, key: str, value: str | dict | BaseModel, message_id: str | None = None
+    ) -> ShortTermMemory:
+        """Creates a short-term memory for the current assistant message asynchronously.
+
+        Args:
+            key (str): The memory key
+            value (str | dict | BaseModel): The memory value
+
+        Returns:
+            ShortTermMemory: The created short-term memory
+
+        Raises:
+            Exception: If the creation fails
+        """
+        return await self.create_message_memory_by_id_async(
+            message_id=message_id or self._assistant_message_id,
+            key=key,
+            value=value,
+        )
+
+    def find_chat_memory(self, *, key: str) -> ShortTermMemory:
+        """Finds the latest short-term memory for the current chat synchronously.
+
+        Args:
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return self.find_chat_memory_by_id(
+            chat_id=self._chat_id,
+            key=key,
+        )
+
+    async def find_chat_memory_async(self, *, key: str) -> ShortTermMemory:
+        """Finds the latest short-term memory for the current chat asynchronously.
+
+        Args:
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return await self.find_chat_memory_by_id_async(
+            chat_id=self._chat_id,
+            key=key,
+        )
+
+    @overload
+    def find_message_memory(self, *, key: str) -> ShortTermMemory: ...
+
+    @overload
+    def find_message_memory(self, *, key: str, message_id: str) -> ShortTermMemory: ...
+
+    def find_message_memory(
+        self, *, key: str, message_id: str | None = None
+    ) -> ShortTermMemory:
+        """Finds the latest short-term memory for the current assistant message synchronously.
+
+        Args:
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return self.find_message_memory_by_id(
+            message_id=message_id or self._assistant_message_id,
+            key=key,
+        )
+
+    @overload
+    async def find_message_memory_async(self, *, key: str) -> ShortTermMemory: ...
+
+    @overload
+    async def find_message_memory_async(
+        self, *, key: str, message_id: str
+    ) -> ShortTermMemory: ...
+
+    async def find_message_memory_async(
+        self, *, key: str, message_id: str | None = None
+    ) -> ShortTermMemory:
+        """Finds the latest short-term memory for the current assistant message asynchronously.
+
+        Args:
+            key (str): The memory key
+
+        Returns:
+            ShortTermMemory: The latest short-term memory
+
+        Raises:
+            Exception: If the retrieval fails
+        """
+        return await self.find_message_memory_by_id_async(
+            message_id=message_id or self._assistant_message_id,
+            key=key,
+        )

unique_toolkit/short_term_memory/functions.py

@@ -2,6 +2,7 @@ import json
 import logging
 
 import unique_sdk
+from typing_extensions import Any
 
 from unique_toolkit.short_term_memory.constants import DOMAIN_NAME
 from unique_toolkit.short_term_memory.schemas import ShortTermMemory
@@ -41,7 +42,7 @@ async def find_latest_memory_async(
             chatId=chat_id,
             messageId=message_id,
         )
-        return ShortTermMemory(**stm)
+        return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
     except Exception as e:
         logger.error(f"Error finding latest short term memory: {e}")
         raise e
@@ -79,7 +80,7 @@ def find_latest_memory(
             chatId=chat_id,
             messageId=message_id,
         )
-        return ShortTermMemory(**stm)
+        return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
     except Exception as e:
         logger.error(f"Error finding latest short term memory: {e}")
         raise e
@@ -124,13 +125,14 @@ async def create_memory_async(
             messageId=message_id,
             data=value,
         )
-        return ShortTermMemory(**stm)
+        return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
     except Exception as e:
         logger.error(f"Error creating short term memory: {e}")
         raise e
 
 
 def create_memory(
+    *,
     user_id: str,
     company_id: str,
     key: str,
@@ -169,7 +171,127 @@ def create_memory(
             messageId=message_id,
             data=value,
         )
-        return ShortTermMemory(**stm)
+        return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
     except Exception as e:
         logger.error(f"Error creating short term memory: {e}")
         raise e
+
+
+def find_last_chat_memory(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    chat_id: str,
+) -> ShortTermMemory:
+    """
+    Find the last chat short term memory.
+    """
+    stm = unique_sdk.ShortTermMemory.find_latest(
+        user_id=user_id,
+        company_id=company_id,
+        memoryName=key,
+        chatId=chat_id,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
+
+
+async def find_last_chat_memory_async(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    chat_id: str,
+) -> ShortTermMemory:
+    """
+    Find the last chat short term memory.
+    """
+    stm = await unique_sdk.ShortTermMemory.find_latest_async(
+        user_id=user_id,
+        company_id=company_id,
+        memoryName=key,
+        chatId=chat_id,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
+
+
+def find_last_message_memory(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    message_id: str,
+) -> ShortTermMemory:
+    """
+    Find the last message short term memory.
+    """
+    stm = unique_sdk.ShortTermMemory.find_latest(
+        user_id=user_id,
+        company_id=company_id,
+        chatId=None,
+        memoryName=key,
+        messageId=message_id,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
+
+
+async def find_last_message_memory_async(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    message_id: str,
+) -> ShortTermMemory:
+    """
+    Find the last message short term memory.
+    """
+    stm = await unique_sdk.ShortTermMemory.find_latest_async(
+        user_id=user_id,
+        company_id=company_id,
+        chatId=None,
+        memoryName=key,
+        messageId=message_id,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
+
+
+def create_chat_memory(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    value: dict[str, Any] | str,
+    chat_id: str,
+) -> ShortTermMemory:
+    """
+    Create a chat short term memory.
+    """
+    stm = unique_sdk.ShortTermMemory.create(
+        user_id=user_id,
+        company_id=company_id,
+        memoryName=key,
+        chatId=chat_id,
+        data=json.dumps(value) if isinstance(value, dict) else value,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)
+
+
+async def create_chat_memory_async(
+    *,
+    user_id: str,
+    company_id: str,
+    key: str,
+    value: dict[str, Any] | str,
+    chat_id: str,
+) -> ShortTermMemory:
+    """
+    Create a chat short term memory.
+    """
+    stm = await unique_sdk.ShortTermMemory.create_async(
+        user_id=user_id,
+        company_id=company_id,
+        memoryName=key,
+        chatId=chat_id,
+        data=json.dumps(value) if isinstance(value, dict) else value,
+    )
+    return ShortTermMemory.model_validate(stm, by_alias=True, by_name=True)

unique_toolkit/short_term_memory/schemas.py

@@ -1,7 +1,8 @@
 import json
+from typing import Any
 
 from humps import camelize
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 model_config = ConfigDict(
     alias_generator=camelize,
@@ -17,27 +18,31 @@ class ShortTermMemory(BaseModel):
     key: str = Field(alias="object")
     chat_id: str | None
     message_id: str | None
-    data: str | dict | int | float | bool | list | None
-
-    @field_validator("chat_id", "message_id", mode="before")
-    def validate_chat_id_or_message_id(cls, v, info):
-        field_name = info.field_name
-        data = info.data
-
-        # Get the other field's value
-        other_field = "message_id" if field_name == "chat_id" else "chat_id"
-        other_value = data.get(other_field)
-
-        # Check if both are None
-        if v is None and other_value is None:
-            camel_name = camelize(field_name)
-            raise ValueError(
-                f"Either {camel_name} or messageId must be provided"
-                if field_name == "chat_id"
-                else f"Either chatId or {camel_name} must be provided"
-            )
-
-        return v
+    data: str | dict | int | float | bool | list | None = Field(deprecated=True)
+    value: str | dict[str, Any] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def _data_to_value(self) -> "ShortTermMemory":
+        if isinstance(self.data, dict):
+            self.value = self.data
+        elif isinstance(self.data, str):
+            try:
+                self.value = json.loads(self.data)
+            except json.JSONDecodeError:
+                self.value = self.data
+        elif self.data is None:
+            self.value = ""
+        else:
+            self.value = str(self.data)
+        return self
+
+    @model_validator(mode="after")
+    def validate_message_id_and_chat_id(self):
+        if (self.message_id is None and self.chat_id is None) or (
+            self.message_id is not None and self.chat_id is not None
+        ):
+            raise ValueError("Either message_id or chat_id must be provided")
+        return self
 
     @field_validator("data", mode="before")
     def validate_data(cls, v):