From 2e68d6e28d1e460ed55b18da1fb9fcea651677b4 Mon Sep 17 00:00:00 2001
From: aries_ckt <916701291@qq.com>
Date: Tue, 4 Jul 2023 11:39:57 +0800
Subject: [PATCH] feat:space return doc_count

---
 pilot/server/knowledge/request/response.py | 14 ++++++++++++++
 pilot/server/knowledge/service.py          | 15 +++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/pilot/server/knowledge/request/response.py b/pilot/server/knowledge/request/response.py
index 7fbf36155..98dcbcc76 100644
--- a/pilot/server/knowledge/request/response.py
+++ b/pilot/server/knowledge/request/response.py
@@ -21,3 +21,17 @@ class DocumentQueryResponse(BaseModel):
     total: int = None
     """page: current page"""
     page: int = None
+
+
+class SpaceQueryResponse(BaseModel):
+    """data: data"""
+
+    name: str = None
+    """vector_type: vector type"""
+    vector_type: str = None
+    """desc: description"""
+    desc: str = None
+    """owner: owner"""
+    owner: str = None
+    """doc_count: doc_count"""
+    doc_count: int = None
diff --git a/pilot/server/knowledge/service.py b/pilot/server/knowledge/service.py
index 4d99be4e3..f076669c9 100644
--- a/pilot/server/knowledge/service.py
+++ b/pilot/server/knowledge/service.py
@@ -27,7 +27,7 @@ from enum import Enum
 
 from pilot.server.knowledge.request.response import (
     ChunkQueryResponse,
-    DocumentQueryResponse,
+    DocumentQueryResponse, SpaceQueryResponse,
 )
 
 knowledge_space_dao = KnowledgeSpaceDao()
@@ -86,7 +86,18 @@ class KnowledgeService:
         query = KnowledgeSpaceEntity(
             name=request.name, vector_type=request.vector_type, owner=request.owner
         )
-        return knowledge_space_dao.get_knowledge_space(query)
+        res = SpaceQueryResponse()
+        space: KnowledgeSpaceEntity = knowledge_space_dao.get_knowledge_space(query)
+        res.name = space.name
+        res.vector_type = space.vector_type
+        res.desc = space.desc
+        query = KnowledgeDocumentEntity(
+            space=space,
+        )
+        res.doc_count = knowledge_document_dao.get_knowledge_documents_count(query)
+        return res
+
+
 
     """get knowledge get_knowledge_documents"""