diff --git a/packages/dbgpt-client/src/dbgpt_client/knowledge.py b/packages/dbgpt-client/src/dbgpt_client/knowledge.py index a391ba6a5..f136edfd6 100644 --- a/packages/dbgpt-client/src/dbgpt_client/knowledge.py +++ b/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: diff --git a/packages/dbgpt-client/src/dbgpt_client/tests/test_knowledge.py b/packages/dbgpt-client/src/dbgpt_client/tests/test_knowledge.py new file mode 100644 index 000000000..9beb47d93 --- /dev/null +++ b/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"