haive.core.schema.state_schema

State schema base class for the Haive framework.

from typing import Any This module provides the StateSchema base class that extends Pydantic’s BaseModel with features specifically designed for AI agent state management and graph-based workflows. The StateSchema class adds powerful capabilities including field sharing between parent and child graphs, reducer functions for state updates, engine I/O tracking, and extensive serialization support.

StateSchema serves as the foundation of the Haive Schema System, enabling fully dynamic and serializable state schemas that can be composed, modified, and extended at runtime. This flexibility makes it ideal for complex agent architectures and nested workflows.

Key features include: - Field sharing: Share state between parent and child graphs with explicit control - Reducer functions: Define how field values should be combined during state updates - Engine I/O tracking: Map which fields are inputs and outputs for specific engines - Message handling: Built-in methods for working with message fields - Serialization: Comprehensive support for converting to/from dictionaries and JSON - State manipulation: Methods for updating, merging, and comparing states - Pretty printing: Rich visualization of state content - Engine integration: Prepare inputs and process outputs for specific engines

Examples

from typing import List from langchain_core.messages import BaseMessage, HumanMessage, AIMessage from pydantic import Field from haive.core.schema import StateSchema from langgraph.graph import add_messages

class ConversationState(StateSchema):

messages: List[BaseMessage] = Field(default_factory=list) query: str = Field(default=””) response: str = Field(default=””) context: List[str] = Field(default_factory=list)

# Define which fields should be shared with parent graphs __shared_fields__ = [“messages”]

# Define reducer functions for each field __reducer_fields__ = {

“messages”: add_messages, “context”: lambda a, b: (a or []) + (b or [])

}

# Define which fields are inputs/outputs for which engines __engine_io_mappings__ = {

“retriever”: {

“inputs”: [“query”], “outputs”: [“context”]

}, “llm”: {

“inputs”: [“query”, “context”, “messages”], “outputs”: [“response”]

}

}

Classes

EngineIOConfig	Configuration for engine input/output mappings.
StateConfig	Configuration for state schema metadata.
StateSchema	Enhanced base class for state schemas in the Haive framework.

Module Contents

class haive.core.schema.state_schema.EngineIOConfig

Bases: typing_extensions.TypedDict

Configuration for engine input/output mappings.

Initialize self. See help(type(self)) for accurate signature.

class haive.core.schema.state_schema.StateConfig

Bases: typing_extensions.TypedDict

Configuration for state schema metadata.

Initialize self. See help(type(self)) for accurate signature.

class haive.core.schema.state_schema.StateSchema(/, **data)

Bases: pydantic.BaseModel, Generic[TEngine, TEngines]

Enhanced base class for state schemas in the Haive framework.

StateSchema extends Pydantic’s BaseModel with features for AI agent state management and graph-based workflows. It serves as the core component of the Haive Schema System, providing extensive capabilities for state management in complex agent architectures.

Key Features:

• Field sharing: Control which fields are shared between parent and child graphs

• Reducer functions: Define how field values are combined during state updates

• Engine I/O tracking: Map which fields are inputs/outputs for which engines

• Message handling: Methods for working with conversation message fields

• Serialization: Convert states to/from dictionaries and JSON

• State manipulation: Update, merge, compare, and diff state objects

• Integration: Support for LangGraph and engine components

• Visualization: Rich display options for state inspection

Special Class Variables:

__shared_fields__ (List[str]): Fields to share with parent graphs __serializable_reducers__ (Dict[str, str]): Serializable reducer function names __engine_io_mappings__ (Dict[str, Dict[str, List[str]]]): Engine I/O mappings __input_fields__ (Dict[str, List[str]]): Input fields for each engine __output_fields__ (Dict[str, List[str]]): Output fields for each engine __structured_models__ (Dict[str, str]): Paths to structured output models __structured_model_fields__ (Dict[str, List[str]]): Fields for structured models __reducer_fields__ (Dict[str, Callable]): Runtime reducer functions (not stored)

Field sharing enables parent and child graphs to maintain synchronized state for specific fields, which is critical for nested graph execution. Reducer functions define how field values are combined during updates, enabling sophisticated state merging operations beyond simple assignment.

Examples

from typing import List from langchain_core.messages import BaseMessage from pydantic import Field from haive.core.schema import StateSchema

class MyState(StateSchema):

messages: List[BaseMessage] = Field(default_factory=list) query: str = Field(default=””) result: str = Field(default=””)

# Share only messages with parent graphs __shared_fields__ = [“messages”]

# Define reducer for messages __reducer_fields__ = {

“messages”: add_messages # From langgraph.graph

}

# Create state instance state = MyState()

# Add a message state.add_message(HumanMessage(content=”Hello”))

# Convert to dictionary state_dict = state.to_dict()

# Create from dictionary new_state = MyState.from_dict(state_dict)

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.

Parameters:

data (Any)

add_engine(name, engine)

Add an engine to the engines registry.

Parameters:

• name (str) – Name/key for the engine

• engine (haive.core.engine.base.Engine) – Engine instance to add

Return type:

None

add_message(message)

Add a single message to the messages field.

Parameters:

message (langchain_core.messages.BaseMessage) – BaseMessage to add

Returns:

Self for chaining

Return type:

StateSchema

add_messages(new_messages)

Add multiple messages to the messages field.

Parameters:

new_messages (list[langchain_core.messages.BaseMessage]) – List of messages to add

Returns:

Self for chaining

Return type:

StateSchema

apply_reducers(other)

Update state applying reducer functions where defined.

This method processes updates with special handling for fields that have reducer functions defined.

Parameters:

other (dict[str, Any] | StateSchema) – Dictionary or StateSchema with update values

Returns:

Self for chaining

Return type:

StateSchema

classmethod as_table()

Create a rich table representation of the schema.

Returns:

Rich Table object

Return type:

rich.table.Table

clear_messages()

Clear all messages in the messages field.

Returns:

Self for chaining

Return type:

StateSchema

combine_with(other)

Combine this state with another, applying reducers for shared fields.

This is more sophisticated than update() or apply_reducers() as it properly handles StateSchema-specific metadata and shared fields.

Parameters:

other (StateSchema | dict[str, Any]) – Other state to combine with

Returns:

New combined state instance

Return type:

StateSchema

classmethod compare_with(other, title=None)

Compare this schema with another in a side-by-side display.

Parameters:

• other (type[StateSchema]) – Other schema to compare with

• title (str | None) – Optional title for the comparison

Return type:

None

copy(**updates)

Create a copy of this state, optionally with updates.

Parameters:

**updates – Field values to update in the copy

Returns:

New StateSchema instance

Return type:

StateSchema

classmethod create_input_schema(engine_name=None, name=None)

Alias for derive_input_schema for backward compatibility.

Parameters:

• engine_name (str | None) – Optional name of the engine to target

• name (str | None) – Optional name for the schema class

Returns:

A BaseModel subclass for input validation

Return type:

type[pydantic.BaseModel]

classmethod create_output_schema(engine_name=None, name=None)

Alias for derive_output_schema for backward compatibility.

Parameters:

• engine_name (str | None) – Optional name of the engine to target

• name (str | None) – Optional name for the schema class

Returns:

A BaseModel subclass for output validation

Return type:

type[pydantic.BaseModel]

deep_copy()

Create a deep copy of this state object.

Returns:

New StateSchema instance with deep-copied values

Return type:

StateSchema

classmethod derive_input_schema(engine_name=None, name=None)

Derive an input schema for the given engine from this state schema.

This method intelligently selects the appropriate base class for the derived schema, using prebuilt states (MessagesState, ToolState) when appropriate instead of just creating a generic BaseModel.

Parameters:

• engine_name (str | None) – Optional name of the engine to target (default: all inputs)

• name (str | None) – Optional name for the schema class

Returns:

A BaseModel subclass for input validation, potentially inheriting from MessagesState or ToolState for better compatibility

Return type:

type[pydantic.BaseModel]

classmethod derive_output_schema(engine_name=None, name=None)

Derive an output schema for the given engine from this state schema.

This method intelligently selects the appropriate base class for the derived schema, using prebuilt states (MessagesState, ToolState) when appropriate instead of just creating a generic BaseModel.

Parameters:

• engine_name (str | None) – Optional name of the engine to target (default: all outputs)

• name (str | None) – Optional name for the schema class

Returns:

A BaseModel subclass for output validation, potentially inheriting from MessagesState or ToolState for better compatibility

Return type:

type[pydantic.BaseModel]

dict(**kwargs)

Backwards compatibility alias for model_dump.

Parameters:

**kwargs – Keyword arguments for model_dump

Returns:

Dictionary representation of the state

Return type:

dict[str, Any]

differences_from(other)

Compare this state with another and return differences.

Parameters:

other (StateSchema | dict[str, Any]) – Other state to compare with

Returns:

Dictionary mapping field names to (self_value, other_value) tuples

Return type:

dict[str, tuple[Any, Any]]

classmethod display_code(title=None)

Display Python code representation of the schema.

Parameters:

title (str | None) – Optional title for the display

Return type:

None

classmethod display_schema(title=None)

Display schema information in a rich format.

Parameters:

title (str | None) – Optional title for the display

Return type:

None

classmethod display_table()

Display schema as a table.

Return type:

None

classmethod extract_values(state, keys=None)

Class method to extract values from a state object or dictionary.

Parameters:

• state (StateSchema | dict[str, Any]) – State object or dictionary to extract values from

• keys (list[str] | dict[str, str] | None) – Can be: - List[str]: List of field names to extract - Dict[str, str]: Mapping of output keys to state field names - None: Extract all fields

Returns:

Dictionary containing the requested values

Return type:

dict[str, Any]

classmethod from_dict(data)

Create a state from a dictionary.

Parameters:

data (FieldMapping) – Dictionary with field values

Returns:

New StateSchema instance

Return type:

Self

classmethod from_json(json_str)

Create state from JSON string.

Parameters:

json_str (str) – JSON string to parse

Returns:

New StateSchema instance

Return type:

StateSchema

classmethod from_partial_dict(data)

Create a state from a partial dictionary, filling in defaults.

Parameters:

data (dict[str, Any]) – Partial dictionary with field values

Returns:

New StateSchema instance with defaults applied

Return type:

StateSchema

classmethod from_runnable_config(config)

Extract state from a RunnableConfig.

Parameters:

config (langchain_core.runnables.RunnableConfig) – RunnableConfig to extract from

Returns:

StateSchema instance or None if no state found

Return type:

StateSchema | None

classmethod from_snapshot(snapshot)

Create a state from a LangGraph StateSnapshot.

Parameters:

snapshot (Any) – StateSnapshot from LangGraph

Returns:

New StateSchema instance

Return type:

StateSchema

get(key, default=None)

Safely get a field value with a default.

Parameters:

• key (str) – Field name to get

• default (Any) – Default value if field doesn’t exist

Returns:

Field value or default

Return type:

Any

classmethod get_all_class_engines()

Get all class-level engines.

Returns:

Dictionary of all engines

Return type:

dict[str, Any]

get_all_instance_engines()

Get all engines from both instance and class level.

Returns:

Dictionary mapping engine names to engine instances

Return type:

dict[str, Any]

classmethod get_class_engine(name)

Get a class-level engine by name.

Parameters:

name (str) – Name of the engine to retrieve

Returns:

Engine instance if found, None otherwise

Return type:

Any | None

get_engine(name)

Get an engine by name.

Parameters:

name (str) – Name of the engine to retrieve

Returns:

Engine instance if found, None otherwise

Return type:

haive.core.engine.base.Engine | None

get_engines()

Get all engines in this state.

Returns:

Dictionary mapping engine names to engine instances

Return type:

dict[str, Any]

get_instance_engine(name)

Get an engine from instance or class level.

Parameters:

name (str) – Name of the engine to retrieve

Returns:

Engine instance if found, None otherwise

Return type:

Any | None

get_last_message()

Get the last message in the messages field.

Returns:

Last message or None if no messages exist

Return type:

langchain_core.messages.BaseMessage | None

get_state_values(keys=None)

Extract specified state values into a dictionary.

Parameters:

keys (list[str] | dict[str, str] | None) – Can be: - List[str]: List of field names to extract - Dict[str, str]: Mapping of output keys to state field names - None: Extract all fields

Returns:

Dictionary containing the requested state values

Return type:

dict[str, Any]

classmethod get_structured_model(model_name)

Get a structured output model class by name.

Parameters:

model_name (str) – Name of the structured model

Returns:

Model class if found, None otherwise

Return type:

type[pydantic.BaseModel] | None

has_engine(name)

Check if an engine exists.

Parameters:

name (str) – Name of the engine to check

Returns:

True if engine exists, False otherwise

Return type:

bool

classmethod is_shared(field_name)

Check if a field is shared with parent graphs.

Parameters:

field_name (str) – Field name to check

Returns:

True if field is shared, False otherwise

Return type:

bool

list_engines()

Get list of all engine names.

Returns:

List of engine names

Return type:

list[str]

classmethod list_structured_models()

List all structured output models in this schema.

Returns:

List of structured model names

Return type:

list[str]

classmethod manager()

Get a manager for this schema (shorthand for to_manager()).

Returns:

StateSchemaManager instance

Return type:

haive.core.schema.schema_manager.StateSchemaManager

merge_engine_output(engine_name, output, apply_reducers=True)

Merge output from an engine into this state.

Parameters:

• engine_name (str) – Name of the engine

• output (dict[str, Any]) – Output data from the engine

• apply_reducers (bool) – Whether to apply reducers during merge

Returns:

Self for chaining

Return type:

StateSchema

merge_messages(new_messages)

Merge new messages with existing messages using appropriate reducer.

Parameters:

new_messages (list[langchain_core.messages.BaseMessage]) – New messages to add

Returns:

Self for chaining

Return type:

StateSchema

model_dump(**kwargs)

Override model_dump to exclude internal fields and handle special types.

Parameters:

**kwargs (Any) – Keyword arguments for model_dump

Returns:

Dictionary representation of the state

Return type:

FieldMapping

model_post_init(__context)

Sync engines from class level to instance level after initialization.

This ensures that engines stored at the class level (via SchemaComposer) are available on state instances.

Parameters:

__context (Any)

Return type:

None

patch(update_data, apply_reducers=True)

Update specific fields in the state.

Parameters:

• update_data (dict[str, Any]) – Dictionary of field updates

• apply_reducers (bool) – Whether to apply reducer functions

Returns:

Self for chaining

Return type:

StateSchema

prepare_for_engine(engine_name)

Prepare state data for a specific engine.

Extracts only fields that are inputs for the specified engine.

Parameters:

engine_name (str) – Name of the engine to prepare for

Returns:

Dictionary with engine-specific inputs

Return type:

dict[str, Any]

pretty_print(title=None)

Print state with rich formatting for easy inspection.

Parameters:

title (str | None) – Optional title for the display

Return type:

None

remove_engine(name)

Remove an engine from the registry.

Parameters:

name (str) – Name of the engine to remove

Returns:

True if engine was removed, False if not found

Return type:

bool

setup_engines_and_tools()

Setup engines and sync their tools, structured output models, and add engine to state.

This validator runs after the model is created and: 1. Finds all engine fields in the state 2. Syncs engine to main engine field and engines dict 3. Syncs tools from engine to state tools field 4. Syncs structured output models 5. Sets up parent-child relationships for nested state schemas

Return type:

Self

classmethod shared_fields()

Get the list of fields shared with parent graphs.

Returns:

List of shared field names

Return type:

list[str]

sync_engine_fields()

Sync between engine and engines dict for backward compatibility.

This validator ensures that: 1. If ‘engine’ is set, it’s available in engines dict 2. If engines dict has items but no engine, set main engine 3. Both access patterns work seamlessly

Return type:

Self

to_command(goto=None, graph=None)

Convert state to a Command object for LangGraph control flow.

Parameters:

• goto (str | None) – Optional next node to go to

• graph (str | None) – Optional graph to target (None for current, PARENT for parent)

Returns:

Command object with state update

Return type:

Any

to_dict()

Convert the state to a clean dictionary.

Returns:

Dictionary representation of the state

Return type:

FieldMapping

to_json()

Convert state to JSON string.

Returns:

JSON string representation of the state

Return type:

str

classmethod to_manager(name=None)

Convert schema class to a StateSchemaManager for further manipulation.

Parameters:

name (str | None) – Optional name for the resulting manager

Returns:

StateSchemaManager instance

Return type:

haive.core.schema.schema_manager.StateSchemaManager

classmethod to_python_code()

Convert schema to Python code representation.

Returns:

String containing Python code representation

Return type:

str

to_runnable_config(thread_id=None, **kwargs)

Convert state to a RunnableConfig.

Parameters:

• thread_id (str | None) – Optional thread ID for the configuration

• **kwargs – Additional configuration parameters

Returns:

RunnableConfig containing state data

Return type:

langchain_core.runnables.RunnableConfig

update(other)

Update the state with values from another state or dictionary.

This method performs a simple update without applying reducers.

Parameters:

other (dict[str, Any] | StateSchema) – Dictionary or StateSchema with update values

Returns:

Self for chaining

Return type:

StateSchema

classmethod validate_engine(v)

Handle both serialized dict and actual Engine instances.

This validator allows the engine field to accept both: - Actual Engine instances (for runtime use) - Serialized dicts (for state passing between agents)

This prevents the “Can’t instantiate abstract class Engine” error when deserializing state in multi-agent systems.

Return type:

Any

classmethod validate_engines(v)

Handle both serialized dicts and actual Engine instances in engines dict.

Similar to validate_engine but for the engines dictionary. Each value can be either a serialized dict or an actual Engine instance.

Return type:

Any

classmethod with_shared_fields(fields)

Create a copy of this schema with specified shared fields.

Parameters:

fields (list[str]) – List of field names to be marked as shared

Returns:

New StateSchema subclass with updated shared fields

Return type:

type[StateSchema]

property llm: haive.core.engine.base.Engine | None

Convenience property to access the LLM engine.

Return type:

haive.core.engine.base.Engine | None

property main_engine: haive.core.engine.base.Engine | None

Convenience property to access the main engine.

Return type:

haive.core.engine.base.Engine | None