KnowledgeGraphAgent

class KnowledgeGraphAgent(ChatAgent):

An agent that can extract node and relationship information for different entities from given Element content.

Attributes: task_prompt (TextPrompt): A prompt for the agent to extract node and relationship information for different entities.

init

def __init__(self, model: Optional[BaseModelBackend] = None):

Initialize the KnowledgeGraphAgent.

Parameters:

  • model (BaseModelBackend, optional): The model backend to use for generating responses. (default: :obj:OpenAIModel with GPT_4O_MINI)

run

def run(
    self,
    element: 'Element',
    parse_graph_elements: bool = False,
    prompt: Optional[str] = None
):

Run the agent to extract node and relationship information.

Parameters:

  • element (Element): The input element.
  • parse_graph_elements (bool, optional): Whether to parse into GraphElement. Defaults to False.
  • prompt (str, optional): The custom prompt to be used. Defaults to None.

Returns:

Union[str, GraphElement]: The extracted node and relationship information. If parse_graph_elements is True then return GraphElement, else return str.

_validate_node

def _validate_node(self, node: Node):

Validate if the object is a valid Node.

Parameters:

  • node (Node): Object to be validated.

Returns:

bool: True if the object is a valid Node, False otherwise.

_validate_relationship

def _validate_relationship(self, relationship: Relationship):

Validate if the object is a valid Relationship.

Parameters:

  • relationship (Relationship): Object to be validated.

Returns:

bool: True if the object is a valid Relationship, False otherwise.

_parse_graph_elements

def _parse_graph_elements(self, input_string: str):

Parses graph elements from given content.

Parameters:

  • input_string (str): The input content.

Returns:

GraphElement: The parsed graph elements.