diff --git a/src/retk/controllers/ai/knowledge.py b/src/retk/controllers/ai/knowledge.py index 84b8ae5..e50d80e 100644 --- a/src/retk/controllers/ai/knowledge.py +++ b/src/retk/controllers/ai/knowledge.py @@ -9,14 +9,24 @@ async def get_extended_nodes( au: AuthedUser, ) -> schemas.ai.GetExtendedNodesResponse: docs = await core.ai.llm.knowledge.extended.get_extended_nodes(uid=au.u.id) - return schemas.ai.GetExtendedNodesResponse( - requestId=au.request_id, - nodes=[schemas.ai.GetExtendedNodesResponse.Node( + nodes = [] + for doc in docs: + res = doc["extendMd"].split("\n", 1) + if len(res) == 2: + title, content = res + else: + title, content = res[0], "" + node = schemas.ai.GetExtendedNodesResponse.Node( id=str(doc["_id"]), sourceNid=doc["sourceNid"], sourceTitle=doc["sourceMd"].split("\n", 1)[0].strip(), - md=doc["extendMd"], - ) for doc in docs] + title=title.strip(), + content=content.strip(), + ) + nodes.append(node) + return schemas.ai.GetExtendedNodesResponse( + requestId=au.request_id, + nodes=nodes ) diff --git a/src/retk/controllers/schemas/ai.py b/src/retk/controllers/schemas/ai.py index 317998f..947e5e6 100644 --- a/src/retk/controllers/schemas/ai.py +++ b/src/retk/controllers/schemas/ai.py @@ -8,7 +8,8 @@ class Node(BaseModel): id: str sourceNid: str sourceTitle: str - md: str + title: str + content: str requestId: str nodes: List[Node] diff --git a/src/retk/controllers/user.py b/src/retk/controllers/user.py index 54c9e84..c610004 100644 --- a/src/retk/controllers/user.py +++ b/src/retk/controllers/user.py @@ -15,6 +15,7 @@ async def get_user( max_space = 0 else: max_space = const.USER_TYPE.id2config(au.u.type).max_store_space + total_nodes = await core.user.get_user_nodes_count(uid=au.u.id, disabled=False, in_trash=False) return schemas.user.UserInfoResponse( requestId=au.request_id, user=schemas.user.UserInfoResponse.User( @@ -38,7 +39,7 @@ async def get_user( editorSepRightWidth=au.u.settings.editor_sep_right_width, editorSideCurrentToolId=au.u.settings.editor_side_current_tool_id, ), - totalNodes=await core.user.get_user_nodes_count(uid=au.u.id, disabled=False, in_trash=False), + totalNodes=total_nodes, ), ) diff --git a/src/retk/routes/ai.py b/src/retk/routes/ai.py index 71f6900..6fe198d 100644 --- a/src/retk/routes/ai.py +++ b/src/retk/routes/ai.py @@ -8,7 +8,7 @@ router = APIRouter( prefix="/api/ai", - tags=["node"], + tags=["ai"], responses={404: {"description": "Not found"}}, ) diff --git a/tests/test_api.py b/tests/test_api.py index ba69bb8..fef2319 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1598,7 +1598,8 @@ async def test_node_extend(self): rj = self.check_ok_response(resp, 200) self.assertEqual(1, len(rj["nodes"])) n = rj["nodes"][0] - self.assertEqual("this is extended md", n["md"]) + self.assertEqual("this is extended md", n["title"]) + self.assertEqual("", n["content"]) self.assertEqual(node["id"], n["sourceNid"]) resp = self.client.post(