agents.memory.kg_generator_agent¶

Knowledge Graph Generator Agent for Memory System.

This agent specializes in extracting and maintaining knowledge graphs from memories, building entity relationships and semantic connections across the memory system.

Classes¶

`KGGeneratorAgent`	Agent specializing in knowledge graph generation from memory system.
`KGGeneratorAgentConfig`	Configuration for Knowledge Graph Generator Agent.
`KnowledgeGraphNode`	Represents a node (entity) in the knowledge graph.
`KnowledgeGraphRelationship`	Represents a relationship (edge) in the knowledge graph.
`MemoryKnowledgeGraph`	Complete knowledge graph structure extracted from memory system.

Module Contents¶

class agents.memory.kg_generator_agent.KGGeneratorAgent(/, **data)¶

Bases: haive.agents.simple.agent.SimpleAgent

Agent specializing in knowledge graph generation from memory system.

The KGGeneratorAgent is a specialized memory agent that extracts entities and relationships from memory content to build comprehensive knowledge graphs. It uses LLM-powered extraction to identify semantic structures and connections across memories.

This agent extends SimpleAgent with knowledge graph capabilities, providing methods for: - Entity extraction with confidence scoring - Relationship discovery between entities - Incremental graph building and updates - Entity neighborhood exploration - Graph-based memory analysis

Parameters:: data (Any)

memory_store¶: Manager for memory storage and retrieval operations

classifier¶: Classifier for analyzing memory content and types

knowledge_graph¶: The maintained knowledge graph structure

graph_transformer¶: Transformer for graph processing operations

extract_batch_size¶: Number of memories to process in each batch

min_confidence_threshold¶: Minimum confidence score for extracted entities/relationships

enable_iterative_refinement¶: Whether to enable iterative graph refinement

entity_types¶: List of entity types the agent should extract

relationship_types¶: List of relationship types the agent should extract

entity_extraction_prompt¶: Prompt template for entity extraction

relationship_extraction_prompt¶: Prompt template for relationship extraction

Examples

Basic usage:

# Create configuration
config = KGGeneratorAgentConfig(
    memory_store_manager=store_manager,
    memory_classifier=classifier,
    engine=AugLLMConfig(temperature=0.1)
)

# Create agent
kg_agent = KGGeneratorAgent(config)

# Extract knowledge graph from memories
graph = await kg_agent.extract_knowledge_graph_from_memories()
print(f"Extracted {len(graph.nodes)} entities and {len(graph.relationships)} relationships")

Entity extraction:

# Extract entities with limit
entities = await kg_agent.extract_entities_from_memories(limit=10)

for entity in entities:
    print(f"Entity: {entity.name} ({entity.type}) - Confidence: {entity.confidence}")

Relationship extraction:

# Extract relationships
relationships = await kg_agent.extract_relationships_from_memories(limit=10)

for rel in relationships:
    print(f"Relationship: {rel.relationship_type} "
          f"({rel.source_id} -> {rel.target_id})")

Entity exploration:

# Get entity context with neighborhood
context = await kg_agent.get_entity_context("Python")

if "entity" in context:
    entity = context["entity"]
    neighborhood = context["neighborhood"]
    print(f"Entity: {entity.name}")
    print(f"Connected nodes: {neighborhood['total_nodes']}")
    print(f"Relationships: {neighborhood['total_relationships']}")

Agent execution:

# Use agent for interactive queries
response = await kg_agent.run("Extract knowledge graph from recent memories")
print(response)

# Get graph statistics
stats = await kg_agent.run("Show me graph statistics")
print(stats)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

async extract_entities_from_memories(limit=None, namespace=None)¶

Extract entities from memories and return them as a list.

Parameters:

limit (int | None) – Maximum number of memories to process
namespace (tuple[str, Ellipsis] | None) – Memory namespace to process

Returns:

List of extracted knowledge graph nodes (entities)

Return type:

list[KnowledgeGraphNode]

async extract_knowledge_graph_from_memories(memory_ids=None, namespace=None, memory_types=None)¶

Extract knowledge graph from specified memories.

Processes memories to extract entities and relationships, building a comprehensive knowledge graph. Uses LLM-powered extraction to identify semantic structures and connections across memory content.

Parameters:

memory_ids (list[str] | None) – Specific memory IDs to process. If None, processes all memories in the specified namespace and types.
namespace (tuple[str, Ellipsis] | None) – Memory namespace to process. If None, processes all namespaces.
memory_types (list[haive.agents.memory.core.types.MemoryType] | None) – Specific memory types to process. If None, processes all types.

Returns:

Updated MemoryKnowledgeGraph containing extracted entities and relationships with metadata about the extraction process.

Raises:

ValueError – If no memories are found to process
RuntimeError – If extraction fails due to LLM or processing errors

Return type:

MemoryKnowledgeGraph

Examples

Extract from all memories:

# Extract from all memories
graph = await kg_agent.extract_knowledge_graph_from_memories()
print(f"Extracted {len(graph.nodes)} entities")
print(f"Extracted {len(graph.relationships)} relationships")

Extract from specific memories:

# Extract from specific memory IDs
memory_ids = ["mem_123", "mem_456", "mem_789"]
graph = await kg_agent.extract_knowledge_graph_from_memories(
    memory_ids=memory_ids
)

Extract from specific namespace:

# Extract from user's personal memories
graph = await kg_agent.extract_knowledge_graph_from_memories(
    namespace=("user", "personal")
)

Extract from specific memory types:

# Extract only from semantic and episodic memories
graph = await kg_agent.extract_knowledge_graph_from_memories(
    memory_types=[MemoryType.SEMANTIC, MemoryType.EPISODIC]
)

Note

The extraction process runs in batches (controlled by extract_batch_size) to manage memory usage and API rate limits. Progress is logged during processing.

async extract_relationships_from_memories(limit=None, namespace=None)¶

Extract relationships from memories and return them as a list.

Parameters:

limit (int | None) – Maximum number of memories to process
namespace (tuple[str, Ellipsis] | None) – Memory namespace to process

Returns:

List of extracted knowledge graph relationships

Return type:

list[KnowledgeGraphRelationship]

async get_entity_context(entity_name)¶

Get context for an entity by name.

Parameters:: entity_name (str) – Name of the entity to get context for
Returns:: Dictionary containing entity context information
Return type:: dict[str, Any]

async get_entity_neighborhood(entity_id, depth=1)¶

Get the neighborhood of an entity up to specified depth.

Parameters:

entity_id (str) – The entity to explore
depth (int) – Maximum depth to traverse

Returns:

Dictionary containing entity neighborhood information

Return type:

dict[str, Any]

async run(user_input)¶

Main execution method for the KG Generator Agent.

Parameters:: user_input (str)
Return type:: str

setup_agent()¶

Setup agent after initialization - Pydantic pattern.

Return type:: None

model_config¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class agents.memory.kg_generator_agent.KGGeneratorAgentConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for Knowledge Graph Generator Agent.

This configuration class defines all parameters needed to create and configure a KGGeneratorAgent, including LLM settings, extraction parameters, and entity/relationship type specifications.

Parameters:: data (Any)

name¶: Unique identifier for the agent instance

memory_store_manager¶: Manager for memory storage and retrieval operations

memory_classifier¶: Classifier for analyzing memory content and types

engine¶: LLM engine configuration for entity and relationship extraction

extract_batch_size¶: Number of memories to process in each batch

min_confidence_threshold¶: Minimum confidence score for extracted entities/relationships

enable_iterative_refinement¶: Whether to enable iterative graph refinement

entity_types¶: List of entity types the agent should extract

relationship_types¶: List of relationship types the agent should extract

Examples

Basic configuration:

config = KGGeneratorAgentConfig(
    name="my_kg_generator",
    memory_store_manager=store_manager,
    memory_classifier=classifier,
    engine=AugLLMConfig(temperature=0.1),
    extract_batch_size=20,
    min_confidence_threshold=0.7
)

Advanced configuration with custom types:

config = KGGeneratorAgentConfig(
    name="custom_kg_generator",
    memory_store_manager=store_manager,
    memory_classifier=classifier,
    engine=AugLLMConfig(
        model="gpt-4",
        temperature=0.1,
        max_tokens=2000
    ),
    extract_batch_size=15,
    min_confidence_threshold=0.8,
    entity_types=[
        "person", "organization", "technology", "concept",
        "project", "skill", "tool", "document"
    ],
    relationship_types=[
        "works_at", "knows", "uses", "creates", "manages",
        "collaborates_with", "depends_on", "teaches"
    ]
)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

model_config¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class agents.memory.kg_generator_agent.KnowledgeGraphNode(/, **data)¶

Bases: pydantic.BaseModel

Represents a node (entity) in the knowledge graph.

A knowledge graph node encapsulates an entity extracted from memory content, including its type, properties, and references to the memories that mention it.

Parameters:: data (Any)

id¶: Unique identifier for the node, typically generated from name and type

type¶: Entity type classification (person, organization, concept, etc.)

name¶: Human-readable display name of the entity

properties¶: Dictionary of additional properties and attributes

memory_references¶: List of memory IDs that reference this entity

confidence¶: Confidence score in entity existence and accuracy (0.0-1.0)

created_at¶: UTC timestamp when the entity was first created

last_updated¶: UTC timestamp of the last update to this entity

Examples

Creating a person entity:

person = KnowledgeGraphNode(
    id="person_alice_smith",
    type="person",
    name="Alice Smith",
    properties={"role": "engineer", "company": "TechCorp"},
    memory_references=["mem_123", "mem_456"],
    confidence=0.95
)

Creating a concept entity:

concept = KnowledgeGraphNode(
    id="concept_machine_learning",
    type="concept",
    name="Machine Learning",
    properties={"category": "technology", "complexity": "high"},
    confidence=0.9
)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

class agents.memory.kg_generator_agent.KnowledgeGraphRelationship(/, **data)¶

Bases: pydantic.BaseModel

Represents a relationship (edge) in the knowledge graph.

A knowledge graph relationship connects two entities with a typed relationship, capturing semantic connections discovered from memory content.

Parameters:: data (Any)

id¶: Unique identifier for the relationship, typically generated from source, target, and type

source_id¶: ID of the source entity in the relationship

target_id¶: ID of the target entity in the relationship

relationship_type¶: Type of relationship (works_at, knows, uses, etc.)

properties¶: Dictionary of additional relationship properties

memory_references¶: List of memory IDs that reference this relationship

confidence¶: Confidence score in relationship existence and accuracy (0.0-1.0)

created_at¶: UTC timestamp when the relationship was first created

last_updated¶: UTC timestamp of the last update to this relationship

Examples

Creating a work relationship:

work_rel = KnowledgeGraphRelationship(
    id="person_alice_smith_works_at_organization_techcorp",
    source_id="person_alice_smith",
    target_id="organization_techcorp",
    relationship_type="works_at",
    properties={"role": "senior_engineer", "start_date": "2020-01-15"},
    memory_references=["mem_123"],
    confidence=0.9
)

Creating a knowledge relationship:

knows_rel = KnowledgeGraphRelationship(
    id="person_alice_smith_knows_concept_machine_learning",
    source_id="person_alice_smith",
    target_id="concept_machine_learning",
    relationship_type="knows",
    properties={"proficiency": "expert", "years_experience": 5},
    confidence=0.85
)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

class agents.memory.kg_generator_agent.MemoryKnowledgeGraph(/, **data)¶

Bases: pydantic.BaseModel

Complete knowledge graph structure extracted from memory system.

The MemoryKnowledgeGraph represents a comprehensive knowledge graph built from memory content, containing entities (nodes) and their relationships (edges). It provides methods for managing the graph structure and querying entity neighborhoods.

Parameters:: data (Any)

nodes¶: Dictionary mapping node IDs to KnowledgeGraphNode objects

relationships¶: Dictionary mapping relationship IDs to KnowledgeGraphRelationship objects

metadata¶: Dictionary containing graph-level metadata and statistics

Examples

Creating and populating a knowledge graph:

graph = MemoryKnowledgeGraph()

# Add entities
person = KnowledgeGraphNode(
    id="person_alice", type="person", name="Alice Smith"
)
company = KnowledgeGraphNode(
    id="org_techcorp", type="organization", name="TechCorp"
)

graph.add_node(person)
graph.add_node(company)

# Add relationship
works_at = KnowledgeGraphRelationship(
    id="alice_works_at_techcorp",
    source_id="person_alice",
    target_id="org_techcorp",
    relationship_type="works_at"
)

graph.add_relationship(works_at)

Querying the graph:

# Get connected nodes
connected = graph.get_connected_nodes("person_alice")

# Get all relationships for a node
relationships = graph.get_relationships_for_node("person_alice")

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

add_node(node)¶

Add or update a node in the graph.

If the node already exists, its properties and memory references are merged with the existing node. Otherwise, a new node is added to the graph.

Parameters:: node (KnowledgeGraphNode) – The KnowledgeGraphNode to add or update
Return type:: None

Examples

Adding a new node:

graph = MemoryKnowledgeGraph()
node = KnowledgeGraphNode(
    id="person_alice", type="person", name="Alice Smith"
)
graph.add_node(node)

Updating an existing node:

# First addition
node1 = KnowledgeGraphNode(
    id="person_alice", type="person", name="Alice Smith",
    properties={"role": "engineer"}
)
graph.add_node(node1)

# Update with additional properties
node2 = KnowledgeGraphNode(
    id="person_alice", type="person", name="Alice Smith",
    properties={"company": "TechCorp"}
)
graph.add_node(node2)
# Result: node has both "role" and "company" properties

add_relationship(relationship)¶

Add or update a relationship in the graph.

If the relationship already exists, its properties and memory references are merged with the existing relationship. Otherwise, a new relationship is added to the graph.

Parameters:: relationship (KnowledgeGraphRelationship) – The KnowledgeGraphRelationship to add or update
Return type:: None

Examples

Adding a new relationship:

graph = MemoryKnowledgeGraph()
rel = KnowledgeGraphRelationship(
    id="alice_works_at_techcorp",
    source_id="person_alice",
    target_id="org_techcorp",
    relationship_type="works_at"
)
graph.add_relationship(rel)

Updating an existing relationship:

# First addition
rel1 = KnowledgeGraphRelationship(
    id="alice_works_at_techcorp",
    source_id="person_alice",
    target_id="org_techcorp",
    relationship_type="works_at",
    properties={"role": "engineer"}
)
graph.add_relationship(rel1)

# Update with additional properties
rel2 = KnowledgeGraphRelationship(
    id="alice_works_at_techcorp",
    source_id="person_alice",
    target_id="org_techcorp",
    relationship_type="works_at",
    properties={"start_date": "2020-01-15"}
)
graph.add_relationship(rel2)
# Result: relationship has both "role" and "start_date" properties

get_connected_nodes(node_id)¶

Get all nodes connected to a given node.

Traverses all relationships to find nodes that are directly connected to the specified node, regardless of relationship direction.

Parameters:: node_id (str) – The ID of the node to find connections for
Returns:: List of KnowledgeGraphNode objects connected to the specified node
Return type:: list[KnowledgeGraphNode]

Examples

Finding connected nodes:

graph = MemoryKnowledgeGraph()
# Assume graph has been populated with nodes and relationships

connected = graph.get_connected_nodes("person_alice")
for node in connected:
    print(f"Connected to: {node.name} ({node.type})")

get_relationships_for_node(node_id)¶

Get all relationships involving a given node.

Finds all relationships where the specified node is either the source or target of the relationship.

Parameters:: node_id (str) – The ID of the node to find relationships for
Returns:: List of KnowledgeGraphRelationship objects involving the specified node
Return type:: list[KnowledgeGraphRelationship]

Examples

Finding node relationships:

graph = MemoryKnowledgeGraph()
# Assume graph has been populated

relationships = graph.get_relationships_for_node("person_alice")
for rel in relationships:
    print(f"Relationship: {rel.relationship_type} "
          f"({rel.source_id} -> {rel.target_id})")