diff --git a/tools/knowlege_init.py b/tools/knowlege_init.py index 56889bd7f..bf1cb1b68 100644 --- a/tools/knowlege_init.py +++ b/tools/knowlege_init.py @@ -5,6 +5,8 @@ import os import sys from pilot.embedding_engine.knowledge_type import KnowledgeType +from pilot.openapi.knowledge.knowledge_service import KnowledgeService +from pilot.openapi.knowledge.request.knowledge_request import KnowledgeSpaceRequest sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) @@ -14,6 +16,7 @@ from pilot.configs.model_config import ( LLM_MODEL_CONFIG, ) from pilot.embedding_engine.knowledge_embedding import KnowledgeEmbedding +knowledge_space_service = KnowledgeService() CFG = Config() @@ -23,7 +26,7 @@ class LocalKnowledgeInit: def __init__(self, vector_store_config) -> None: self.vector_store_config = vector_store_config - self.model_name = LLM_MODEL_CONFIG["text2vec"] + self.model_name = LLM_MODEL_CONFIG[CFG.EMBEDDING_MODEL] def knowledge_persist(self, file_path): """knowledge persist""" @@ -39,6 +42,12 @@ class LocalKnowledgeInit: ) client = ke.init_knowledge_embedding() client.source_embedding() + print(f"""begin create {self.vector_store_config["vector_store_name"]} space""") + space = KnowledgeSpaceRequest + space.name = self.vector_store_config["vector_store_name"] + space.desc = "" + space.owner = "knowledge_init.py" + knowledge_space_service.create_knowledge_space(space) if __name__ == "__main__":