diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/community_summary.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/community_summary.py index 8c29bb6a8..1bd45bd95 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/community_summary.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/community_summary.py @@ -22,7 +22,6 @@ from dbgpt_ext.storage.knowledge_graph.community.community_store import Communit from dbgpt_ext.storage.knowledge_graph.knowledge_graph import ( GRAPH_PARAMETERS, BuiltinKnowledgeGraph, - BuiltinKnowledgeGraphConfig, ) logger = logging.getLogger(__name__) @@ -146,14 +145,34 @@ logger = logging.getLogger(__name__) description=_("Community Summary Knowledge Graph."), parameters=[ Parameter.build_from( - _("Community Summary Knowledge Graph Config."), + _("Graph Store Config"), "config", - BuiltinKnowledgeGraphConfig, - description=_("Community Summary Knowledge Graph Config."), + GraphStoreConfig, + description=_("graph store config."), + ), + Parameter.build_from( + _("Graph Store Name"), + "name", + str, + optional=True, + default="dbgpt", + description=_("Graph Store Name"), + ), + Parameter.build_from( + _("LLM Client"), + "llm_client", + LLMClient, + description=_("llm client for extract graph triplets."), + ), + Parameter.build_from( + _("LLM Model Name"), + "llm_model", + str, + description=_("kg extract llm model name."), optional=True, default=None, ), - ], + ] ) class CommunitySummaryKnowledgeGraph(BuiltinKnowledgeGraph): """Community summary knowledge graph class.""" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py index c8396a718..3184af298 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/knowledge_graph/knowledge_graph.py @@ -3,7 +3,6 @@ import asyncio import logging import os -from dataclasses import dataclass, field from typing import List, Optional, Any from dbgpt.core import Chunk, Embeddings, LLMClient @@ -11,7 +10,7 @@ from dbgpt.core.awel.flow import Parameter, ResourceCategory, register_resource from dbgpt.rag.transformer.keyword_extractor import KeywordExtractor from dbgpt.storage.graph_store.base import GraphStoreBase, GraphStoreConfig from dbgpt.storage.graph_store.graph import Graph -from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase, KnowledgeGraphConfig +from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase from dbgpt.storage.vector_store.filters import MetadataFilters from dbgpt.util.i18n_utils import _ from dbgpt_ext.rag.transformer.triplet_extractor import TripletExtractor