mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-27 13:57:46 +00:00
48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
"""Knowledge graph base class."""
|
|
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from typing import List, Optional
|
|
|
|
from pydantic import Field
|
|
|
|
from dbgpt._private.pydantic import ConfigDict
|
|
from dbgpt.core import Chunk
|
|
from dbgpt.rag.index.base import IndexStoreBase, IndexStoreConfig
|
|
from dbgpt.storage.graph_store.graph import Graph
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KnowledgeGraphConfig(IndexStoreConfig):
|
|
"""Knowledge graph config."""
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
|
|
|
|
|
|
class KnowledgeGraphBase(IndexStoreBase, ABC):
|
|
"""Knowledge graph base class."""
|
|
|
|
@abstractmethod
|
|
def get_config(self) -> KnowledgeGraphConfig:
|
|
"""Get the knowledge graph config."""
|
|
|
|
@abstractmethod
|
|
def query_graph(self, limit: Optional[int] = None) -> Graph:
|
|
"""Get graph data."""
|
|
|
|
@abstractmethod
|
|
def delete_by_ids(self, ids: str) -> List[str]:
|
|
"""Delete document by ids."""
|
|
|
|
|
|
class ParagraphChunk(Chunk):
|
|
"""Loaded chunk, used in GraphRAG."""
|
|
|
|
chunk_parent_id: str = Field(default=None, description="id of parent chunk")
|
|
chunk_parent_name: str = Field(default=None, description="parent chunk name")
|
|
parent_content: str = Field(default=None, description="parent chunk text content")
|
|
parent_is_document: bool = Field(
|
|
default=False, description="is parent chunk a document"
|
|
)
|