diff --git a/.env.template b/.env.template index 799113fc2..38ec1c61f 100644 --- a/.env.template +++ b/.env.template @@ -70,6 +70,7 @@ EMBEDDING_MODEL=text2vec #EMBEDDING_MODEL=bge-large-zh KNOWLEDGE_CHUNK_SIZE=500 KNOWLEDGE_SEARCH_TOP_SIZE=5 +KNOWLEDGE_GRAPH_SEARCH_TOP_SIZE=50 ## Maximum number of chunks to load at once, if your single document is too large, ## you can set this value to a higher value for better performance. ## if out of memory when load large document, you can set this value to a lower value. @@ -138,10 +139,12 @@ LOCAL_DB_TYPE=sqlite EXECUTE_LOCAL_COMMANDS=False #*******************************************************************# -#** VECTOR STORE SETTINGS **# +#** VECTOR STORE / KNOWLEDGE GRAPH SETTINGS **# #*******************************************************************# -### Chroma vector db config VECTOR_STORE_TYPE=Chroma +GRAPH_STORE_TYPE=TuGraph + +### Chroma vector db config #CHROMA_PERSIST_PATH=/root/DB-GPT/pilot/data ### Milvus vector db config @@ -163,6 +166,15 @@ ElasticSearch_PORT=9200 ElasticSearch_USERNAME=elastic ElasticSearch_PASSWORD=i=+iLw9y0Jduq86XTi6W +### TuGraph config +#TUGRAPH_HOST=127.0.0.1 +#TUGRAPH_PORT=7687 +#TUGRAPH_USERNAME=admin +#TUGRAPH_PASSWORD=73@TuGraph +#TUGRAPH_VERTEX_TYPE=entity +#TUGRAPH_EDGE_TYPE=relation +#TUGRAPH_EDGE_NAME_KEY=label + #*******************************************************************# #** WebServer Language Support **# #*******************************************************************# diff --git a/.mypy.ini b/.mypy.ini index ae26c8dc3..e2c2bc3ab 100644 --- a/.mypy.ini +++ b/.mypy.ini @@ -112,3 +112,6 @@ ignore_missing_imports = True [mypy-ollama.*] ignore_missing_imports = True + +[mypy-networkx.*] +ignore_missing_imports = True diff --git a/dbgpt/_private/config.py b/dbgpt/_private/config.py index 60c3908f7..7c0279a11 100644 --- a/dbgpt/_private/config.py +++ b/dbgpt/_private/config.py @@ -254,6 +254,9 @@ class Config(metaclass=Singleton): self.KNOWLEDGE_CHUNK_SIZE = int(os.getenv("KNOWLEDGE_CHUNK_SIZE", 100)) self.KNOWLEDGE_CHUNK_OVERLAP = int(os.getenv("KNOWLEDGE_CHUNK_OVERLAP", 50)) self.KNOWLEDGE_SEARCH_TOP_SIZE = int(os.getenv("KNOWLEDGE_SEARCH_TOP_SIZE", 5)) + self.KNOWLEDGE_GRAPH_SEARCH_TOP_SIZE = int( + os.getenv("KNOWLEDGE_GRAPH_SEARCH_TOP_SIZE", 50) + ) self.KNOWLEDGE_MAX_CHUNKS_ONCE_LOAD = int( os.getenv("KNOWLEDGE_MAX_CHUNKS_ONCE_LOAD", 10) ) diff --git a/dbgpt/app/knowledge/api.py b/dbgpt/app/knowledge/api.py index a891c46ea..2203d4c2c 100644 --- a/dbgpt/app/knowledge/api.py +++ b/dbgpt/app/knowledge/api.py @@ -13,6 +13,7 @@ from dbgpt.app.knowledge.request.request import ( DocumentSummaryRequest, DocumentSyncRequest, EntityExtractRequest, + GraphVisRequest, KnowledgeDocumentRequest, KnowledgeQueryRequest, KnowledgeSpaceRequest, @@ -75,7 +76,7 @@ def space_delete(request: KnowledgeSpaceRequest): try: return Result.succ(knowledge_space_service.delete_space(request.name)) except Exception as e: - return Result.failed(code="E000X", msg=f"space list error {e}") + return Result.failed(code="E000X", msg=f"space delete error {e}") @router.post("/knowledge/{space_name}/arguments") @@ -84,7 +85,7 @@ def arguments(space_name: str): try: return Result.succ(knowledge_space_service.arguments(space_name)) except Exception as e: - return Result.failed(code="E000X", msg=f"space list error {e}") + return Result.failed(code="E000X", msg=f"space arguments error {e}") @router.post("/knowledge/{space_name}/argument/save") @@ -95,7 +96,7 @@ def arguments_save(space_name: str, argument_request: SpaceArgumentRequest): knowledge_space_service.argument_save(space_name, argument_request) ) except Exception as e: - return Result.failed(code="E000X", msg=f"space list error {e}") + return Result.failed(code="E000X", msg=f"space save error {e}") @router.post("/knowledge/{space_name}/document/add") @@ -156,6 +157,20 @@ def document_list(space_name: str, query_request: DocumentQueryRequest): return Result.failed(code="E000X", msg=f"document list error {e}") +@router.post("/knowledge/{space_name}/graphvis") +def graph_vis(space_name: str, query_request: GraphVisRequest): + print(f"/document/list params: {space_name}, {query_request}") + print(query_request.limit) + try: + return Result.succ( + knowledge_space_service.query_graph( + space_name=space_name, limit=query_request.limit + ) + ) + except Exception as e: + return Result.failed(code="E000X", msg=f"get graph vis error {e}") + + @router.post("/knowledge/{space_name}/document/delete") def document_delete(space_name: str, query_request: DocumentQueryRequest): print(f"/document/list params: {space_name}, {query_request}") @@ -164,7 +179,7 @@ def document_delete(space_name: str, query_request: DocumentQueryRequest): knowledge_space_service.delete_document(space_name, query_request.doc_name) ) except Exception as e: - return Result.failed(code="E000X", msg=f"document list error {e}") + return Result.failed(code="E000X", msg=f"document delete error {e}") @router.post("/knowledge/{space_name}/document/upload") @@ -232,7 +247,7 @@ def document_sync(space_name: str, request: DocumentSyncRequest): @router.post("/knowledge/{space_name}/document/sync_batch") -def batch_document_sync( +async def batch_document_sync( space_name: str, request: List[KnowledgeSyncRequest], service: Service = Depends(get_rag_service), @@ -242,13 +257,13 @@ def batch_document_sync( space = service.get({"name": space_name}) for sync_request in request: sync_request.space_id = space.id - doc_ids = service.sync_document(requests=request) + doc_ids = await service.sync_document(requests=request) # doc_ids = service.sync_document( # space_name=space_name, sync_requests=request # ) return Result.succ({"tasks": doc_ids}) except Exception as e: - return Result.failed(code="E000X", msg=f"document sync error {e}") + return Result.failed(code="E000X", msg=f"document sync batch error {e}") @router.post("/knowledge/{space_name}/chunk/list") diff --git a/dbgpt/app/knowledge/request/request.py b/dbgpt/app/knowledge/request/request.py index 7c4897e03..ba2bb2c15 100644 --- a/dbgpt/app/knowledge/request/request.py +++ b/dbgpt/app/knowledge/request/request.py @@ -53,6 +53,10 @@ class DocumentQueryRequest(BaseModel): page_size: int = 20 +class GraphVisRequest(BaseModel): + limit: int = 100 + + class DocumentSyncRequest(BaseModel): """Sync request""" diff --git a/dbgpt/app/knowledge/service.py b/dbgpt/app/knowledge/service.py index 9119d3981..9755ea04b 100644 --- a/dbgpt/app/knowledge/service.py +++ b/dbgpt/app/knowledge/service.py @@ -25,8 +25,9 @@ from dbgpt.app.knowledge.request.response import ( ) from dbgpt.component import ComponentType from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG -from dbgpt.core import Chunk +from dbgpt.core import Chunk, LLMClient from dbgpt.model import DefaultLLMClient +from dbgpt.model.cluster import WorkerManagerFactory from dbgpt.rag.assembler.embedding import EmbeddingAssembler from dbgpt.rag.assembler.summary import SummaryAssembler from dbgpt.rag.chunk_manager import ChunkParameters @@ -39,7 +40,7 @@ from dbgpt.rag.text_splitter.text_splitter import ( ) from dbgpt.serve.rag.api.schemas import KnowledgeSyncRequest from dbgpt.serve.rag.models.models import KnowledgeSpaceDao, KnowledgeSpaceEntity -from dbgpt.serve.rag.service.service import Service, SyncStatus +from dbgpt.serve.rag.service.service import SyncStatus from dbgpt.storage.vector_store.base import VectorStoreConfig from dbgpt.storage.vector_store.connector import VectorStoreConnector from dbgpt.util.executor_utils import ExecutorFactory, blocking_func_to_async @@ -52,7 +53,6 @@ document_chunk_dao = DocumentChunkDao() logger = logging.getLogger(__name__) CFG = Config() - # default summary max iteration call with llm. DEFAULT_SUMMARY_MAX_ITERATION = 5 # default summary concurrency call with llm. @@ -70,6 +70,13 @@ class KnowledgeService: def __init__(self): pass + @property + def llm_client(self) -> LLMClient: + worker_manager = CFG.SYSTEM_APP.get_component( + ComponentType.WORKER_MANAGER_FACTORY, WorkerManagerFactory + ).create() + return DefaultLLMClient(worker_manager, True) + def create_knowledge_space(self, request: KnowledgeSpaceRequest): """create knowledge space Args: @@ -332,16 +339,23 @@ class KnowledgeService: embedding_fn = embedding_factory.create( model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL] ) + + spaces = self.get_knowledge_space(KnowledgeSpaceRequest(name=space_name)) + if len(spaces) != 1: + raise Exception(f"invalid space name:{space_name}") + space = spaces[0] + from dbgpt.storage.vector_store.base import VectorStoreConfig config = VectorStoreConfig( - name=space_name, + name=space.name, embedding_fn=embedding_fn, max_chunks_once_load=CFG.KNOWLEDGE_MAX_CHUNKS_ONCE_LOAD, + llm_client=self.llm_client, + model_name=self.model_name, ) vector_store_connector = VectorStoreConnector( - vector_store_type=CFG.VECTOR_STORE_TYPE, - vector_store_config=config, + vector_store_type=space.vector_type, vector_store_config=config ) knowledge = KnowledgeFactory.create( datasource=doc.content, @@ -442,21 +456,27 @@ class KnowledgeService: Args: - space_name: knowledge space name """ - query = KnowledgeSpaceEntity(name=space_name) - spaces = knowledge_space_dao.get_knowledge_space(query) - if len(spaces) == 0: - raise Exception(f"delete error, no space name:{space_name} in database") + spaces = knowledge_space_dao.get_knowledge_space( + KnowledgeSpaceEntity(name=space_name) + ) + if len(spaces) != 1: + raise Exception(f"invalid space name:{space_name}") space = spaces[0] + embedding_factory = CFG.SYSTEM_APP.get_component( "embedding_factory", EmbeddingFactory ) embedding_fn = embedding_factory.create( model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL] ) - config = VectorStoreConfig(name=space.name, embedding_fn=embedding_fn) + config = VectorStoreConfig( + name=space.name, + embedding_fn=embedding_fn, + llm_client=self.llm_client, + model_name=None, + ) vector_store_connector = VectorStoreConnector( - vector_store_type=CFG.VECTOR_STORE_TYPE, - vector_store_config=config, + vector_store_type=space.vector_type, vector_store_config=config ) # delete vectors vector_store_connector.delete_vector_name(space.name) @@ -480,6 +500,12 @@ class KnowledgeService: documents = knowledge_document_dao.get_documents(document_query) if len(documents) != 1: raise Exception(f"there are no or more than one document called {doc_name}") + + spaces = self.get_knowledge_space(KnowledgeSpaceRequest(name=space_name)) + if len(spaces) != 1: + raise Exception(f"invalid space name:{space_name}") + space = spaces[0] + vector_ids = documents[0].vector_ids if vector_ids is not None: embedding_factory = CFG.SYSTEM_APP.get_component( @@ -488,10 +514,14 @@ class KnowledgeService: embedding_fn = embedding_factory.create( model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL] ) - config = VectorStoreConfig(name=space_name, embedding_fn=embedding_fn) + config = VectorStoreConfig( + name=space.name, + embedding_fn=embedding_fn, + llm_client=self.llm_client, + model_name=None, + ) vector_store_connector = VectorStoreConnector( - vector_store_type=CFG.VECTOR_STORE_TYPE, - vector_store_config=config, + vector_store_type=space.vector_type, vector_store_config=config ) # delete vector by ids vector_store_connector.delete_by_ids(vector_ids) @@ -535,7 +565,7 @@ class KnowledgeService: """ logger.info( - f"async doc embedding sync, doc:{doc.doc_name}, chunks length is {len(chunk_docs)}, begin embedding to vector store-{CFG.VECTOR_STORE_TYPE}" + f"async doc embedding sync, doc:{doc.doc_name}, chunks length is {len(chunk_docs)}" ) try: with root_tracer.start_span( @@ -645,3 +675,40 @@ class KnowledgeService: **{"chat_param": chat_param}, ) return chat + + def query_graph(self, space_name, limit): + embedding_factory = CFG.SYSTEM_APP.get_component( + "embedding_factory", EmbeddingFactory + ) + embedding_fn = embedding_factory.create( + model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL] + ) + spaces = self.get_knowledge_space(KnowledgeSpaceRequest(name=space_name)) + if len(spaces) != 1: + raise Exception(f"invalid space name:{space_name}") + space = spaces[0] + print(CFG.LLM_MODEL) + config = VectorStoreConfig( + name=space.name, + embedding_fn=embedding_fn, + max_chunks_once_load=CFG.KNOWLEDGE_MAX_CHUNKS_ONCE_LOAD, + llm_client=self.llm_client, + model_name=None, + ) + + vector_store_connector = VectorStoreConnector( + vector_store_type=space.vector_type, vector_store_config=config + ) + graph = vector_store_connector.client.query_graph(limit=limit) + res = {"nodes": [], "edges": []} + for node in graph.vertices(): + res["nodes"].append({"vid": node.vid}) + for edge in graph.edges(): + res["edges"].append( + { + "src": edge.sid, + "dst": edge.tid, + "label": edge.props[graph.edge_label], + } + ) + return res diff --git a/dbgpt/app/scene/chat_knowledge/v1/chat.py b/dbgpt/app/scene/chat_knowledge/v1/chat.py index 297d9dd04..2cfef3eed 100644 --- a/dbgpt/app/scene/chat_knowledge/v1/chat.py +++ b/dbgpt/app/scene/chat_knowledge/v1/chat.py @@ -9,6 +9,7 @@ from dbgpt.app.knowledge.document_db import ( KnowledgeDocumentDao, KnowledgeDocumentEntity, ) +from dbgpt.app.knowledge.request.request import KnowledgeSpaceRequest from dbgpt.app.knowledge.service import KnowledgeService from dbgpt.app.scene import BaseChat, ChatScene from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG @@ -50,7 +51,7 @@ class ChatKnowledge(BaseChat): ) self.space_context = self.get_space_context(self.knowledge_space) self.top_k = ( - CFG.KNOWLEDGE_SEARCH_TOP_SIZE + self.get_knowledge_search_top_size(self.knowledge_space) if self.space_context is None else int(self.space_context["embedding"]["topk"]) ) @@ -73,12 +74,27 @@ class ChatKnowledge(BaseChat): embedding_fn = embedding_factory.create( model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL] ) + from dbgpt.serve.rag.models.models import ( + KnowledgeSpaceDao, + KnowledgeSpaceEntity, + ) from dbgpt.storage.vector_store.base import VectorStoreConfig - config = VectorStoreConfig(name=self.knowledge_space, embedding_fn=embedding_fn) + spaces = KnowledgeSpaceDao().get_knowledge_space( + KnowledgeSpaceEntity(name=self.knowledge_space) + ) + if len(spaces) != 1: + raise Exception(f"invalid space name:{self.knowledge_space}") + space = spaces[0] + + config = VectorStoreConfig( + name=space.name, + embedding_fn=embedding_fn, + llm_client=self.llm_client, + llm_model=self.llm_model, + ) vector_store_connector = VectorStoreConnector( - vector_store_type=CFG.VECTOR_STORE_TYPE, - vector_store_config=config, + vector_store_type=space.vector_type, vector_store_config=config ) query_rewrite = None if CFG.KNOWLEDGE_SEARCH_REWRITE: @@ -239,6 +255,18 @@ class ChatKnowledge(BaseChat): service = KnowledgeService() return service.get_space_context(space_name) + def get_knowledge_search_top_size(self, space_name) -> int: + service = KnowledgeService() + request = KnowledgeSpaceRequest(name=space_name) + spaces = service.get_knowledge_space(request) + if len(spaces) == 1: + from dbgpt.storage import vector_store + + if spaces[0].vector_type in vector_store.__knowledge_graph__: + return CFG.KNOWLEDGE_GRAPH_SEARCH_TOP_SIZE + + return CFG.KNOWLEDGE_SEARCH_TOP_SIZE + async def execute_similar_search(self, query): """execute similarity search""" with root_tracer.start_span( diff --git a/dbgpt/app/static/404.html b/dbgpt/app/static/404.html index 184e3b581..822d93636 100644 --- a/dbgpt/app/static/404.html +++ b/dbgpt/app/static/404.html @@ -1 +1 @@ -