fix: handle knowledge space id responses (#3070)

Co-authored-by: wolfkill <27876473+wolfkill@users.noreply.github.com>

packages/dbgpt-client/src/dbgpt_client/knowledge.py

@@ -1,7 +1,7 @@
 """Knowledge API client."""
 
 import json
-from typing import List
+from typing import Any, List
 
 from dbgpt._private.pydantic import model_to_dict, model_to_json
 from dbgpt.core.schema.api import Result
@@ -10,6 +10,14 @@ from .client import Client, ClientException
 from .schema import DocumentModel, SpaceModel, SyncModel
 
 
+def _space_model_from_create_response(data: Any, request: SpaceModel) -> SpaceModel:
+    if isinstance(data, dict):
+        return SpaceModel(**data)
+    space = model_to_dict(request)
+    space["id"] = data
+    return SpaceModel(**space)
+
+
 async def create_space(client: Client, space_model: SpaceModel) -> SpaceModel:
     """Create a new space.
 
@@ -25,7 +33,7 @@ async def create_space(client: Client, space_model: SpaceModel) -> SpaceModel:
         res = await client.post("/knowledge/spaces", model_to_dict(space_model))
         result: Result = res.json()
         if result["success"]:
-            return SpaceModel(**result["data"])
+            return _space_model_from_create_response(result["data"], space_model)
         else:
             raise ClientException(status=result["err_code"], reason=result)
     except Exception as e:

packages/dbgpt-client/src/dbgpt_client/tests/test_knowledge.py

@@ -0,0 +1,69 @@
+import pytest
+
+from dbgpt_client.knowledge import create_space
+from dbgpt_client.schema import SpaceModel
+
+
+class _Response:
+    def __init__(self, payload):
+        self._payload = payload
+
+    def json(self):
+        return self._payload
+
+
+class _CreateSpaceClient:
+    def __init__(self, payload):
+        self.payload = payload
+        self.requests = []
+
+    async def post(self, path, body):
+        self.requests.append((path, body))
+        return _Response(self.payload)
+
+
+@pytest.mark.asyncio
+async def test_create_space_accepts_primary_key_response():
+    client = _CreateSpaceClient(
+        {"success": True, "err_code": None, "err_msg": None, "data": 2}
+    )
+    request = SpaceModel(
+        name="test_space",
+        vector_type="Chroma",
+        desc="for client space",
+        owner="dbgpt",
+    )
+
+    created = await create_space(client, request)
+
+    assert created.id == 2
+    assert created.name == "test_space"
+    assert created.vector_type == "Chroma"
+    assert created.desc == "for client space"
+    assert created.owner == "dbgpt"
+    assert client.requests[0][0] == "/knowledge/spaces"
+
+
+@pytest.mark.asyncio
+async def test_create_space_preserves_mapping_response():
+    client = _CreateSpaceClient(
+        {
+            "success": True,
+            "err_code": None,
+            "err_msg": None,
+            "data": {
+                "id": 2,
+                "name": "server_space",
+                "vector_type": "Milvus",
+                "desc": "from server",
+                "owner": "dbgpt",
+            },
+        }
+    )
+
+    created = await create_space(client, SpaceModel(name="request_space"))
+
+    assert created.id == 2
+    assert created.name == "server_space"
+    assert created.vector_type == "Milvus"
+    assert created.desc == "from server"