cheshirecat-python-sdk 1.2.1__py3-none-any.whl → 1.8.4__py3-none-any.whl

cheshirecat_python_sdk/endpoints/rabbit_hole.py

@@ -19,8 +19,9 @@ class RabbitHoleEndpoint(AbstractEndpoint):
     def post_file(
         self,
         file_path,
+        agent_id: str,
+        chat_id: str | None = None,
         file_name: str | None = None,
-        agent_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> UploadSingleFileResponse:
         """
@@ -29,8 +30,9 @@ class RabbitHoleEndpoint(AbstractEndpoint):
         The process is asynchronous and the results are returned in a batch.
         The CheshireCat processes the injection in background and the client will be informed at the end of the process.
         :param file_path: The path to the file to upload.
-        :param file_name: The name of the file.
         :param agent_id: The ID of the agent.
+        :param chat_id: The ID of the chat (optional).
+        :param file_name: The name of the file.
         :param metadata: The metadata to include with the file.
         :return: The response from the RabbitHole API.
         """
@@ -40,24 +42,28 @@ class RabbitHoleEndpoint(AbstractEndpoint):
         if metadata is not None:
             payload.data["metadata"] = json.dumps(metadata)
 
+        endpoint = self.prefix if not chat_id else self.format_url(chat_id)
+
         with open(file_path, "rb") as file:
             payload.files = [("file", file_attributes(file_name, file))]
-            result = self.post_multipart(self.prefix, UploadSingleFileResponse, payload, agent_id)
+            result = self.post_multipart(endpoint, agent_id, output_class=UploadSingleFileResponse, payload=payload)
 
         return result
 
     def post_files(
         self,
         file_paths: List[str],
-        agent_id: str | None = None,
+        agent_id: str,
+        chat_id: str | None = None,
         metadata: Dict[str, Any] | None = None
-    ) -> Dict[str, UploadSingleFileResponse]:
+    ) -> Dict[str, UploadSingleFileResponse]:  # type: ignore
         """
         Posts multiple files to the RabbitHole API. The files are uploaded to the RabbitHole server and
         ingested into the RAG system. The files are processed in a batch. The process is asynchronous.
         The CheshireCat processes the injection in background and the client will be informed at the end of the process.
         :param file_paths: The paths to the files to upload.
         :param agent_id: The ID of the agent.
+        :param chat_id: The ID of the chat (optional).
         :param metadata: The metadata to include with the files.
         :return: The response from the RabbitHole API.
         """
@@ -67,13 +73,16 @@ class RabbitHoleEndpoint(AbstractEndpoint):
 
         files = []
         file_handles = []
+
+        endpoint = self.format_url("/batch") if not chat_id else self.format_url(f"/batch/{chat_id}")
         try:
             for file_path in file_paths:
                 file = open(file_path, "rb")
                 file_handles.append(file)
                 files.append(("files", file_attributes(Path(file_path).name, file)))
 
-            response = self.get_http_client(agent_id).post(self.format_url("/batch"), data=data, files=files)
+            response = self.get_http_client(agent_id).post(endpoint, data=data, files=files)
+            response.raise_for_status()
 
             result = {}
             for key, item in response.json().items():
@@ -86,39 +95,43 @@ class RabbitHoleEndpoint(AbstractEndpoint):
     def post_web(
         self,
         web_url: str,
-        agent_id: str | None = None,
+        agent_id: str,
+        chat_id: str | None = None,
         metadata: Dict[str, Any] | None = None
     ) -> UploadUrlResponse:
         """
         Posts a web URL to the RabbitHole API. The web URL is ingested into the RAG system. The web URL is
-        processed by the RAG system by Web scraping and the results are stored in the RAG database.
+        processed by the RAG system by Web scraping, and the results are stored in the RAG database.
         The process is asynchronous.
-        The CheshireCat processes the injection in background and the client will be informed at the end of the process.
+        The CheshireCat processes the injection in background, and the client will be informed at the end of the process.
         :param web_url: The URL of the website to ingest.
         :param agent_id: The ID of the agent.
+        :param chat_id: The ID of the chat (optional).
         :param metadata: The metadata to include with the files.
         :return: The response from the RabbitHole API.
         """
         payload = {"url": web_url}
         if metadata is not None:
-            payload["metadata"] = metadata
+            payload["metadata"] = metadata  # type: ignore
+
+        endpoint = self.format_url("/web") if not chat_id else self.format_url(f"/web/{chat_id}")
 
-        return self.post_json(f"{self.prefix}/web", UploadUrlResponse, payload, agent_id)
+        return self.post_json(endpoint, agent_id, output_class=UploadUrlResponse, payload=payload)
 
     def post_memory(
         self,
         file_path: str,
+        agent_id: str,
         file_name: str | None = None,
-        agent_id: str | None = None
     ) -> UploadSingleFileResponse:
         """
-        Posts a memory point, either for the agent identified by the agent_id parameter (for multi-agent
-        installations) or for the default agent (for single-agent installations). The memory point is ingested into the
-        RAG system. The process is asynchronous. The provided file must be in JSON format.
-        The CheshireCat processes the injection in background and the client will be informed at the end of the process.
+        Posts a memory point, for the agent identified by the agent_id parameter.
+        The memory point is ingested into the RAG system. The process is asynchronous. The provided file must be in JSON
+        format. The CheshireCat processes the injection in the background, and the client will be informed at the end of
+        the process.
         :param file_path: The path to the file to upload.
-        :param file_name: The name of the file.
         :param agent_id: The ID of the agent.
+        :param file_name: The name of the file (optional).
         :return: The response from the RabbitHole API.
         """
         file_name = file_name or Path(file_path).name
@@ -126,18 +139,35 @@ class RabbitHoleEndpoint(AbstractEndpoint):
         payload = MultipartPayload()
         with open(file_path, "rb") as file:
             payload.files = [("file", file_attributes(file_name, file))]
-            result = self.post_multipart(f"{self.prefix}/memory", UploadSingleFileResponse, payload, agent_id)
+            result = self.post_multipart(
+                self.format_url("/memory"), agent_id, output_class=UploadSingleFileResponse, payload=payload
+            )
 
         return result
 
-    def get_allowed_mime_types(self, agent_id: str | None = None) -> AllowedMimeTypesOutput:
+    def get_allowed_mime_types(self, agent_id: str) -> AllowedMimeTypesOutput:
         """
         Retrieves the allowed MIME types for the RabbitHole API. The allowed MIME types are the MIME types
         that are allowed to be uploaded to the RabbitHole API. The allowed MIME types are returned in a list.
-        If the agent_id parameter is provided, the allowed MIME types are retrieved for the agent identified by the
-        agent_id parameter (for multi-agent installations). If the agent_id parameter is not provided, the allowed MIME
-        types are retrieved for the default agent (for single-agent installations).
         :param agent_id: The ID of the agent.
         :return: AllowedMimeTypesOutput, the details of the allowed MIME types.
         """
-        return self.get(f"{self.prefix}/allowed-mimetypes", AllowedMimeTypesOutput, agent_id)
+        return self.get(
+            self.format_url("/allowed-mimetypes"),
+            agent_id,
+            output_class=AllowedMimeTypesOutput
+        )
+
+    def get_web_sources(self, agent_id: str, chat_id: str | None = None) -> List[str]:
+        """
+        This method retrieves the web sources for the RabbitHole API. The web sources are the web URLs that are allowed
+        to be uploaded to the RabbitHole API. The web sources are returned in a list.
+        :param agent_id: The ID of the agent.
+        :param chat_id: The chat id, optional
+        :return: List[str]
+        """
+        return self.get(
+            self.format_url("/web"),
+            agent_id,
+            chat_id=chat_id,
+        )

cheshirecat_python_sdk/endpoints/users.py

@@ -1,7 +1,6 @@
 from typing import Any, List, Dict
 
 from cheshirecat_python_sdk.endpoints.base import AbstractEndpoint
-from cheshirecat_python_sdk.models.api.tokens import TokenOutput
 from cheshirecat_python_sdk.models.api.users import UserOutput
 from cheshirecat_python_sdk.utils import deserialize
 
@@ -11,48 +10,24 @@ class UsersEndpoint(AbstractEndpoint):
         super().__init__(client)
         self.prefix = "/users"
 
-    def token(self, username: str, password: str) -> TokenOutput:
-        """
-        This endpoint is used to get a token for the user. The token is used to authenticate the user in the system. When
-        the token expires, the user must request a new token.
-        """
-        http_client = self.client.http_client.get_client()
-
-        response = http_client.post("/auth/token", json={
-            "username": username,
-            "password": password,
-        })
-
-        result = deserialize(response.json(), TokenOutput)
-        self.client.add_token(result.access_token)
-
-        return result
-
-    def get_available_permissions(self, agent_id: str | None = None) -> dict[int | str, Any]:
-        """
-        This endpoint is used to get a list of available permissions in the system. The permissions are used to define
-        the access rights of the users in the system. The permissions are defined by the system administrator.
-        The endpoint can be used either for the agent identified by the agentId parameter (for multi-agent installations)
-        or for the default agent (for single-agent installations).
-        :param agent_id: The id of the agent to get settings for (optional)
-        :return array<int|string, Any>, the available permissions
-        """
-        return self.get("/auth/available-permissions", agent_id=agent_id)
-
     def post_user(
-        self, username: str, password: str, permissions: dict[str, Any] | None = None, agent_id: str | None = None
+        self,
+        agent_id: str,
+        username: str,
+        password: str,
+        permissions: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> UserOutput:
         """
         This endpoint is used to create a new user in the system. The user is created with the specified username and
         password. The user is assigned the specified permissions. The permissions are used to define the access rights
         of the user in the system and are defined by the system administrator.
-        The endpoint can be used either for the
-        agent identified by the agentId parameter (for multi-agent installations) or for the default agent (for
-        single-agent installations).
+        The endpoint can be used for the agent identified by the agentId parameter.
+        :param agent_id: The id of the agent to create the user for
         :param username: The username of the user to create
         :param password: The password of the user to create
         :param permissions: The permissions of the user to create (optional)
-        :param agent_id: The id of the agent to create the user for (optional)
+        :param metadata: The metadata of the user to create (optional)
         :return UserOutput, the created user
         """
         payload = {
@@ -60,66 +35,69 @@ class UsersEndpoint(AbstractEndpoint):
             "password": password,
         }
         if permissions is not None:
-            payload["permissions"] = permissions
+            payload["permissions"] = permissions  # type: ignore
+
+        if metadata is not None:
+            payload["metadata"] = metadata  # type: ignore
 
         return self.post_json(
             self.prefix,
-            UserOutput,
-            payload,
             agent_id,
+            output_class=UserOutput,
+            payload=payload,
         )
 
-    def get_users(self, agent_id: str | None = None) -> List[UserOutput]:
+    def get_users(self, agent_id: str) -> List[UserOutput]:
         """
         This endpoint is used to get a list of users in the system. The list includes the username and the permissions of
         each user. The permissions are used to define the access rights of the users in the system and are defined by the
         system administrator.
-        The endpoint can be used either for the agent identified by the agentId parameter (for multi-agent installations)
-        or for the default agent (for single-agent installations).
-        :param agent_id: The id of the agent to get users for (optional)
+        The endpoint can be used for the agent identified by the agentId parameter.
+        :param agent_id: The id of the agent to get users for
         :return List[UserOutput], the users in the system with their permissions for the agent identified by agent_id
         """
         response = self.get_http_client(agent_id).get(self.prefix)
+        response.raise_for_status()
 
         result = []
         for item in response.json():
             result.append(deserialize(item, UserOutput))
         return result
 
-    def get_user(self, user_id: str, agent_id: str | None = None) -> UserOutput:
+    def get_user(self, user_id: str, agent_id: str) -> UserOutput:
         """
         This endpoint is used to get a user in the system. The user is identified by the userId parameter, previously
         provided by the CheshireCat API when the user was created. The endpoint returns the username and the permissions
         of the user. The permissions are used to define the access rights of the user in the system and are defined by
         the system administrator.
-        The endpoint can be used either for the agent identified by the agentId parameter (for multi-agent installations)
-        or for the default agent (for single-agent installations).
+        The endpoint can be used for the agent identified by the agentId parameter.
         :param user_id: The id of the user to get
-        :param agent_id: The id of the agent to get the user for (optional)
+        :param agent_id: The id of the agent to get the user for
         :return UserOutput, the user
         """
-        return self.get(self.format_url(user_id), UserOutput, agent_id)
+        return self.get(self.format_url(user_id), agent_id, output_class=UserOutput)
 
     def put_user(
         self,
         user_id: str,
+        agent_id: str,
         username: str | None = None,
         password: str | None = None,
         permissions: Dict[str, Any] | None = None,
-        agent_id: str | None = None
+        metadata: Dict[str, Any] | None = None,
     ) -> UserOutput:
         """
         The endpoint is used to update the user in the system. The user is identified by the userId parameter, previously
         provided by the CheshireCat API when the user was created. The endpoint updates the username, the password, and
         the permissions of the user. The permissions are used to define the access rights of the user in the system and
         are defined by the system administrator.
-        The endpoint can be used either for the agent identified by the agentId parameter (for multi-agent installations)
-        or for the default agent (for single-agent installations).
+        The endpoint can be used for the agent identified by the agentId parameter.
         :param user_id: The id of the user to update
+        :param agent_id: The id of the agent to update the user for
         :param username: The new username of the user (optional)
         :param password: The new password of the user (optional)
         :param permissions: The new permissions of the user (optional)
-        :param agent_id: The id of the agent to update the user for (optional)
+        :param metadata: The new metadata of the user (optional)
         :return UserOutput, the updated user
         """
         payload = {}
@@ -129,17 +107,18 @@ class UsersEndpoint(AbstractEndpoint):
             payload["password"] = password
         if permissions is not None:
             payload["permissions"] = permissions
+        if metadata is not None:
+            payload["metadata"] = metadata
 
-        return self.put(self.format_url(user_id), UserOutput, payload, agent_id)
+        return self.put(self.format_url(user_id), agent_id, output_class=UserOutput, payload=payload)
 
-    def delete_user(self, user_id: str, agent_id: str | None = None) -> UserOutput:
+    def delete_user(self, user_id: str, agent_id: str) -> UserOutput:
         """
         This endpoint is used to delete the user in the system. The user is identified by the userId parameter, previously
         provided by the CheshireCat API when the user was created.
-        The endpoint can be used either for the agent identified by the agentId parameter (for multi-agent installations)
-        or for the default agent (for single-agent installations).
+        The endpoint can be used for the agent identified by the agentId parameter.
         :param user_id: The id of the user to delete
-        :param agent_id: The id of the agent to delete the user for (optional)
+        :param agent_id: The id of the agent to delete the user for
         :return UserOutput, the deleted user
         """
-        return self.delete(self.format_url(user_id), UserOutput, agent_id)
+        return self.delete(self.format_url(user_id), agent_id, output_class=UserOutput)

cheshirecat_python_sdk/endpoints/utils.py

@@ -0,0 +1,71 @@
+from typing import List
+
+from cheshirecat_python_sdk.endpoints.base import AbstractEndpoint
+from cheshirecat_python_sdk.models.api.admins import ResetOutput, ClonedOutput, CreatedOutput
+from cheshirecat_python_sdk.utils import deserialize
+
+
+class UtilsEndpoint(AbstractEndpoint):
+    def __init__(self, client: "CheshireCatClient"):
+        super().__init__(client)
+        self.prefix = "/utils"
+
+    def post_factory_reset(self) -> ResetOutput:
+        """
+        Reset the system to the factory settings.
+        :return: ResetOutput, the details of the reset.
+        """
+        return self.post_json(self.format_url("/factory/reset/"), self.system_id, output_class=ResetOutput)
+
+    def get_agents(self) -> List[str]:
+        """
+        Get a list of all agents.
+        :return: List[str], the IDs of the agents.
+        """
+        return self.get(self.format_url("/agents/"), self.system_id)
+
+    def post_agent_create(self, agent_id: str) -> CreatedOutput:
+        """
+        Create a new agent.
+        :param agent_id: The ID of the agent.
+        :return: CreatedOutput, the details of the agent.
+        """
+        response = self.get_http_client().post(
+            self.format_url("/agents/create/"),
+            json={
+                "agent_id": agent_id,
+            },
+        )
+        response.raise_for_status()
+
+        return deserialize(response.json(), CreatedOutput)
+
+    def post_agent_reset(self, agent_id: str) -> ResetOutput:
+        """
+        Reset an agent to the factory settings.
+        :param agent_id: The ID of the agent.
+        :return: ResetOutput, the details of the reset.
+        """
+        return self.post_json(self.format_url("/agents/reset/"), agent_id, output_class=ResetOutput)
+
+    def post_agent_destroy(self, agent_id: str) -> ResetOutput:
+        """
+        Destroy an agent.
+        :param agent_id: The ID of the agent.
+        :return: ResetOutput, the details of the reset.
+        """
+        return self.post_json(self.format_url("/agents/destroy/"), agent_id, output_class=ResetOutput)
+
+    def post_agent_clone(self, agent_id: str, new_agent_id: str) -> ClonedOutput:
+        """
+        Destroy an agent.
+        :param agent_id: The ID of the agent.
+        :param new_agent_id: The ID of the new cloned agent.
+        :return: ClonedOutput, the details of the cloning.
+        """
+        return self.post_json(
+            self.format_url("/agents/clone/"),
+            agent_id,
+            payload={"agent_id": new_agent_id},
+            output_class=ClonedOutput
+        )

cheshirecat_python_sdk/endpoints/vector_database.py

@@ -0,0 +1,52 @@
+from typing import Dict, Any
+
+from cheshirecat_python_sdk.endpoints.base import AbstractEndpoint
+from cheshirecat_python_sdk.models.api.factories import FactoryObjectSettingsOutput, FactoryObjectSettingOutput
+
+
+class VectorDatabaseEndpoint(AbstractEndpoint):
+    def __init__(self, client: "CheshireCatClient"):
+        super().__init__(client)
+        self.prefix = "/vector_database"
+
+    def get_vector_databases_settings(self, agent_id: str) -> FactoryObjectSettingsOutput:
+        """
+        Get all vector databases settings for the agent specified by agent_id
+        :param agent_id: The agent id
+        :return: FactoryObjectSettingsOutput, a list of vector database settings
+        """
+        return self.get(
+            self.format_url("/settings"),
+            agent_id,
+            output_class=FactoryObjectSettingsOutput,
+        )
+
+    def get_vector_database_settings(self, vector_database: str, agent_id: str) -> FactoryObjectSettingOutput:
+        """
+        Get the vector database settings for the vector database specified by vector_database and agent_id
+        :param vector_database: The name of the vector database
+        :param agent_id: The agent id
+        :return: FactoryObjectSettingOutput, the vector database settings
+        """
+        return self.get(
+            self.format_url(f"/settings/{vector_database}"),
+            agent_id,
+            output_class=FactoryObjectSettingOutput,
+        )
+
+    def put_vector_database_settings(
+        self, vector_database: str, agent_id: str, values: Dict[str, Any]
+    ) -> FactoryObjectSettingOutput:
+        """
+        Update the vector database settings for the vector database specified by vector_database and agent_id
+        :param vector_database: The name of the vector database
+        :param agent_id: The agent id
+        :param values: The new settings
+        :return: FactoryObjectSettingOutput, the updated vector database settings
+        """
+        return self.put(
+            self.format_url(f"/settings/{vector_database}"),
+            agent_id,
+            output_class=FactoryObjectSettingOutput,
+            payload=values,
+        )

cheshirecat_python_sdk/enums.py

@@ -28,14 +28,3 @@ class Enum(BaseEnum, metaclass=MetaEnum):
 
     def __hash__(self):
         return hash(self.value)
-
-
-class Collection(str, Enum):
-    DECLARATIVE = "declarative"
-    PROCEDURAL = "procedural"
-    EPISODIC = "episodic"
-
-
-class Role(str, Enum):
-    AI = "AI"
-    HUMAN = "Human"

cheshirecat_python_sdk/models/api/admins.py

@@ -1,15 +1,9 @@
-from typing import Dict, List, Any
+from typing import Dict, Any
 from pydantic import BaseModel
 
 from cheshirecat_python_sdk.models.api.plugins import PluginToggleOutput
 
 
-class AdminOutput(BaseModel):
-    username: str
-    permissions: Dict[str, List[str]]
-    id: str
-
-
 class CreatedOutput(BaseModel):
     created: bool
 
@@ -36,3 +30,7 @@ class ResetOutput(BaseModel):
     deleted_settings: bool
     deleted_memories: bool
     deleted_plugin_folders: bool
+
+
+class ClonedOutput(BaseModel):
+    cloned: bool = False

cheshirecat_python_sdk/models/api/conversations.py

@@ -0,0 +1,24 @@
+from typing import List
+from pydantic import BaseModel
+
+from cheshirecat_python_sdk.models.api.nested.memories import ConversationMessage
+
+
+class ConversationDeleteOutput(BaseModel):
+    deleted: bool
+
+
+class ConversationHistoryOutput(BaseModel):
+    history: List[ConversationMessage]
+
+
+class ConversationsResponse(BaseModel):
+    chat_id: str
+    name: str
+    num_messages: int
+    created_at: float | None
+    updated_at: float | None
+
+
+class ConversationNameChangeOutput(BaseModel):
+    changed: bool

cheshirecat_python_sdk/models/api/factories.py

@@ -7,6 +7,12 @@ class FactoryObjectSettingOutput(BaseModel):
     value: Dict[str, Any]
     scheme: Dict[str, Any] | None = None
 
+    def __init__(self, /, **data: Any) -> None:
+        # if tags is a list, convert it to a comma-separated string
+        if "scheme" in data and isinstance(data["scheme"], Dict) and not data["scheme"]:
+            data["scheme"] = None
+        super().__init__(**data)
+
 
 class FactoryObjectSettingsOutput(BaseModel):
     settings: List[FactoryObjectSettingOutput]

cheshirecat_python_sdk/models/api/file_managers.py

@@ -0,0 +1,18 @@
+from typing import List
+from pydantic import BaseModel
+
+
+class FileResponse(BaseModel):
+    path: str
+    name: str
+    size: int
+    last_modified: str
+
+
+class FileManagerAttributes(BaseModel):
+    files: List[FileResponse]
+    size: int
+
+
+class FileManagerDeletedFiles(BaseModel):
+    deleted: bool

cheshirecat_python_sdk/models/api/memories.py

@@ -1,9 +1,8 @@
-from typing import Dict, List
+from typing import Dict, List, Any
 from pydantic import BaseModel
 
 from cheshirecat_python_sdk.models.api.nested.memories import (
     CollectionsItem,
-    ConversationHistoryItem,
     MemoryPointsDeleteByMetadataInfo,
     Record,
     MemoryRecallQuery,
@@ -19,13 +18,6 @@ class CollectionPointsDestroyOutput(BaseModel):
 class CollectionsOutput(BaseModel):
     collections: List[CollectionsItem]
 
-class ConversationHistoryDeleteOutput(BaseModel):
-    deleted: bool
-
-
-class ConversationHistoryOutput(BaseModel):
-    history: List[ConversationHistoryItem]
-
 
 class MemoryPointDeleteOutput(BaseModel):
     deleted: str
@@ -33,7 +25,7 @@ class MemoryPointDeleteOutput(BaseModel):
 
 class MemoryPointOutput(MemoryPoint):
     id: str
-    vector: List[float]
+    vector: List[float] | List[List[float]] | Dict[str, Any]
 
 
 class MemoryPointsDeleteByMetadataOutput(BaseModel):

cheshirecat_python_sdk/models/api/messages.py

@@ -1,14 +1,16 @@
-from pydantic import Field
+from pydantic import Field, BaseModel
 
 from cheshirecat_python_sdk.models.dtos import MessageBase, Why
 
 
 class MessageOutput(MessageBase):
-    why: Why = Field(default_factory=Why)  # Assuming Why has a no-args constructor
+    why: Why | None = Field(default_factory=Why)  # Assuming Why has a no-args constructor
     type: str | None = "chat"  # Default argument
     error: bool | None = False  # Default argument
-    content: str = Field(init=False)  # Field without a default value
 
-    def __init__(self, **data):
-        super().__init__(**data)
-        self.content = self.text
+
+class ChatOutput(BaseModel):
+    agent_id: str
+    user_id: str
+    chat_id: str
+    message: MessageOutput

cheshirecat_python_sdk/models/api/nested/memories.py

@@ -13,7 +13,7 @@ class ConversationHistoryItemContent(MessageBase):
     why: Why | None = None
 
 
-class ConversationHistoryItem(BaseModel):
+class ConversationMessage(BaseModel):
     who: str
     when: float
     content: ConversationHistoryItemContent
@@ -26,7 +26,7 @@ class MemoryPointsDeleteByMetadataInfo(BaseModel):
 
 class MemoryRecallQuery(BaseModel):
     text: str
-    vector: List[float]
+    vector: List[float] | List[List[float]] | Dict[str, Any]
 
 
 class MemoryRecallVectors(BaseModel):
@@ -37,6 +37,6 @@ class MemoryRecallVectors(BaseModel):
 class Record(BaseModel):
     id: str
     payload: Dict[str, Any] | None = None
-    vector: List[float] | None = None
-    shard_key: str | None = None
-    order_value: float | None = None
+    vector: List[float] | List[List[float]] | Dict[str, Any] | None = None
+    shard_key: int | str | None = None
+    order_value: int | float | None = None