feat: add assistant memory

1. add memory to assistant schema and related APIs
2. add description and operationId to action schemas

Author: jameszyao

examples/assistant/chat_with_assistant.ipynb

@@ -61,7 +61,8 @@
     "        \"paths\": {\n",
     "            \"/{number}\": {\n",
     "            \"get\": {\n",
-    "                \"summary\": \"Get fact about a number\",\n",
+    "                \"description\": \"Get a fact about a number\",\n",
+    "                \"operationId\": \"getNumberFact\",\n",
     "                \"parameters\": [\n",
     "                {\n",
     "                    \"name\": \"number\",\n",
@@ -109,6 +110,7 @@
    "outputs": [],
    "source": [
     "from taskingai.assistant import Assistant, Chat, AssistantTool, AssistantToolType\n",
+    "from taskingai.assistant.memory import AssistantMessageWindowMemory\n",
     "\n",
     "# choose an available chat_completion model from your project\n",
     "model_id = \"YOUR_MODEL_ID\"\n",
@@ -117,6 +119,10 @@
     "    model_id=model_id,\n",
     "    name=\"My Assistant\",\n",
     "    description=\"A assistant who knows the meaning of various numbers.\",\n",
+    "    memory=AssistantMessageWindowMemory(\n",
+    "        max_messages=20,\n",
+    "        max_tokens=1000\n",
+    "    ),\n",
     "    system_prompt_template=[\n",
     "        \"You know the meaning of various numbers.\",\n",
     "        \"No matter what the user's language is, you will use the {{langugae}} to explain.\"\n",
@@ -255,8 +261,7 @@
     "    chat_id=chat.chat_id,\n",
     ")\n",
     "for message in messages:\n",
-    "    print(f\"{message.role}: {message.content.text}\")\n",
-    "    print(\"-\"*100)"
+    "    print(f\"{message.role}: {message.content.text}\")"
    ],
    "metadata": {
     "collapsed": false
@@ -265,7 +270,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "outputs": [],
    "source": [
     "# delete assistant\n",
@@ -277,16 +282,6 @@
     "collapsed": false
    },
    "id": "ed39836bbfdc7a4e"
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [],
-   "metadata": {
-    "collapsed": false
-   },
-   "id": "325965f154e14f42"
   }
  ],
  "metadata": {

examples/crud/assistant_crud.ipynb

@@ -26,6 +26,7 @@
    "outputs": [],
    "source": [
     "from taskingai.assistant import Assistant, Chat\n",
+    "from taskingai.assistant.memory import AssistantNaiveMemory\n",
     "\n",
     "# choose an available chat_completion model from your project\n",
     "model_id = \"YOUR_MODEL_ID\""
@@ -53,6 +54,7 @@
     "        name=\"My Assistant\",\n",
     "        description=\"This is my assistant\",\n",
     "        system_prompt_template=[\"You are a professional assistant speaking {{language}}.\"],\n",
+    "        memory=AssistantNaiveMemory(),\n",
     "        tools=[],\n",
     "        retrievals=[],\n",
     "        metadata={\"foo\": \"bar\"},\n",
@@ -222,6 +224,18 @@
    "metadata": {
     "collapsed": false
    }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# delete assistant\n",
+    "taskingai.assistant.delete_assistant(assistant_id=assistant.assistant_id)"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
   }
  ],
  "metadata": {

examples/crud/retrieval_crud.ipynb

@@ -161,9 +161,10 @@
    "source": [
     "# create a new collection\n",
     "collection: Collection = create_collection()\n",
+    "print(collection)\n",
     "\n",
     "# wait for the collection creation to finish\n",
-    "time.sleep(10)"
+    "time.sleep(3)"
    ],
    "metadata": {
     "collapsed": false

examples/crud/tool_crud.ipynb

@@ -69,7 +69,8 @@
     "        \"paths\": {\n",
     "            \"/{number}\": {\n",
     "            \"get\": {\n",
-    "                \"summary\": \"Get fact about a number\",\n",
+    "                \"description\": \"Get fact about a number\",\n",
+    "                \"operationId\": \"getNumberFact\",\n",
     "                \"parameters\": [\n",
     "                {\n",
     "                    \"name\": \"number\",\n",
@@ -197,14 +198,6 @@
    },
    "id": "5a1a36d15055918f"
   },
-  {
-   "cell_type": "markdown",
-   "source": [],
-   "metadata": {
-    "collapsed": false
-   },
-   "id": "2f288bf5d1988887"
-  },
   {
    "cell_type": "markdown",
    "source": [],

taskingai/_version.py

@@ -1,2 +1,2 @@
 __title__ = "taskingai"
-__version__ = "0.0.2"
+__version__ = "0.0.5"

taskingai/assistant/__init__.py

@@ -1,3 +1,3 @@
 from .assistant import *
 from .chat import *
-from .message import *
+from .message import *

taskingai/assistant/assistant.py

@@ -1,10 +1,23 @@
 from typing import Optional, List, Dict
 
 from taskingai.client.utils import get_api_instance, ModuleType
-from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool, AssistantToolType, AssistantRetrievalType
-from taskingai.client.models import AssistantCreateRequest, AssistantCreateResponse,\
-    AssistantUpdateRequest, AssistantUpdateResponse,\
-    AssistantGetResponse, AssistantListResponse
+from taskingai.client.models import (
+    Assistant,
+    AssistantMemory,
+    AssistantRetrieval,
+    AssistantTool,
+    AssistantToolType,
+    AssistantRetrievalType
+)
+
+from taskingai.client.models import (
+    AssistantCreateRequest,
+    AssistantCreateResponse,
+    AssistantUpdateRequest,
+    AssistantUpdateResponse,
+    AssistantGetResponse,
+    AssistantListResponse,
+)
 
 __all__ = [
     "Assistant",
@@ -121,6 +134,7 @@ async def a_get_assistant(assistant_id: str) -> Assistant:
 
 def create_assistant(
     model_id: str,
+    memory: AssistantMemory,
     name: Optional[str] = None,
     description: Optional[str] = None,
     system_prompt_template: Optional[List[str]] = None,
@@ -130,8 +144,9 @@ def create_assistant(
 ) -> Assistant:
     """
     Create an assistant.
-    
+
     :param model_id: The ID of an available chat completion model in your project.
+    :param memory: The assistant memory.
     :param name: The assistant name.
     :param description: The assistant description.
     :param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
@@ -142,10 +157,12 @@ def create_assistant(
     """
 
     api_instance = get_api_instance(ModuleType.assistant)
+    memory_dict = memory.model_dump()
     body = AssistantCreateRequest(
         model_id=model_id,
         name=name,
         description=description,
+        memory=memory_dict,
         system_prompt_template=system_prompt_template,
         tools=tools,
         retrievals=retrievals,
@@ -158,6 +175,7 @@ def create_assistant(
 
 async def a_create_assistant(
     model_id: str,
+    memory: AssistantMemory,
     name: Optional[str] = None,
     description: Optional[str] = None,
     system_prompt_template: Optional[List[str]] = None,
@@ -169,6 +187,7 @@ async def a_create_assistant(
     Create an assistant in async mode.
 
     :param model_id: The ID of an available chat completion model in your project.
+    :param memory: The assistant memory.
     :param name: The assistant name.
     :param description: The assistant description.
     :param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
@@ -179,10 +198,12 @@ async def a_create_assistant(
     """
 
     api_instance = get_api_instance(ModuleType.assistant, async_client=True)
+    memory_dict = memory.model_dump()
     body = AssistantCreateRequest(
         model_id=model_id,
         name=name,
         description=description,
+        memory=memory_dict,
         system_prompt_template=system_prompt_template,
         tools=tools,
         retrievals=retrievals,
@@ -199,6 +220,7 @@ def update_assistant(
     name: Optional[str] = None,
     description: Optional[str] = None,
     system_prompt_template: Optional[List[str]] = None,
+    memory: Optional[AssistantMemory] = None,
     tools: Optional[List[AssistantTool]] = None,
     retrievals: Optional[List[AssistantRetrieval]] = None,
     metadata: Optional[Dict[str, str]] = None,
@@ -211,6 +233,7 @@ def update_assistant(
     :param name: The assistant name.
     :param description: The assistant description.
     :param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
+    :param memory: The assistant memory.
     :param tools: The assistant tools.
     :param retrievals: The assistant retrievals.
     :param metadata: The assistant metadata. It can store up to 16 key-value pairs where each key's length is less than 64 and value's length is less than 512.
@@ -223,6 +246,7 @@ def update_assistant(
         name=name,
         description=description,
         system_prompt_template=system_prompt_template,
+        memory=memory,
         tools=tools,
         retrievals=retrievals,
         metadata=metadata,
@@ -238,6 +262,7 @@ async def a_update_assistant(
     name: Optional[str] = None,
     description: Optional[str] = None,
     system_prompt_template: Optional[List[str]] = None,
+    memory: Optional[AssistantMemory] = None,
     tools: Optional[List[AssistantTool]] = None,
     retrievals: Optional[List[AssistantRetrieval]] = None,
     metadata: Optional[Dict[str, str]] = None,
@@ -250,6 +275,7 @@ async def a_update_assistant(
     :param name: The assistant name.
     :param description: The assistant description.
     :param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
+    :param memory: The assistant memory.
     :param tools: The assistant tools.
     :param retrievals: The assistant retrievals.
     :param metadata: The assistant metadata. It can store up to 16 key-value pairs where each key's length is less than 64 and value's length is less than 512.
@@ -262,6 +288,7 @@ async def a_update_assistant(
         name=name,
         description=description,
         system_prompt_template=system_prompt_template,
+        memory=memory,
         tools=tools,
         retrievals=retrievals,
         metadata=metadata,

taskingai/assistant/memory.py

@@ -0,0 +1,15 @@
+from taskingai.client.models import (
+    AssistantMemoryType,
+    AssistantMemory,
+    AssistantNaiveMemory,
+    AssistantZeroMemory,
+    AssistantMessageWindowMemory,
+)
+
+__all__ = [
+    "AssistantMemory",
+    "AssistantMemoryType",
+    "AssistantNaiveMemory",
+    "AssistantZeroMemory",
+    "AssistantMessageWindowMemory",
+]

taskingai/client/models/entity/assistant/assistant.py

@@ -1,13 +1,20 @@
 from typing import List, Dict
+from pydantic import field_validator
 from .._base import TaskingaiBaseModel
+from .assistant_memory import *
 from enum import Enum
 
 __all__ = [
     "Assistant",
+    "AssistantMemoryType",
+    "AssistantNaiveMemory",
+    "AssistantZeroMemory",
+    "AssistantMessageWindowMemory",
     "AssistantTool",
     "AssistantRetrieval",
     "AssistantToolType",
     "AssistantRetrievalType",
+    "AssistantMemory"
 ]
 
 
@@ -37,7 +44,16 @@ class Assistant(TaskingaiBaseModel):
     name: str
     description: str
     system_prompt_template: List[str]
+    memory: AssistantMemory
     tools: List[AssistantTool]
     retrievals: List[AssistantRetrieval]
     metadata: Dict[str, str]
     created_timestamp: int
+
+
+    @field_validator('memory', mode='before')
+    def validate_memory(cls, memory_dict: Dict):
+        memory: AssistantMemory = build_assistant_memory(memory_dict)
+        if not memory:
+            raise ValueError('Invalid input memory')
+        return memory

taskingai/client/models/entity/assistant/assistant_memory.py

@@ -0,0 +1,61 @@
+from enum import Enum
+from pydantic import BaseModel, Field
+from abc import ABC
+from typing import Optional, Dict
+
+
+class AssistantMemoryType(str, Enum):
+    zero = "zero"
+    naive = "naive"
+    message_window = "message_window"
+
+
+class AssistantMemory(BaseModel, ABC):
+    type: AssistantMemoryType = Field(..., description="The type of the memory.")
+
+
+class AssistantMessageWindowMemory(AssistantMemory):
+    type: AssistantMemoryType = Field(AssistantMemoryType.message_window, Literal=AssistantMemoryType.message_window)
+    max_messages: int = Field(...)
+    max_tokens: int = Field(...)
+
+
+class AssistantNaiveMemory(AssistantMemory):
+    type: AssistantMemoryType = Field(AssistantMemoryType.naive, Literal=AssistantMemoryType.naive)
+
+
+class AssistantZeroMemory(AssistantMemory):
+    type: AssistantMemoryType = Field(AssistantMemoryType.zero, Literal=AssistantMemoryType.zero)
+
+
+def build_assistant_memory(memory_dict: Dict) -> Optional[AssistantMemory]:
+    # Check if the memory dictionary is provided and has the 'type' key
+    if not memory_dict or 'type' not in memory_dict:
+        return None
+
+    memory_type = memory_dict['type']
+
+    if memory_type == AssistantMemoryType.zero.value:
+        # For zero memory, no additional information is needed
+        return AssistantZeroMemory()
+
+    elif memory_type == AssistantMemoryType.naive.value:
+        # For naive memory, no additional configuration is needed
+        return AssistantNaiveMemory()
+
+    elif memory_type == AssistantMemoryType.message_window.value:
+        # For message window memory, additional configuration is needed
+        max_messages = memory_dict.get('max_messages')
+        max_tokens = memory_dict.get('max_tokens')
+
+        # Validate that required fields are present
+        if max_messages is None or max_tokens is None:
+            return None
+
+        return AssistantMessageWindowMemory(max_messages=max_messages, max_tokens=max_tokens)
+
+    else:
+        # If the memory type is unknown, return None
+        return None
+
+