test: add test for assistant and inference with dict

Author: taskingaijc

taskingai/client/models/entities/action.py

@@ -31,11 +31,7 @@ class Action(BaseModel):
     description: str = Field(..., min_length=1, max_length=512)
     url: str = Field(...)
     method: ActionMethod = Field(...)
-    path_param_schema: Optional[Dict[str, ActionParam]] = Field(None)
-    query_param_schema: Optional[Dict[str, ActionParam]] = Field(None)
     body_type: ActionBodyType = Field(...)
-    body_param_schema: Optional[Dict[str, ActionParam]] = Field(None)
-    function_def: ChatCompletionFunction = Field(...)
     openapi_schema: Dict[str, Any] = Field(...)
     authentication: ActionAuthentication = Field(...)
     updated_timestamp: int = Field(..., ge=0)

taskingai/inference/chat_completion.py

@@ -59,21 +59,25 @@ def _validate_chat_completion_params(
     def _validate_message(msg: Union[SystemMessage, UserMessage, AssistantMessage, FunctionMessage, Dict[str, Any]]):
         if isinstance(msg, Dict):
             if msg["role"] == ChatCompletionRole.SYSTEM.value:
+                msg.pop("role")
                 return SystemMessage(**msg)
             elif msg["role"] == ChatCompletionRole.USER.value:
+                msg.pop("role")
                 return UserMessage(**msg)
             elif msg["role"] == ChatCompletionRole.ASSISTANT.value:
+                msg.pop("role")
                 return AssistantMessage(**msg)
             elif msg["role"] == ChatCompletionRole.FUNCTION.value:
+                msg.pop("role")
                 return FunctionMessage(**msg)
             else:
                 raise ValueError("Invalid message role.")
 
         elif (
-            isinstance(msg, SystemMessage)
-            or isinstance(msg, UserMessage)
-            or isinstance(msg, AssistantMessage)
-            or isinstance(msg, FunctionMessage)
+            isinstance(msg, ChatCompletionSystemMessage)
+            or isinstance(msg, ChatCompletionUserMessage)
+            or isinstance(msg, ChatCompletionAssistantMessage)
+            or isinstance(msg, ChatCompletionFunctionMessage)
         ):
             return msg
 

test/common/utils.py

@@ -135,11 +135,14 @@ def assume_assistant_result(assistant_dict: dict, res: dict):
     for key, value in assistant_dict.items():
         if key == 'system_prompt_template' and isinstance(value, str):
             pytest.assume(res[key] == [assistant_dict[key]])
-        elif key in ["memory", "tool", "retrievals"]:
+        elif key in ['retrieval_configs']:
+            if isinstance(value, dict):
+                pytest.assume(vars(res[key]) == assistant_dict[key])
+            else:
+                pytest.assume(res[key] == assistant_dict[key])
+        elif key in ["memory", "tools", "retrievals"]:
             continue
         else:
-            if key == 'retrieval_configs':
-                res[key] = vars(res[key])
             pytest.assume(res[key] == assistant_dict[key])
 
 

test/testcase/test_async/test_async_assistant.py

@@ -2,7 +2,7 @@
 
 from taskingai.assistant import *
 from taskingai.client.models import ToolRef, ToolType, RetrievalRef, RetrievalType, RetrievalConfig
-from taskingai.assistant.memory import AssistantNaiveMemory
+from taskingai.assistant.memory import AssistantNaiveMemory, AssistantZeroMemory
 from test.config import Config
 from test.common.logger import logger
 from test.common.utils import list_to_dict
@@ -13,16 +13,6 @@
 @pytest.mark.test_async
 class TestAssistant(Base):
 
-    retrieval_configs_list = [
-       {"method": "memory", "top_k": 2, "max_tokens": 4000},
-       RetrievalConfig(
-            method="memory",
-            top_k=1,
-            max_tokens=5000,
-
-        )
-    ]
-
     @pytest.mark.run(order=51)
     @pytest.mark.asyncio
     async def test_a_create_assistant(self):
@@ -62,7 +52,11 @@ async def test_a_create_assistant(self):
         }
         for i in range(4):
             if i == 0:
+                assistant_dict.update({"memory": {"type": "naive"}})
+                assistant_dict.update({"retrievals": [{"type": "collection", "id": self.collection_id}]})
                 assistant_dict.update({"retrieval_configs": {"method": "memory", "top_k": 2, "max_tokens": 4000}})
+                assistant_dict.update({"tools": [{"type": "action", "id": self.action_id},
+                                                 {"type": "plugin", "id": "open_weather/get_hourly_forecast"}]})
             res = await a_create_assistant(**assistant_dict)
             res_dict = vars(res)
             logger.info(f'response_dict:{res_dict}, except_dict:{assistant_dict}')
@@ -106,22 +100,54 @@ async def test_a_get_assistant(self):
 
     @pytest.mark.run(order=54)
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("retrieval_configs", retrieval_configs_list)
-    async def test_a_update_assistant(self, retrieval_configs):
+    async def test_a_update_assistant(self):
 
         # Update an assistant.
 
-        name = "openai"
-        description = "test for openai"
+        update_data_list = [
+            {
+                "name": "openai",
+                "description": "test for openai",
+                "memory": AssistantZeroMemory(),
+                "retrievals": [
+                    RetrievalRef(
+                        type=RetrievalType.COLLECTION,
+                        id=self.collection_id,
+                    ),
+                ],
+                "retrieval_configs": RetrievalConfig(
+                    method="memory",
+                    top_k=2,
+                    max_tokens=4000,
 
-        res = await a_update_assistant(assistant_id=self.assistant_id, name=name, description=description, retrieval_configs=retrieval_configs)
-        res_dict = vars(res)
-        pytest.assume(res_dict["name"] == name)
-        pytest.assume(res_dict["description"] == description)
-        if isinstance(retrieval_configs, dict):
-            pytest.assume(vars(res_dict["retrieval_configs"]) == retrieval_configs)
-        else:
-            pytest.assume(res_dict["retrieval_configs"] == retrieval_configs)
+                ),
+                "tools": [
+                    ToolRef(
+                        type=ToolType.ACTION,
+                        id=self.action_id,
+                    ),
+                    ToolRef(
+                        type=ToolType.PLUGIN,
+                        id="open_weather/get_hourly_forecast",
+                    )
+                ]
+            },
+            {
+                "name": "openai",
+                "description": "test for openai",
+                "memory": {"type": "naive"},
+                "retrievals": [{"type": "collection", "id": self.collection_id}],
+                "retrieval_configs": {"method": "memory", "top_k": 2, "max_tokens": 4000},
+                "tools": [{"type": "action", "id": self.action_id},
+                          {"type": "plugin", "id": "open_weather/get_hourly_forecast"}]
+
+            }
+        ]
+        for update_data in update_data_list:
+            res = await a_update_assistant(assistant_id=self.assistant_id, **update_data)
+            res_dict = vars(res)
+            logger.info(f'response_dict:{res_dict}, except_dict:{update_data}')
+            assume_assistant_result(update_data, res_dict)
 
     @pytest.mark.run(order=66)
     @pytest.mark.asyncio