diff --git a/assets/schema/knowledge_management.sql b/assets/schema/knowledge_management.sql index 2c50e42b4..ceaa566da 100644 --- a/assets/schema/knowledge_management.sql +++ b/assets/schema/knowledge_management.sql @@ -8,7 +8,7 @@ CREATE TABLE `knowledge_space` ( `gmt_modified` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time', PRIMARY KEY (`id`), KEY `idx_name` (`name`) COMMENT 'index:idx_name' -) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='knowledge space table'; +) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COMMENT='knowledge space table'; CREATE TABLE `knowledge_document` ( `id` int NOT NULL AUTO_INCREMENT COMMENT 'auto increment id', @@ -25,7 +25,7 @@ CREATE TABLE `knowledge_document` ( `gmt_modified` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time', PRIMARY KEY (`id`), KEY `idx_doc_name` (`doc_name`) COMMENT 'index:idx_doc_name' -) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='knowledge document table'; +) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COMMENT='knowledge document table'; CREATE TABLE `document_chunk` ( `id` int NOT NULL AUTO_INCREMENT COMMENT 'auto increment id', @@ -38,4 +38,4 @@ CREATE TABLE `document_chunk` ( `gmt_modified` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time', PRIMARY KEY (`id`), KEY `idx_document_id` (`document_id`) COMMENT 'index:document_id' -) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='knowledge document chunk detail' \ No newline at end of file +) ENGINE=InnoDB AUTO_INCREMENT=100001 DEFAULT CHARSET=utf8mb4 COMMENT='knowledge document chunk detail' \ No newline at end of file diff --git a/pilot/embedding_engine/markdown_embedding.py b/pilot/embedding_engine/markdown_embedding.py index 0d70ba34f..225cdad82 100644 --- a/pilot/embedding_engine/markdown_embedding.py +++ b/pilot/embedding_engine/markdown_embedding.py @@ -6,7 +6,7 @@ from typing import List import markdown from bs4 import BeautifulSoup from langchain.schema import Document -from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter +from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter, RecursiveCharacterTextSplitter from pilot.configs.config import Config from pilot.embedding_engine import SourceEmbedding, register @@ -37,11 +37,14 @@ class MarkdownEmbedding(SourceEmbedding): length_function=len, ) else: - text_splitter = SpacyTextSplitter( - pipeline="zh_core_web_sm", - chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, - chunk_overlap=100, - ) + try: + text_splitter = SpacyTextSplitter( + pipeline="zh_core_web_sm", + chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, + chunk_overlap=100, + ) + except Exception: + text_splitter = RecursiveCharacterTextSplitter(chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50) return loader.load_and_split(text_splitter) @register diff --git a/pilot/embedding_engine/pdf_embedding.py b/pilot/embedding_engine/pdf_embedding.py index 2b8f244e3..3a47e115e 100644 --- a/pilot/embedding_engine/pdf_embedding.py +++ b/pilot/embedding_engine/pdf_embedding.py @@ -4,7 +4,7 @@ from typing import List from langchain.document_loaders import PyPDFLoader from langchain.schema import Document -from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter +from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter, RecursiveCharacterTextSplitter from pilot.configs.config import Config from pilot.embedding_engine import SourceEmbedding, register @@ -40,11 +40,14 @@ class PDFEmbedding(SourceEmbedding): length_function=len, ) else: - text_splitter = SpacyTextSplitter( - pipeline="zh_core_web_sm", - chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, - chunk_overlap=100, - ) + try: + text_splitter = SpacyTextSplitter( + pipeline="zh_core_web_sm", + chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, + chunk_overlap=100, + ) + except Exception: + text_splitter = RecursiveCharacterTextSplitter(chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50) return loader.load_and_split(text_splitter) @register diff --git a/pilot/embedding_engine/ppt_embedding.py b/pilot/embedding_engine/ppt_embedding.py index da4390849..2ddc583f9 100644 --- a/pilot/embedding_engine/ppt_embedding.py +++ b/pilot/embedding_engine/ppt_embedding.py @@ -4,7 +4,7 @@ from typing import List from langchain.document_loaders import UnstructuredPowerPointLoader from langchain.schema import Document -from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter +from langchain.text_splitter import SpacyTextSplitter, CharacterTextSplitter, RecursiveCharacterTextSplitter from pilot.configs.config import Config from pilot.embedding_engine import SourceEmbedding, register @@ -38,11 +38,14 @@ class PPTEmbedding(SourceEmbedding): length_function=len, ) else: - text_splitter = SpacyTextSplitter( - pipeline="zh_core_web_sm", - chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, - chunk_overlap=100, - ) + try: + text_splitter = SpacyTextSplitter( + pipeline="zh_core_web_sm", + chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, + chunk_overlap=100, + ) + except Exception: + text_splitter = RecursiveCharacterTextSplitter(chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50) return loader.load_and_split(text_splitter) @register diff --git a/pilot/embedding_engine/source_embedding.py b/pilot/embedding_engine/source_embedding.py index 372e35c22..6d7500007 100644 --- a/pilot/embedding_engine/source_embedding.py +++ b/pilot/embedding_engine/source_embedding.py @@ -59,7 +59,7 @@ class SourceEmbedding(ABC): self.vector_client = VectorStoreConnector( CFG.VECTOR_STORE_TYPE, self.vector_store_config ) - self.vector_client.load_document(docs) + return self.vector_client.load_document(docs) @register def similar_search(self, doc, topk): diff --git a/pilot/embedding_engine/url_embedding.py b/pilot/embedding_engine/url_embedding.py index ce9e5360d..273b5e777 100644 --- a/pilot/embedding_engine/url_embedding.py +++ b/pilot/embedding_engine/url_embedding.py @@ -3,7 +3,7 @@ from typing import List from bs4 import BeautifulSoup from langchain.document_loaders import WebBaseLoader from langchain.schema import Document -from langchain.text_splitter import CharacterTextSplitter, SpacyTextSplitter +from langchain.text_splitter import CharacterTextSplitter, SpacyTextSplitter, RecursiveCharacterTextSplitter from pilot.configs.config import Config from pilot.configs.model_config import KNOWLEDGE_CHUNK_SPLIT_SIZE @@ -33,11 +33,14 @@ class URLEmbedding(SourceEmbedding): length_function=len, ) else: - text_splitter = SpacyTextSplitter( - pipeline="zh_core_web_sm", - chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, - chunk_overlap=100, - ) + try: + text_splitter = SpacyTextSplitter( + pipeline="zh_core_web_sm", + chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, + chunk_overlap=100, + ) + except Exception: + text_splitter = RecursiveCharacterTextSplitter(chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50) return loader.load_and_split(text_splitter) @register diff --git a/pilot/embedding_engine/word_embedding.py b/pilot/embedding_engine/word_embedding.py index 55988a240..efb337a9a 100644 --- a/pilot/embedding_engine/word_embedding.py +++ b/pilot/embedding_engine/word_embedding.py @@ -4,7 +4,7 @@ from typing import List from langchain.document_loaders import UnstructuredWordDocumentLoader from langchain.schema import Document -from langchain.text_splitter import CharacterTextSplitter, SpacyTextSplitter +from langchain.text_splitter import CharacterTextSplitter, SpacyTextSplitter, RecursiveCharacterTextSplitter from pilot.configs.config import Config from pilot.embedding_engine import SourceEmbedding, register @@ -32,11 +32,14 @@ class WordEmbedding(SourceEmbedding): length_function=len, ) else: - text_splitter = SpacyTextSplitter( - pipeline="zh_core_web_sm", - chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, - chunk_overlap=100, - ) + try: + text_splitter = SpacyTextSplitter( + pipeline="zh_core_web_sm", + chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, + chunk_overlap=100, + ) + except Exception: + text_splitter = RecursiveCharacterTextSplitter(chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50) return loader.load_and_split(text_splitter) @register diff --git a/pilot/openapi/api_v1/api_v1.py b/pilot/openapi/api_v1/api_v1.py index c9a4e0b87..990ebeefb 100644 --- a/pilot/openapi/api_v1/api_v1.py +++ b/pilot/openapi/api_v1/api_v1.py @@ -26,8 +26,8 @@ from pilot.openapi.api_v1.api_view_model import ( ChatSceneVo, ) from pilot.configs.config import Config -from pilot.openapi.knowledge.knowledge_service import KnowledgeService -from pilot.openapi.knowledge.request.knowledge_request import KnowledgeSpaceRequest +from pilot.server.knowledge.service import KnowledgeService +from pilot.server.knowledge.request.request import KnowledgeSpaceRequest from pilot.scene.base_chat import BaseChat from pilot.scene.base import ChatScene diff --git a/pilot/server/dbgpt_server.py b/pilot/server/dbgpt_server.py index 032b2b896..86cb8b4bb 100644 --- a/pilot/server/dbgpt_server.py +++ b/pilot/server/dbgpt_server.py @@ -23,7 +23,7 @@ from fastapi import FastAPI, applications from fastapi.openapi.docs import get_swagger_ui_html from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware -from pilot.openapi.knowledge.knowledge_controller import router as knowledge_router +from pilot.server.knowledge.api import router as knowledge_router from pilot.openapi.api_v1.api_v1 import router as api_v1, validation_exception_handler diff --git a/pilot/openapi/knowledge/__init__.py b/pilot/server/knowledge/__init__.py similarity index 100% rename from pilot/openapi/knowledge/__init__.py rename to pilot/server/knowledge/__init__.py diff --git a/pilot/openapi/knowledge/knowledge_controller.py b/pilot/server/knowledge/api.py similarity index 94% rename from pilot/openapi/knowledge/knowledge_controller.py rename to pilot/server/knowledge/api.py index 26e79609f..887bbadb2 100644 --- a/pilot/openapi/knowledge/knowledge_controller.py +++ b/pilot/server/knowledge/api.py @@ -2,7 +2,7 @@ import os import shutil from tempfile import NamedTemporaryFile -from fastapi import APIRouter, File, UploadFile, Request, Form +from fastapi import APIRouter, File, UploadFile, Form from langchain.embeddings import HuggingFaceEmbeddings @@ -12,8 +12,8 @@ from pilot.configs.model_config import LLM_MODEL_CONFIG, KNOWLEDGE_UPLOAD_ROOT_P from pilot.openapi.api_v1.api_view_model import Result from pilot.embedding_engine.knowledge_embedding import KnowledgeEmbedding -from pilot.openapi.knowledge.knowledge_service import KnowledgeService -from pilot.openapi.knowledge.request.knowledge_request import ( +from pilot.server.knowledge.service import KnowledgeService +from pilot.server.knowledge.request.request import ( KnowledgeQueryRequest, KnowledgeQueryResponse, KnowledgeDocumentRequest, @@ -22,7 +22,7 @@ from pilot.openapi.knowledge.request.knowledge_request import ( DocumentQueryRequest, ) -from pilot.openapi.knowledge.request.knowledge_request import KnowledgeSpaceRequest +from pilot.server.knowledge.request.request import KnowledgeSpaceRequest CFG = Config() router = APIRouter() diff --git a/pilot/openapi/knowledge/document_chunk_dao.py b/pilot/server/knowledge/chunk_db.py similarity index 100% rename from pilot/openapi/knowledge/document_chunk_dao.py rename to pilot/server/knowledge/chunk_db.py diff --git a/pilot/openapi/knowledge/knowledge_document_dao.py b/pilot/server/knowledge/document_db.py similarity index 100% rename from pilot/openapi/knowledge/knowledge_document_dao.py rename to pilot/server/knowledge/document_db.py diff --git a/pilot/openapi/knowledge/request/__init__.py b/pilot/server/knowledge/request/__init__.py similarity index 100% rename from pilot/openapi/knowledge/request/__init__.py rename to pilot/server/knowledge/request/__init__.py diff --git a/pilot/openapi/knowledge/request/knowledge_request.py b/pilot/server/knowledge/request/request.py similarity index 100% rename from pilot/openapi/knowledge/request/knowledge_request.py rename to pilot/server/knowledge/request/request.py diff --git a/pilot/openapi/knowledge/request/knowledge_response.py b/pilot/server/knowledge/request/response.py similarity index 58% rename from pilot/openapi/knowledge/request/knowledge_response.py rename to pilot/server/knowledge/request/response.py index 7fbf36155..98dcbcc76 100644 --- a/pilot/openapi/knowledge/request/knowledge_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/openapi/knowledge/knowledge_service.py b/pilot/server/knowledge/service.py similarity index 94% rename from pilot/openapi/knowledge/knowledge_service.py rename to pilot/server/knowledge/service.py index b6308a81a..ef8cd3d89 100644 --- a/pilot/openapi/knowledge/knowledge_service.py +++ b/pilot/server/knowledge/service.py @@ -5,19 +5,19 @@ from pilot.configs.config import Config from pilot.configs.model_config import LLM_MODEL_CONFIG from pilot.embedding_engine.knowledge_embedding import KnowledgeEmbedding from pilot.logs import logger -from pilot.openapi.knowledge.document_chunk_dao import ( +from pilot.server.knowledge.chunk_db import ( DocumentChunkEntity, DocumentChunkDao, ) -from pilot.openapi.knowledge.knowledge_document_dao import ( +from pilot.server.knowledge.document_db import ( KnowledgeDocumentDao, KnowledgeDocumentEntity, ) -from pilot.openapi.knowledge.knowledge_space_dao import ( +from pilot.server.knowledge.space_db import ( KnowledgeSpaceDao, KnowledgeSpaceEntity, ) -from pilot.openapi.knowledge.request.knowledge_request import ( +from pilot.server.knowledge.request.request import ( KnowledgeSpaceRequest, KnowledgeDocumentRequest, DocumentQueryRequest, @@ -25,9 +25,9 @@ from pilot.openapi.knowledge.request.knowledge_request import ( ) from enum import Enum -from pilot.openapi.knowledge.request.knowledge_response import ( +from pilot.server.knowledge.request.response import ( ChunkQueryResponse, - DocumentQueryResponse, + DocumentQueryResponse, SpaceQueryResponse, ) knowledge_space_dao = KnowledgeSpaceDao() @@ -195,7 +195,8 @@ class KnowledgeService: vector_ids = client.knowledge_embedding_batch(chunk_docs) doc.status = SyncStatus.FINISHED.name doc.result = "document embedding success" - doc.vector_ids = ",".join(vector_ids) + if vector_ids is not None: + doc.vector_ids = ",".join(vector_ids) logger.info(f"async document embedding, success:{doc.doc_name}") except Exception as e: doc.status = SyncStatus.FAILED.name diff --git a/pilot/openapi/knowledge/knowledge_space_dao.py b/pilot/server/knowledge/space_db.py similarity index 97% rename from pilot/openapi/knowledge/knowledge_space_dao.py rename to pilot/server/knowledge/space_db.py index 16def1e99..d4f6fbba6 100644 --- a/pilot/openapi/knowledge/knowledge_space_dao.py +++ b/pilot/server/knowledge/space_db.py @@ -5,7 +5,7 @@ from sqlalchemy.ext.declarative import declarative_base from pilot.configs.config import Config -from pilot.openapi.knowledge.request.knowledge_request import KnowledgeSpaceRequest +from pilot.server.knowledge.request.request import KnowledgeSpaceRequest from sqlalchemy.orm import sessionmaker CFG = Config() diff --git a/tools/knowledge_init.py b/tools/knowledge_init.py index 34460e398..6f46b713f 100644 --- a/tools/knowledge_init.py +++ b/tools/knowledge_init.py @@ -6,8 +6,8 @@ import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) 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 +from pilot.server.knowledge.service import KnowledgeService +from pilot.server.knowledge.request.request import KnowledgeSpaceRequest from pilot.configs.config import Config @@ -46,12 +46,18 @@ class LocalKnowledgeInit: docs.extend(doc) embedding_engine.index_to_store(docs) 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 = "knowledge_init.py" - space.vector_type = CFG.VECTOR_STORE_TYPE - space.owner = "DB-GPT" - knowledge_space_service.create_knowledge_space(space) + try: + space = KnowledgeSpaceRequest + space.name = self.vector_store_config["vector_store_name"] + space.desc = "knowledge_init.py" + space.vector_type = CFG.VECTOR_STORE_TYPE + space.owner = "DB-GPT" + knowledge_space_service.create_knowledge_space(space) + except Exception as e: + if "have already named" in str(e): + print(f"Warning: you have already named {space.name}") + else: + raise e if __name__ == "__main__":