"""Define the CommunitySummaryKnowledgeGraph.""" import logging import os from typing import List, Optional from dbgpt._private.pydantic import ConfigDict, Field from dbgpt.core import Chunk from dbgpt.rag.transformer.community_summarizer import CommunitySummarizer from dbgpt.rag.transformer.graph_extractor import GraphExtractor from dbgpt.storage.knowledge_graph.community.community_store import CommunityStore from dbgpt.storage.knowledge_graph.community.factory import CommunityStoreAdapterFactory from dbgpt.storage.knowledge_graph.knowledge_graph import ( BuiltinKnowledgeGraph, BuiltinKnowledgeGraphConfig, ) from dbgpt.storage.vector_store.base import VectorStoreConfig from dbgpt.storage.vector_store.factory import VectorStoreFactory from dbgpt.storage.vector_store.filters import MetadataFilters logger = logging.getLogger(__name__) class CommunitySummaryKnowledgeGraphConfig(BuiltinKnowledgeGraphConfig): """Community summary knowledge graph config.""" model_config = ConfigDict(arbitrary_types_allowed=True) vector_store_type: str = Field( default="Chroma", description="The type of vector store." ) user: Optional[str] = Field( default=None, description="The user of vector store, if not set, will use the default user.", ) password: Optional[str] = Field( default=None, description=( "The password of vector store, if not set, will use the default password." ), ) extract_topk: int = Field( default=5, description="Topk of knowledge graph extract", ) extract_score_threshold: float = Field( default=0.3, description="Recall score of knowledge graph extract", ) community_topk: int = Field( default=50, description="Topk of community search in knowledge graph", ) community_score_threshold: float = Field( default=0.0, description="Recall score of community search in knowledge graph", ) class CommunitySummaryKnowledgeGraph(BuiltinKnowledgeGraph): """Community summary knowledge graph class.""" def __init__(self, config: CommunitySummaryKnowledgeGraphConfig): """Initialize community summary knowledge graph class.""" super().__init__(config) self._config = config self._vector_store_type = os.getenv( "VECTOR_STORE_TYPE", config.vector_store_type ) self._extract_topk = int( os.getenv("KNOWLEDGE_GRAPH_EXTRACT_SEARCH_TOP_SIZE", config.extract_topk) ) self._extract_score_threshold = float( os.getenv( "KNOWLEDGE_GRAPH_EXTRACT_SEARCH_RECALL_SCORE", config.extract_score_threshold, ) ) self._community_topk = int( os.getenv( "KNOWLEDGE_GRAPH_COMMUNITY_SEARCH_TOP_SIZE", config.community_topk ) ) self._community_score_threshold = float( os.getenv( "KNOWLEDGE_GRAPH_COMMUNITY_SEARCH_RECALL_SCORE", config.community_score_threshold, ) ) def extractor_configure(name: str, cfg: VectorStoreConfig): cfg.name = name cfg.embedding_fn = config.embedding_fn cfg.max_chunks_once_load = config.max_chunks_once_load cfg.max_threads = config.max_threads cfg.user = config.user cfg.password = config.password cfg.topk = self._extract_topk cfg.score_threshold = self._extract_score_threshold self._graph_extractor = GraphExtractor( self._llm_client, self._model_name, VectorStoreFactory.create( self._vector_store_type, config.name + "_CHUNK_HISTORY", extractor_configure, ), ) def community_store_configure(name: str, cfg: VectorStoreConfig): cfg.name = name cfg.embedding_fn = config.embedding_fn cfg.max_chunks_once_load = config.max_chunks_once_load cfg.max_threads = config.max_threads cfg.user = config.user cfg.password = config.password cfg.topk = self._community_topk cfg.score_threshold = self._community_score_threshold self._community_store = CommunityStore( CommunityStoreAdapterFactory.create(self._graph_store), CommunitySummarizer(self._llm_client, self._model_name), VectorStoreFactory.create( self._vector_store_type, config.name + "_COMMUNITY_SUMMARY", community_store_configure, ), ) def get_config(self) -> BuiltinKnowledgeGraphConfig: """Get the knowledge graph config.""" return self._config async def aload_document(self, chunks: List[Chunk]) -> List[str]: """Extract and persist graph.""" # todo add doc node for chunk in chunks: # todo add chunk node # todo add relation doc-chunk # extract graphs and save graphs = await self._graph_extractor.extract(chunk.content) for graph in graphs: self._graph_store.insert_graph(graph) # build communities and save await self._community_store.build_communities() return [chunk.chunk_id for chunk in chunks] async def asimilar_search_with_scores( self, text, topk, score_threshold: float, filters: Optional[MetadataFilters] = None, ) -> List[Chunk]: """Retrieve relevant community summaries.""" # global search: retrieve relevant community summaries communities = await self._community_store.search_communities(text) summaries = [ f"Section {i + 1}:\n{community.summary}" for i, community in enumerate(communities) ] context = "\n".join(summaries) if summaries else "" # local search: extract keywords and explore subgraph keywords = await self._keyword_extractor.extract(text) subgraph = self._graph_store.explore(keywords, limit=topk).format() logger.info(f"Search subgraph from {len(keywords)} keywords") if not summaries and not subgraph: return [] # merge search results into context content = HYBRID_SEARCH_PT_CN.format(context=context, graph=subgraph) return [Chunk(content=content)] def truncate(self) -> List[str]: """Truncate knowledge graph.""" logger.info("Truncate community store") self._community_store.truncate() logger.info("Truncate keyword extractor") self._keyword_extractor.truncate() logger.info("Truncate triplet extractor") self._graph_extractor.truncate() return [self._config.name] def delete_vector_name(self, index_name: str): """Delete knowledge graph.""" logger.info("Drop community store") self._community_store.drop() logger.info("Drop keyword extractor") self._keyword_extractor.drop() logger.info("Drop triplet extractor") self._graph_extractor.drop() HYBRID_SEARCH_PT_CN = ( "## 角色\n" "你非常擅长结合提示词模板提供的[上下文]信息与[知识图谱]信息，" "准确恰当地回答用户的问题，并保证不会输出与上下文和知识图谱无关的信息。" "\n" "## 技能\n" "### 技能 1: 上下文理解\n" "- 准确地理解[上下文]提供的信息，上下文信息可能被拆分为多个章节。\n" "- 上下文的每个章节内容都会以[Section]开始，并按需进行了编号。\n" "- 上下文信息提供了与用户问题相关度最高的总结性描述，请合理使用它们。" "### 技能 2: 知识图谱理解\n" "- 准确地识别[知识图谱]中提供的[Entities:]章节中的实体信息" "和[Relationships:]章节中的关系信息，实体和关系信息的一般格式为：\n" "```" "* 实体信息格式:\n" "- (实体名)\n" "- (实体名:实体描述)\n" "- (实体名:实体属性表)\n" "- (文本块ID:文档块内容)\n" "- (目录ID:目录名)\n" "- (文档ID:文档名称)\n" "\n" "* 关系信息的格式:\n" "- (来源实体名)-[关系名]->(目标实体名)\n" "- (来源实体名)-[关系名:关系描述]->(目标实体名)\n" "- (来源实体名)-[关系名:关系属性表]->(目标实体名)\n" "- (文本块实体)-[包含]->(实体名)\n" "- (目录ID)-[包含]->(文本块实体)\n" "- (目录ID)-[包含]->(子目录ID)\n" "- (文档ID)-[包含]->(文本块实体)\n" "- (文档ID)-[包含]->(目录ID)\n" "```" "- 正确地将关系信息中的实体名/ID与实体信息关联，还原出图结构。" "- 将图结构所表达的信息作为用户提问的明细上下文，辅助生成更好的答案。\n" "\n" "## 约束条件\n" "- 不要在答案中描述你的思考过程，直接给出用户问题的答案，不要生成无关信息。\n" "- 若[知识图谱]没有提供信息，此时应根据[上下文]提供的信息回答问题。" "- 确保以第三人称书写，从客观角度结合[上下文]和[知识图谱]表达的信息回答问题。\n" "- 若提供的信息相互矛盾，请解决矛盾并提供一个单一、连贯的描述。\n" "- 避免使用停用词和过于常见的词汇。\n" "\n" "## 参考案例\n" "```\n" "[上下文]:\n" "Section 1:\n" "菲尔・贾伯的大儿子叫雅各布・贾伯。\n" "Section 2:\n" "菲尔・贾伯的小儿子叫比尔・贾伯。\n" "[知识图谱]:\n" "Entities:\n" "(菲尔・贾伯#菲尔兹咖啡创始人)\n" "(菲尔兹咖啡#加利福尼亚州伯克利创立的咖啡品牌)\n" "(雅各布・贾伯#菲尔・贾伯的儿子)\n" "(美国多地#菲尔兹咖啡的扩展地区)\n" "\n" "Relationships:\n" "(菲尔・贾伯#创建#菲尔兹咖啡#1978年在加利福尼亚州伯克利创立)\n" "(菲尔兹咖啡#位于#加利福尼亚州伯克利#菲尔兹咖啡的创立地点)\n" "(菲尔・贾伯#拥有#雅各布・贾伯#菲尔・贾伯的儿子)\n" "(雅各布・贾伯#担任#首席执行官#在2005年成为菲尔兹咖啡的首席执行官)\n" "(菲尔兹咖啡#扩展至#美国多地#菲尔兹咖啡的扩展范围)\n" "```\n" "\n" "----\n" "\n" "接下来的[上下文]和[知识图谱]的信息，可以帮助你回答更好地用户的问题。\n" "\n" "[上下文]:\n" "{context}\n" "\n" "[知识图谱]:\n" "{graph}\n" "\n" ) HYBRID_SEARCH_PT_EN = ( "## Role\n" "You excel at combining the information provided in the [Context] with " "information from the [KnowledgeGraph] to accurately and appropriately " "answer user questions, ensuring that you do not output information " "unrelated to the context and knowledge graph.\n" "\n" "## Skills\n" "### Skill 1: Context Understanding\n" "- Accurately understand the information provided in the [Context], " "which may be divided into several sections.\n" "- Each section in the context will start with [Section] " "and may be numbered as needed.\n" "- The context provides a summary description most relevant to the user’s " "question, and it should be used wisely." "### Skill 2: Knowledge Graph Understanding\n" "- Accurately identify entity information in the [Entities:] section and " "relationship information in the [Relationships:] section " "of the [KnowledgeGraph]. The general format for entity " "and relationship information is:\n" "```" "* Entity Information Format:\n" "- (entity_name)\n" "- (entity_name: entity_description)\n" "- (entity_name: entity_property_map)\n" "- (chunk_id: chunk_content)\n" "- (catalog_id: catalog_name)\n" "- (document_id: document_name)\n" "\n" "* Relationship Information Format:\n" "- (source_entity_name)-[relationship_name]->(target_entity_name)\n" "- (source_entity_name)-[relationship_name: relationship_description]->" "(target_entity_name)\n" "- (source_entity_name)-[relationship_name: relationship_property_map]->" "(target_entity_name)\n" "- (chunk_id)-[Contains]->(entity_name)\n" "- (catalog_id)-[Contains]->(chunk_id)\n" "- (catalog_id)-[Contains]->(sub_catalog_id)\n" "- (document_id)-[Contains]->(chunk_id)\n" "- (document_id)-[Contains]->(catalog_id)\n" "```" "- Correctly associate entity names/IDs in the relationship information " "with entity information to restore the graph structure." "- Use the information expressed by the graph structure as detailed " "context for the user's query to assist in generating better answers.\n" "\n" "## Constraints\n" "- Don't describe your thought process in the answer, provide the answer " "to the user's question directly without generating irrelevant information." "- If the [KnowledgeGraph] does not provide information, you should answer " "the question based on the information provided in the [Context]." "- Ensure to write in the third person, responding to questions from " "an objective perspective based on the information combined from the " "[Context] and the [KnowledgeGraph].\n" "- If the provided information is contradictory, resolve the " "contradictions and provide a single, coherent description.\n" "- Avoid using stop words and overly common vocabulary.\n" "\n" "## Reference Example\n" "```\n" "[Context]:\n" "Section 1:\n" "Phil Schiller's eldest son is Jacob Schiller.\n" "Section 2:\n" "Phil Schiller's youngest son is Bill Schiller.\n" "[KnowledgeGraph]:\n" "Entities:\n" "(Phil Jaber#Founder of Philz Coffee)\n" "(Philz Coffee#Coffee brand founded in Berkeley, California)\n" "(Jacob Jaber#Son of Phil Jaber)\n" "(Multiple locations in the USA#Expansion regions of Philz Coffee)\n" "\n" "Relationships:\n" "(Phil Jaber#Created#Philz Coffee" "#Founded in Berkeley, California in 1978)\n" "(Philz Coffee#Located in#Berkeley, California" "#Founding location of Philz Coffee)\n" "(Phil Jaber#Has#Jacob Jaber#Son of Phil Jaber)\n" "(Jacob Jaber#Serves as#CEO#Became CEO of Philz Coffee in 2005)\n" "(Philz Coffee#Expanded to#Multiple locations in the USA" "#Expansion regions of Philz Coffee)\n" "```\n" "\n" "----\n" "\n" "The following information from the [Context] and [KnowledgeGraph] can " "help you better answer user questions.\n" "\n" "[Context]:\n" "{context}\n" "\n" "[KnowledgeGraph]:\n" "{graph}\n" "\n" )