hindsight-client 0.1.16__tar.gz → 0.2.1__tar.gz

hindsight_client-0.2.1/.openapi-generator/FILES

@@ -0,0 +1,58 @@
+hindsight_client_api/__init__.py
+hindsight_client_api/api/__init__.py
+hindsight_client_api/api/banks_api.py
+hindsight_client_api/api/documents_api.py
+hindsight_client_api/api/entities_api.py
+hindsight_client_api/api/memory_api.py
+hindsight_client_api/api/monitoring_api.py
+hindsight_client_api/api/operations_api.py
+hindsight_client_api/api_client.py
+hindsight_client_api/api_response.py
+hindsight_client_api/configuration.py
+hindsight_client_api/exceptions.py
+hindsight_client_api/models/__init__.py
+hindsight_client_api/models/add_background_request.py
+hindsight_client_api/models/background_response.py
+hindsight_client_api/models/bank_list_item.py
+hindsight_client_api/models/bank_list_response.py
+hindsight_client_api/models/bank_profile_response.py
+hindsight_client_api/models/bank_stats_response.py
+hindsight_client_api/models/budget.py
+hindsight_client_api/models/cancel_operation_response.py
+hindsight_client_api/models/chunk_data.py
+hindsight_client_api/models/chunk_include_options.py
+hindsight_client_api/models/chunk_response.py
+hindsight_client_api/models/create_bank_request.py
+hindsight_client_api/models/delete_document_response.py
+hindsight_client_api/models/delete_response.py
+hindsight_client_api/models/disposition_traits.py
+hindsight_client_api/models/document_response.py
+hindsight_client_api/models/entity_detail_response.py
+hindsight_client_api/models/entity_include_options.py
+hindsight_client_api/models/entity_input.py
+hindsight_client_api/models/entity_list_item.py
+hindsight_client_api/models/entity_list_response.py
+hindsight_client_api/models/entity_observation_response.py
+hindsight_client_api/models/entity_state_response.py
+hindsight_client_api/models/graph_data_response.py
+hindsight_client_api/models/http_validation_error.py
+hindsight_client_api/models/include_options.py
+hindsight_client_api/models/list_documents_response.py
+hindsight_client_api/models/list_memory_units_response.py
+hindsight_client_api/models/memory_item.py
+hindsight_client_api/models/operation_response.py
+hindsight_client_api/models/operations_list_response.py
+hindsight_client_api/models/recall_request.py
+hindsight_client_api/models/recall_response.py
+hindsight_client_api/models/recall_result.py
+hindsight_client_api/models/reflect_fact.py
+hindsight_client_api/models/reflect_include_options.py
+hindsight_client_api/models/reflect_request.py
+hindsight_client_api/models/reflect_response.py
+hindsight_client_api/models/retain_request.py
+hindsight_client_api/models/retain_response.py
+hindsight_client_api/models/update_disposition_request.py
+hindsight_client_api/models/validation_error.py
+hindsight_client_api/models/validation_error_loc_inner.py
+hindsight_client_api/rest.py
+hindsight_client_api_README.md

hindsight_client-0.2.1/.openapi-generator/VERSION

@@ -0,0 +1 @@
+7.10.0

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hindsight-client
-Version: 0.1.16
+Version: 0.2.1
 Summary: Python client for Hindsight - Semantic memory system with personality-driven thinking
 Author: Hindsight Team
 Requires-Python: >=3.10

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client/hindsight_client.py

@@ -112,6 +112,7 @@ class Hindsight:
         context: Optional[str] = None,
         document_id: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
+        entities: Optional[List[Dict[str, str]]] = None,
     ) -> RetainResponse:
         """
         Store a single memory (simplified interface).
@@ -123,13 +124,14 @@ class Hindsight:
             context: Optional context description
             document_id: Optional document ID for grouping
             metadata: Optional user-defined metadata
+            entities: Optional list of entities [{"text": "...", "type": "..."}]
 
         Returns:
             RetainResponse with success status
         """
         return self.retain_batch(
             bank_id=bank_id,
-            items=[{"content": content, "timestamp": timestamp, "context": context, "metadata": metadata}],
+            items=[{"content": content, "timestamp": timestamp, "context": context, "metadata": metadata, "entities": entities}],
             document_id=document_id,
         )
 
@@ -145,24 +147,34 @@ class Hindsight:
 
         Args:
             bank_id: The memory bank ID
-            items: List of memory items with 'content' and optional 'timestamp', 'context', 'metadata', 'document_id'
+            items: List of memory items with 'content' and optional 'timestamp', 'context', 'metadata', 'document_id', 'entities'
             document_id: Optional document ID for grouping memories (applied to items that don't have their own)
             retain_async: If True, process asynchronously in background (default: False)
 
         Returns:
             RetainResponse with success status and item count
         """
-        memory_items = [
-            memory_item.MemoryItem(
-                content=item["content"],
-                timestamp=item.get("timestamp"),
-                context=item.get("context"),
-                metadata=item.get("metadata"),
-                # Use item's document_id if provided, otherwise fall back to batch-level document_id
-                document_id=item.get("document_id") or document_id,
+        from hindsight_client_api.models.entity_input import EntityInput
+
+        memory_items = []
+        for item in items:
+            entities = None
+            if item.get("entities"):
+                entities = [
+                    EntityInput(text=e["text"], type=e.get("type"))
+                    for e in item["entities"]
+                ]
+            memory_items.append(
+                memory_item.MemoryItem(
+                    content=item["content"],
+                    timestamp=item.get("timestamp"),
+                    context=item.get("context"),
+                    metadata=item.get("metadata"),
+                    # Use item's document_id if provided, otherwise fall back to batch-level document_id
+                    document_id=item.get("document_id") or document_id,
+                    entities=entities,
+                )
             )
-            for item in items
-        ]
 
         request_obj = retain_request.RetainRequest(
             items=memory_items,
@@ -229,6 +241,8 @@ class Hindsight:
         query: str,
         budget: str = "low",
         context: Optional[str] = None,
+        max_tokens: Optional[int] = None,
+        response_schema: Optional[Dict[str, Any]] = None,
     ) -> ReflectResponse:
         """
         Generate a contextual answer based on bank identity and memories.
@@ -238,14 +252,21 @@ class Hindsight:
             query: The question or prompt
             budget: Budget level for reflection - "low", "mid", or "high" (default: "low")
             context: Optional additional context
+            max_tokens: Maximum tokens for the response (server default: 4096)
+            response_schema: Optional JSON Schema for structured output. When provided,
+                the response will include a 'structured_output' field with the LLM
+                response parsed according to this schema.
 
         Returns:
-            ReflectResponse with answer text and optionally facts used
+            ReflectResponse with answer text, optionally facts used, and optionally
+            structured_output if response_schema was provided
         """
         request_obj = reflect_request.ReflectRequest(
             query=query,
             budget=budget,
             context=context,
+            max_tokens=max_tokens,
+            response_schema=response_schema,
         )
 
         return _run_async(self._memory_api.reflect(bank_id, request_obj))
@@ -303,24 +324,34 @@ class Hindsight:
 
         Args:
             bank_id: The memory bank ID
-            items: List of memory items with 'content' and optional 'timestamp', 'context', 'metadata', 'document_id'
+            items: List of memory items with 'content' and optional 'timestamp', 'context', 'metadata', 'document_id', 'entities'
             document_id: Optional document ID for grouping memories (applied to items that don't have their own)
             retain_async: If True, process asynchronously in background (default: False)
 
         Returns:
             RetainResponse with success status and item count
         """
-        memory_items = [
-            memory_item.MemoryItem(
-                content=item["content"],
-                timestamp=item.get("timestamp"),
-                context=item.get("context"),
-                metadata=item.get("metadata"),
-                # Use item's document_id if provided, otherwise fall back to batch-level document_id
-                document_id=item.get("document_id") or document_id,
+        from hindsight_client_api.models.entity_input import EntityInput
+
+        memory_items = []
+        for item in items:
+            entities = None
+            if item.get("entities"):
+                entities = [
+                    EntityInput(text=e["text"], type=e.get("type"))
+                    for e in item["entities"]
+                ]
+            memory_items.append(
+                memory_item.MemoryItem(
+                    content=item["content"],
+                    timestamp=item.get("timestamp"),
+                    context=item.get("context"),
+                    metadata=item.get("metadata"),
+                    # Use item's document_id if provided, otherwise fall back to batch-level document_id
+                    document_id=item.get("document_id") or document_id,
+                    entities=entities,
+                )
             )
-            for item in items
-        ]
 
         request_obj = retain_request.RetainRequest(
             items=memory_items,
@@ -337,6 +368,7 @@ class Hindsight:
         context: Optional[str] = None,
         document_id: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
+        entities: Optional[List[Dict[str, str]]] = None,
     ) -> RetainResponse:
         """
         Store a single memory (async).
@@ -348,13 +380,14 @@ class Hindsight:
             context: Optional context description
             document_id: Optional document ID for grouping
             metadata: Optional user-defined metadata
+            entities: Optional list of entities [{"text": "...", "type": "..."}]
 
         Returns:
             RetainResponse with success status
         """
         return await self.aretain_batch(
             bank_id=bank_id,
-            items=[{"content": content, "timestamp": timestamp, "context": context, "metadata": metadata}],
+            items=[{"content": content, "timestamp": timestamp, "context": context, "metadata": metadata, "entities": entities}],
             document_id=document_id,
         )
 

hindsight_client-0.2.1/hindsight_client_api/__init__.py

@@ -0,0 +1,81 @@
+# coding: utf-8
+
+# flake8: noqa
+
+"""
+    Hindsight HTTP API
+
+    HTTP API for Hindsight
+
+    The version of the OpenAPI document: 0.1.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+__version__ = "0.0.7"
+
+# import apis into sdk package
+from hindsight_client_api.api.banks_api import BanksApi
+from hindsight_client_api.api.documents_api import DocumentsApi
+from hindsight_client_api.api.entities_api import EntitiesApi
+from hindsight_client_api.api.memory_api import MemoryApi
+from hindsight_client_api.api.monitoring_api import MonitoringApi
+from hindsight_client_api.api.operations_api import OperationsApi
+
+# import ApiClient
+from hindsight_client_api.api_response import ApiResponse
+from hindsight_client_api.api_client import ApiClient
+from hindsight_client_api.configuration import Configuration
+from hindsight_client_api.exceptions import OpenApiException
+from hindsight_client_api.exceptions import ApiTypeError
+from hindsight_client_api.exceptions import ApiValueError
+from hindsight_client_api.exceptions import ApiKeyError
+from hindsight_client_api.exceptions import ApiAttributeError
+from hindsight_client_api.exceptions import ApiException
+
+# import models into sdk package
+from hindsight_client_api.models.add_background_request import AddBackgroundRequest
+from hindsight_client_api.models.background_response import BackgroundResponse
+from hindsight_client_api.models.bank_list_item import BankListItem
+from hindsight_client_api.models.bank_list_response import BankListResponse
+from hindsight_client_api.models.bank_profile_response import BankProfileResponse
+from hindsight_client_api.models.bank_stats_response import BankStatsResponse
+from hindsight_client_api.models.budget import Budget
+from hindsight_client_api.models.cancel_operation_response import CancelOperationResponse
+from hindsight_client_api.models.chunk_data import ChunkData
+from hindsight_client_api.models.chunk_include_options import ChunkIncludeOptions
+from hindsight_client_api.models.chunk_response import ChunkResponse
+from hindsight_client_api.models.create_bank_request import CreateBankRequest
+from hindsight_client_api.models.delete_document_response import DeleteDocumentResponse
+from hindsight_client_api.models.delete_response import DeleteResponse
+from hindsight_client_api.models.disposition_traits import DispositionTraits
+from hindsight_client_api.models.document_response import DocumentResponse
+from hindsight_client_api.models.entity_detail_response import EntityDetailResponse
+from hindsight_client_api.models.entity_include_options import EntityIncludeOptions
+from hindsight_client_api.models.entity_input import EntityInput
+from hindsight_client_api.models.entity_list_item import EntityListItem
+from hindsight_client_api.models.entity_list_response import EntityListResponse
+from hindsight_client_api.models.entity_observation_response import EntityObservationResponse
+from hindsight_client_api.models.entity_state_response import EntityStateResponse
+from hindsight_client_api.models.graph_data_response import GraphDataResponse
+from hindsight_client_api.models.http_validation_error import HTTPValidationError
+from hindsight_client_api.models.include_options import IncludeOptions
+from hindsight_client_api.models.list_documents_response import ListDocumentsResponse
+from hindsight_client_api.models.list_memory_units_response import ListMemoryUnitsResponse
+from hindsight_client_api.models.memory_item import MemoryItem
+from hindsight_client_api.models.operation_response import OperationResponse
+from hindsight_client_api.models.operations_list_response import OperationsListResponse
+from hindsight_client_api.models.recall_request import RecallRequest
+from hindsight_client_api.models.recall_response import RecallResponse
+from hindsight_client_api.models.recall_result import RecallResult
+from hindsight_client_api.models.reflect_fact import ReflectFact
+from hindsight_client_api.models.reflect_include_options import ReflectIncludeOptions
+from hindsight_client_api.models.reflect_request import ReflectRequest
+from hindsight_client_api.models.reflect_response import ReflectResponse
+from hindsight_client_api.models.retain_request import RetainRequest
+from hindsight_client_api.models.retain_response import RetainResponse
+from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
+from hindsight_client_api.models.validation_error import ValidationError
+from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client_api/api_client.py

@@ -21,7 +21,6 @@ import mimetypes
 import os
 import re
 import tempfile
-import uuid
 
 from urllib.parse import quote
 from typing import Tuple, Optional, List, Dict, Union
@@ -360,8 +359,6 @@ class ApiClient:
             return obj.get_secret_value()
         elif isinstance(obj, self.PRIMITIVE_TYPES):
             return obj
-        elif isinstance(obj, uuid.UUID):
-            return str(obj)
         elif isinstance(obj, list):
             return [
                 self.sanitize_for_serialization(sub_obj) for sub_obj in obj
@@ -388,10 +385,6 @@ class ApiClient:
             else:
                 obj_dict = obj.__dict__
 
-        if isinstance(obj_dict, list):
-            # here we handle instances that can either be a list or something else, and only became a real list by calling to_dict()
-            return self.sanitize_for_serialization(obj_dict)
-
         return {
             key: self.sanitize_for_serialization(val)
             for key, val in obj_dict.items()
@@ -414,7 +407,7 @@ class ApiClient:
                 data = json.loads(response_text)
             except ValueError:
                 data = response_text
-        elif re.match(r'^application/(json|[\w!#$&.+\-^_]+\+json)\s*(;|$)', content_type, re.IGNORECASE):
+        elif re.match(r'^application/(json|[\w!#$&.+-^_]+\+json)\s*(;|$)', content_type, re.IGNORECASE):
             if response_text == "":
                 data = ""
             else:
@@ -463,13 +456,13 @@ class ApiClient:
 
         if klass in self.PRIMITIVE_TYPES:
             return self.__deserialize_primitive(data, klass)
-        elif klass is object:
+        elif klass == object:
             return self.__deserialize_object(data)
-        elif klass is datetime.date:
+        elif klass == datetime.date:
             return self.__deserialize_date(data)
-        elif klass is datetime.datetime:
+        elif klass == datetime.datetime:
             return self.__deserialize_datetime(data)
-        elif klass is decimal.Decimal:
+        elif klass == decimal.Decimal:
             return decimal.Decimal(data)
         elif issubclass(klass, Enum):
             return self.__deserialize_enum(data, klass)
@@ -527,7 +520,7 @@ class ApiClient:
             if k in collection_formats:
                 collection_format = collection_formats[k]
                 if collection_format == 'multi':
-                    new_params.extend((k, quote(str(value))) for value in v)
+                    new_params.extend((k, str(value)) for value in v)
                 else:
                     if collection_format == 'ssv':
                         delimiter = ' '

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client_api/configuration.py

@@ -17,7 +17,7 @@ import http.client as httplib
 import logging
 from logging import FileHandler
 import sys
-from typing import Any, ClassVar, Dict, List, Literal, Optional, TypedDict, Union
+from typing import Any, ClassVar, Dict, List, Literal, Optional, TypedDict
 from typing_extensions import NotRequired, Self
 
 import urllib3
@@ -159,10 +159,6 @@ class Configuration:
     :param ssl_ca_cert: str - the path to a file of concatenated CA certificates
       in PEM format.
     :param retries: Number of retries for API requests.
-    :param ca_cert_data: verify the peer using concatenated CA certificate data
-      in PEM (str) or DER (bytes) format.
-    :param cert_file: the path to a client certificate file, for mTLS.
-    :param key_file: the path to a client key file, for mTLS. 
 
     """
 
@@ -176,16 +172,13 @@ class Configuration:
         username: Optional[str]=None,
         password: Optional[str]=None,
         access_token: Optional[str]=None,
-        server_index: Optional[int]=None,
+        server_index: Optional[int]=None, 
         server_variables: Optional[ServerVariablesT]=None,
         server_operation_index: Optional[Dict[int, int]]=None,
         server_operation_variables: Optional[Dict[int, ServerVariablesT]]=None,
         ignore_operation_servers: bool=False,
         ssl_ca_cert: Optional[str]=None,
         retries: Optional[int] = None,
-        ca_cert_data: Optional[Union[str, bytes]] = None,
-        cert_file: Optional[str]=None,
-        key_file: Optional[str]=None,
         *,
         debug: Optional[bool] = None,
     ) -> None:
@@ -263,14 +256,10 @@ class Configuration:
         self.ssl_ca_cert = ssl_ca_cert
         """Set this to customize the certificate file to verify the peer.
         """
-        self.ca_cert_data = ca_cert_data
-        """Set this to verify the peer using PEM (str) or DER (bytes)
-           certificate data.
-        """
-        self.cert_file = cert_file
+        self.cert_file = None
         """client certificate file
         """
-        self.key_file = key_file
+        self.key_file = None
         """client key file
         """
         self.assert_hostname = None

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client_api/exceptions.py

@@ -150,13 +150,6 @@ class ApiException(OpenApiException):
         if http_resp.status == 404:
             raise NotFoundException(http_resp=http_resp, body=body, data=data)
 
-        # Added new conditions for 409 and 422
-        if http_resp.status == 409:
-            raise ConflictException(http_resp=http_resp, body=body, data=data)
-
-        if http_resp.status == 422:
-            raise UnprocessableEntityException(http_resp=http_resp, body=body, data=data)
-
         if 500 <= http_resp.status <= 599:
             raise ServiceException(http_resp=http_resp, body=body, data=data)
         raise ApiException(http_resp=http_resp, body=body, data=data)
@@ -169,11 +162,8 @@ class ApiException(OpenApiException):
             error_message += "HTTP response headers: {0}\n".format(
                 self.headers)
 
-        if self.body:
-            error_message += "HTTP response body: {0}\n".format(self.body)
-
-        if self.data:
-            error_message += "HTTP response data: {0}\n".format(self.data)
+        if self.data or self.body:
+            error_message += "HTTP response body: {0}\n".format(self.data or self.body)
 
         return error_message
 
@@ -198,16 +188,6 @@ class ServiceException(ApiException):
     pass
 
 
-class ConflictException(ApiException):
-    """Exception for HTTP 409 Conflict."""
-    pass
-
-
-class UnprocessableEntityException(ApiException):
-    """Exception for HTTP 422 Unprocessable Entity."""
-    pass
-
-
 def render_path(path_to_item):
     """Returns a string representation of a path"""
     result = ""

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client_api/models/__init__.py

@@ -12,6 +12,7 @@
     Do not edit the class manually.
 """  # noqa: E501
 
+
 # import models into model package
 from hindsight_client_api.models.add_background_request import AddBackgroundRequest
 from hindsight_client_api.models.background_response import BackgroundResponse
@@ -31,6 +32,7 @@ from hindsight_client_api.models.disposition_traits import DispositionTraits
 from hindsight_client_api.models.document_response import DocumentResponse
 from hindsight_client_api.models.entity_detail_response import EntityDetailResponse
 from hindsight_client_api.models.entity_include_options import EntityIncludeOptions
+from hindsight_client_api.models.entity_input import EntityInput
 from hindsight_client_api.models.entity_list_item import EntityListItem
 from hindsight_client_api.models.entity_list_response import EntityListResponse
 from hindsight_client_api.models.entity_observation_response import EntityObservationResponse
@@ -55,4 +57,3 @@ from hindsight_client_api.models.retain_response import RetainResponse
 from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
 from hindsight_client_api.models.validation_error import ValidationError
 from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner
-

hindsight_client-0.2.1/hindsight_client_api/models/entity_input.py

@@ -0,0 +1,94 @@
+# coding: utf-8
+
+"""
+    Hindsight HTTP API
+
+    HTTP API for Hindsight
+
+    The version of the OpenAPI document: 0.1.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+import pprint
+import re  # noqa: F401
+import json
+
+from pydantic import BaseModel, ConfigDict, Field, StrictStr
+from typing import Any, ClassVar, Dict, List, Optional
+from typing import Optional, Set
+from typing_extensions import Self
+
+class EntityInput(BaseModel):
+    """
+    Entity to associate with retained content.
+    """ # noqa: E501
+    text: StrictStr = Field(description="The entity name/text")
+    type: Optional[StrictStr] = None
+    __properties: ClassVar[List[str]] = ["text", "type"]
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        protected_namespaces=(),
+    )
+
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.model_dump(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Optional[Self]:
+        """Create an instance of EntityInput from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the dictionary representation of the model using alias.
+
+        This has the following differences from calling pydantic's
+        `self.model_dump(by_alias=True)`:
+
+        * `None` is only added to the output dict for nullable fields that
+          were set at model initialization. Other fields with value `None`
+          are ignored.
+        """
+        excluded_fields: Set[str] = set([
+        ])
+
+        _dict = self.model_dump(
+            by_alias=True,
+            exclude=excluded_fields,
+            exclude_none=True,
+        )
+        # set to None if type (nullable) is None
+        # and model_fields_set contains the field
+        if self.type is None and "type" in self.model_fields_set:
+            _dict['type'] = None
+
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
+        """Create an instance of EntityInput from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return cls.model_validate(obj)
+
+        _obj = cls.model_validate({
+            "text": obj.get("text"),
+            "type": obj.get("type")
+        })
+        return _obj
+
+

{hindsight_client-0.1.16 → hindsight_client-0.2.1}/hindsight_client_api/models/memory_item.py

@@ -20,6 +20,7 @@ import json
 from datetime import datetime
 from pydantic import BaseModel, ConfigDict, StrictStr
 from typing import Any, ClassVar, Dict, List, Optional
+from hindsight_client_api.models.entity_input import EntityInput
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -32,7 +33,8 @@ class MemoryItem(BaseModel):
     context: Optional[StrictStr] = None
     metadata: Optional[Dict[str, StrictStr]] = None
     document_id: Optional[StrictStr] = None
-    __properties: ClassVar[List[str]] = ["content", "timestamp", "context", "metadata", "document_id"]
+    entities: Optional[List[EntityInput]] = None
+    __properties: ClassVar[List[str]] = ["content", "timestamp", "context", "metadata", "document_id", "entities"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -73,6 +75,13 @@ class MemoryItem(BaseModel):
             exclude=excluded_fields,
             exclude_none=True,
         )
+        # override the default output from pydantic by calling `to_dict()` of each item in entities (list)
+        _items = []
+        if self.entities:
+            for _item_entities in self.entities:
+                if _item_entities:
+                    _items.append(_item_entities.to_dict())
+            _dict['entities'] = _items
         # set to None if timestamp (nullable) is None
         # and model_fields_set contains the field
         if self.timestamp is None and "timestamp" in self.model_fields_set:
@@ -93,6 +102,11 @@ class MemoryItem(BaseModel):
         if self.document_id is None and "document_id" in self.model_fields_set:
             _dict['document_id'] = None
 
+        # set to None if entities (nullable) is None
+        # and model_fields_set contains the field
+        if self.entities is None and "entities" in self.model_fields_set:
+            _dict['entities'] = None
+
         return _dict
 
     @classmethod
@@ -109,7 +123,8 @@ class MemoryItem(BaseModel):
             "timestamp": obj.get("timestamp"),
             "context": obj.get("context"),
             "metadata": obj.get("metadata"),
-            "document_id": obj.get("document_id")
+            "document_id": obj.get("document_id"),
+            "entities": [EntityInput.from_dict(_item) for _item in obj["entities"]] if obj.get("entities") is not None else None
         })
         return _obj