Skip to content

Context

The Context class in YosrAI provides a flexible and type-safe way to manage state and contextual information across agents and actions. It integrates with LangGraph's state management system and provides various methods for updating and accessing context data.

Class Overview

from yosrai.core import Context

# Create a context with initial values
context = Context(
    topic="AI Agents",
    user_name="John",
    messages=[]
)

Key Features

Dynamic Type Safety

The Context class automatically creates typed dictionaries based on your context structure, ensuring type safety while maintaining flexibility:

# Context adapts to new attributes
context.update(
    temperature=0.7,
    max_tokens=1000
)

Message Management

Built-in support for LangChain message types and automatic message handling:

# Messages are automatically handled with proper typing
context(messages=[
    HumanMessage(content="Hello"),
    AIMessage(content="Hi there!")
])

Context Functions

The ContextFunctions enum provides various methods for handling context updates:

from yosrai.core import ContextFunctions

# Available functions
DEFAULT = 'default'              # Store value as is
APPEND_MESSAGES = 'append_messages'  # Append to message list
LATEST = 'latest'               # Keep only latest values
JSON = 'json'                   # Parse and store as JSON
REPLIES = 'replies'             # Store only replies
APPEND_OUTPUT_CONTENT = 'append_output_content'  # Append output content

Usage Examples

Basic Context Management

# Create context
context = Context(
    system_prompt="You are a helpful assistant",
    temperature=0.7
)

# Update values
context.update(
    max_tokens=1000,
    model="gpt-4"
)

# Access context
current_state = context()

Using with Actions

from yosrai.core import Action, Context

# Create context for an action
action = Action("Customer Support")
action.Context(
    customer_id="12345",
    priority="high",
    messages=[]
)

Dynamic Updates

# Context automatically adapts to new fields
context.update(
    new_field="value",
    another_field=123
)

# Get typed dictionary for current state
state_type = context.get_type()

Integration with LangGraph

The Context class is designed to work seamlessly with LangGraph's state management:

from langgraph.graph import StateGraph

# Create state graph with context
graph = StateGraph(context.get_type())

Methods

init

def __init__(self, **kwargs)

Initialize context with optional key-value pairs.

update

def update(self, **kwargs)

Update context with new key-value pairs.

call

def __call__(self, **kwargs)

Get current context state, optionally updating it first.

get_type

def get_type(self)

Get the TypedDict representation of the context structure.

Best Practices

  1. Initialize Early: Set up your context structure early with all expected fields.
  2. Type Safety: Use the context's type system to catch errors early.
  3. Message Handling: Let the context handle message management automatically.
  4. State Management: Use context functions for different state management needs.

Advanced Usage

Custom Type Handling

# Create context with custom types
from typing import List, Dict

context = Context(
    custom_data=List[Dict[str, Any]],
    structured_data={"key": "value"}
)

Integration with Custom Classes

class CustomState:
    def __init__(self, data):
        self.data = data

# Use with context
context = Context(
    state=CustomState({"key": "value"})
)

Thread Safety

The Context class is designed to be thread-safe for concurrent access:

# Safe to use in async environments
async def update_context(context, data):
    context.update(**data)