Release v3.10.4

api.md

@@ -127,7 +127,7 @@ Methods:
 
 Types:
 ```python
-from praisonaiagents import Task
+from praisonaiagents import EmbeddingResult, Task, aembedding, aembeddings, embedding, embeddings, get_dimensions
 ```
 
 Methods:
@@ -136,6 +136,9 @@ Methods:
 * <code title="class Task">Task.<a href="./src/praisonai-agents/praisonaiagents/task/task.py">execute_callback_sync</a>(task_output: TaskOutput) -> None</code>
 * <code title="class Task">Task.<a href="./src/praisonai-agents/praisonaiagents/task/task.py">initialize_memory</a>()</code>
 * <code title="class Task">Task.<a href="./src/praisonai-agents/praisonaiagents/task/task.py">store_in_memory</a>(content: str, agent_name: str = None, task_id: str = None)</code>
+* <code title="function">praisonaiagents.<a href="./src/praisonai-agents/praisonaiagents/embedding/embed.py">aembedding</a>(input: Union[str, List[str]], model: str = 'text-embedding-3-small', dimensions: Optional[int] = None, encoding_format: str = 'float', timeout: float = 600.0, api_key: Optional[str] = None, api_base: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> EmbeddingResult</code>
+* <code title="function">praisonaiagents.<a href="./src/praisonai-agents/praisonaiagents/embedding/embed.py">embedding</a>(input: Union[str, List[str]], model: str = 'text-embedding-3-small', dimensions: Optional[int] = None, encoding_format: str = 'float', timeout: float = 600.0, api_key: Optional[str] = None, api_base: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> EmbeddingResult</code>
+* <code title="function">praisonaiagents.<a href="./src/praisonai-agents/praisonaiagents/embedding/dimensions.py">get_dimensions</a>(model_name: str) -> int</code>
 
 # Wrapper (praisonai)
 
@@ -246,8 +249,10 @@ Methods:
 * <code title="cli">praisonai lsp status <a href="./src/praisonai/praisonai/cli/commands/lsp.py">--help</a></code>
 * <code title="cli">praisonai lsp stop <a href="./src/praisonai/praisonai/cli/commands/lsp.py">--help</a></code>
 * <code title="cli">praisonai mcp add <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
+* <code title="cli">praisonai mcp auth <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
 * <code title="cli">praisonai mcp describe <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
 * <code title="cli">praisonai mcp list <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
+* <code title="cli">praisonai mcp logout <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
 * <code title="cli">praisonai mcp mcp-callback <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
 * <code title="cli">praisonai mcp remove <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
 * <code title="cli">praisonai mcp run <a href="./src/praisonai/praisonai/cli/commands/mcp.py">--help</a></code>
@@ -281,6 +286,12 @@ Methods:
 * <code title="cli">praisonai recipe run <a href="./src/praisonai/praisonai/cli/commands/recipe.py">--help</a></code>
 * <code title="cli">praisonai registry list <a href="./src/praisonai/praisonai/cli/commands/registry.py">--help</a></code>
 * <code title="cli">praisonai registry serve <a href="./src/praisonai/praisonai/cli/commands/registry.py">--help</a></code>
+* <code title="cli">praisonai replay cleanup <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
+* <code title="cli">praisonai replay context <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
+* <code title="cli">praisonai replay delete <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
+* <code title="cli">praisonai replay flow <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
+* <code title="cli">praisonai replay list <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
+* <code title="cli">praisonai replay show <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
 * <code title="cli">praisonai research research-main <a href="./src/praisonai/praisonai/cli/commands/research.py">--help</a></code>
 * <code title="cli">praisonai retrieval index <a href="./src/praisonai/praisonai/cli/commands/retrieval.py">--help</a></code>
 * <code title="cli">praisonai retrieval query <a href="./src/praisonai/praisonai/cli/commands/retrieval.py">--help</a></code>
@@ -334,7 +345,11 @@ Methods:
 * <code title="cli">praisonai traces enable <a href="./src/praisonai/praisonai/cli/commands/traces.py">--help</a></code>
 * <code title="cli">praisonai traces list <a href="./src/praisonai/praisonai/cli/commands/traces.py">--help</a></code>
 * <code title="cli">praisonai traces status <a href="./src/praisonai/praisonai/cli/commands/traces.py">--help</a></code>
-* <code title="cli">praisonai train train-main <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
+* <code title="cli">praisonai train agents <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
+* <code title="cli">praisonai train list <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
+* <code title="cli">praisonai train llm <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
+* <code title="cli">praisonai train show <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
+* <code title="cli">praisonai train train-callback <a href="./src/praisonai/praisonai/cli/commands/train.py">--help</a></code>
 * <code title="cli">praisonai ui chat <a href="./src/praisonai/praisonai/cli/commands/ui.py">--help</a></code>
 * <code title="cli">praisonai ui code <a href="./src/praisonai/praisonai/cli/commands/ui.py">--help</a></code>
 * <code title="cli">praisonai ui gradio <a href="./src/praisonai/praisonai/cli/commands/ui.py">--help</a></code>

docker/Dockerfile.chat

@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
 # Install Python packages (using latest versions)
 RUN pip install --no-cache-dir \
     praisonai_tools \
-    "praisonai>=3.9.13" \
+    "praisonai>=3.10.4" \
     "praisonai[chat]" \
     "embedchain[github,youtube]"
 

docker/Dockerfile.dev

@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
 # Install Python packages (using latest versions)
 RUN pip install --no-cache-dir \
     praisonai_tools \
-    "praisonai>=3.9.13" \
+    "praisonai>=3.10.4" \
     "praisonai[ui]" \
     "praisonai[chat]" \
     "praisonai[realtime]" \

docker/Dockerfile.ui

@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
 # Install Python packages (using latest versions)
 RUN pip install --no-cache-dir \
     praisonai_tools \
-    "praisonai>=3.9.13" \
+    "praisonai>=3.10.4" \
     "praisonai[ui]" \
     "praisonai[crewai]"
 

examples/capabilities/embeddings_example.py

@@ -2,12 +2,14 @@
 Embeddings Capability Example
 
 Demonstrates text embedding generation using PraisonAI capabilities.
+Both `embed` and `embedding` work identically - use whichever you prefer.
 """
 
-from praisonai.capabilities import embed
+from praisonai import embed, embedding  # Both work from top-level
+# Or: from praisonai.capabilities import embed, embedding
 
-# Single text embedding
-print("=== Single Text Embedding ===")
+# Single text embedding using embed()
+print("=== Single Text Embedding (using embed) ===")
 result = embed(
     input="Hello, world!",
     model="text-embedding-3-small"
@@ -16,6 +18,14 @@ print(f"Embedding dimensions: {len(result.embeddings[0])}")
 print(f"First 5 values: {result.embeddings[0][:5]}")
 print(f"Usage: {result.usage}")
 
+# Same thing using embedding() alias
+print("\n=== Single Text Embedding (using embedding alias) ===")
+result = embedding(
+    input="Hello, world!",
+    model="text-embedding-3-small"
+)
+print(f"Embedding dimensions: {len(result.embeddings[0])}")
+
 # Multiple text embeddings
 print("\n=== Multiple Text Embeddings ===")
 result = embed(
@@ -25,3 +35,12 @@ result = embed(
 print(f"Number of embeddings: {len(result.embeddings)}")
 for i, emb in enumerate(result.embeddings):
     print(f"  Text {i+1}: {len(emb)} dimensions")
+
+# With custom dimensions (for models that support it)
+print("\n=== Custom Dimensions ===")
+result = embed(
+    input="Hello world",
+    model="text-embedding-3-large",
+    dimensions=256
+)
+print(f"Reduced dimensions: {len(result.embeddings[0])}")

examples/context/01_basic_context.py

@@ -0,0 +1,59 @@
+"""
+Basic Context Management Example
+
+Demonstrates enabling context management with ContextConfig.
+This manages token budgets, auto-compaction, and prevents context overflow.
+"""
+
+from praisonaiagents import Agent
+from praisonaiagents.context import ContextConfig, OptimizerStrategy
+
+# Example 1: Enable context with defaults
+agent_basic = Agent(
+    name="Assistant",
+    instructions="You are a helpful assistant",
+    context=True  # Auto-compact at 80% utilization
+)
+
+# Example 2: Custom configuration
+agent_custom = Agent(
+    name="CustomAgent",
+    instructions="You are a helpful assistant with custom context settings",
+    context=ContextConfig(
+        auto_compact=True,
+        compact_threshold=0.8,      # Trigger at 80% utilization
+        strategy=OptimizerStrategy.SMART,
+        output_reserve=8000,        # Reserve tokens for output
+        keep_recent_turns=5,        # Always keep last 5 turns
+        tool_output_max=10000,      # Max tokens per tool output
+    )
+)
+
+# Example 3: Sliding window strategy
+agent_sliding = Agent(
+    name="SlidingAgent",
+    instructions="You are a helpful assistant using sliding window",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.SLIDING_WINDOW,
+        keep_recent_turns=10,       # Keep last 10 turns
+    )
+)
+
+if __name__ == "__main__":
+    print("=== Basic Context Example ===")
+    print("Context management is automatically enabled.")
+    print()
+    
+    # Test basic agent
+    result = agent_basic.chat("Hello! What is 2 + 2?")
+    print(f"Response: {result}")
+    print()
+    
+    # Show context config
+    print("=== Context Configuration ===")
+    config = ContextConfig(
+        auto_compact=True,
+        compact_threshold=0.8,
+        strategy=OptimizerStrategy.SMART,
+    )
+    print(f"Config: {config.to_dict()}")

examples/context/02_session_tracking.py

@@ -0,0 +1,71 @@
+"""
+Session Tracking Example
+
+Demonstrates tracking conversation state (goal, plan, progress) across turns.
+Inspired by Agno's SessionContextStore pattern.
+"""
+
+from praisonaiagents.context import SessionContextTracker
+
+# Create session tracker
+tracker = SessionContextTracker(
+    session_id="user123",
+    track_summary=True,
+    track_goal=True,
+    track_plan=True,
+    track_progress=True,
+)
+
+# Set the user's goal
+tracker.update_goal("Build a Python web application with Flask")
+
+# Define the plan
+tracker.update_plan([
+    "Create Flask project structure",
+    "Add routes for API endpoints",
+    "Implement database models",
+    "Add authentication",
+    "Write tests",
+    "Deploy to cloud"
+])
+
+# Mark completed steps
+tracker.add_progress("Created Flask project structure")
+tracker.add_progress("Added routes for API endpoints")
+
+# Update summary
+tracker.update_summary(
+    "User is building a Flask web app. Completed project setup and API routes. "
+    "Next: implement database models."
+)
+
+if __name__ == "__main__":
+    print("=== Session Tracking Example ===")
+    print()
+    
+    # Show current state
+    print(f"Session ID: {tracker.session_id}")
+    print(f"Goal: {tracker.goal}")
+    print(f"Plan: {tracker.plan}")
+    print(f"Progress: {tracker.progress}")
+    print()
+    
+    # Get context string (simple format)
+    print("=== Context String ===")
+    print(tracker.to_context_string())
+    print()
+    
+    # Get system prompt section (XML format with guidelines)
+    print("=== System Prompt Section ===")
+    print(tracker.to_system_prompt_section())
+    print()
+    
+    # Export to dict for persistence
+    print("=== Export as Dict ===")
+    data = tracker.to_dict()
+    print(f"Exported: {data.keys()}")
+    print()
+    
+    # Recreate from dict
+    new_tracker = SessionContextTracker.from_dict(data)
+    print(f"Restored goal: {new_tracker.goal}")

examples/context/03_multi_memory.py

@@ -0,0 +1,81 @@
+"""
+Multi-Memory Aggregation Example
+
+Demonstrates concurrent fetching from multiple context sources.
+Inspired by CrewAI's ContextualMemory pattern.
+"""
+
+import asyncio
+from praisonaiagents.context import ContextAggregator, create_aggregator_from_config
+
+# Example 1: Basic aggregator with mock sources
+def mock_memory_search(query: str) -> str:
+    """Mock memory search - returns recent conversation context."""
+    return f"[Memory] User previously asked about: {query}. They prefer Python."
+
+def mock_knowledge_search(query: str) -> str:
+    """Mock knowledge search - returns facts from knowledge base."""
+    return f"[Knowledge] Relevant facts: Python is a high-level programming language."
+
+def mock_rag_retrieve(query: str) -> str:
+    """Mock RAG retrieval - returns document chunks."""
+    return f"[RAG] Retrieved: Documentation about {query} from the codebase."
+
+# Create aggregator
+aggregator = ContextAggregator(
+    max_tokens=4000,
+    separator="\n\n---\n\n",
+    include_source_labels=True
+)
+
+# Register sources with priorities (lower = higher priority)
+aggregator.register_source("memory", mock_memory_search, priority=10)
+aggregator.register_source("knowledge", mock_knowledge_search, priority=20)
+aggregator.register_source("rag", mock_rag_retrieve, priority=30)
+
+async def demo_async():
+    """Demo async aggregation."""
+    print("=== Async Aggregation ===")
+    result = await aggregator.aggregate("Python web development")
+    print(f"Sources used: {result.sources_used}")
+    print(f"Tokens used: {result.tokens_used}")
+    print(f"Fetch times: {result.fetch_times}")
+    print()
+    print("Context:")
+    print(result.context)
+    return result
+
+def demo_sync():
+    """Demo sync aggregation."""
+    print("=== Sync Aggregation ===")
+    result = aggregator.aggregate_sync("Flask API development")
+    print(f"Sources used: {result.sources_used}")
+    print(f"Tokens used: {result.tokens_used}")
+    print()
+    print("Context:")
+    print(result.context)
+    return result
+
+if __name__ == "__main__":
+    print("=== Multi-Memory Aggregation Example ===")
+    print()
+    
+    # Show registered sources
+    print(f"Registered sources: {aggregator.sources}")
+    print()
+    
+    # Demo sync aggregation
+    demo_sync()
+    print()
+    
+    # Demo async aggregation
+    asyncio.run(demo_async())
+    print()
+    
+    # Demo selective sources
+    print("=== Selective Sources ===")
+    result = aggregator.aggregate_sync(
+        "Find database info",
+        sources=["memory", "knowledge"]  # Only use these
+    )
+    print(f"Sources used: {result.sources_used}")

examples/context/04_compaction_strategies.py

@@ -0,0 +1,76 @@
+"""
+Compaction Strategies Example
+
+Demonstrates different context compaction strategies.
+"""
+
+from praisonaiagents import Agent
+from praisonaiagents.context import ContextConfig, OptimizerStrategy
+
+# Strategy 1: TRUNCATE - Fast, removes oldest messages
+agent_truncate = Agent(
+    name="TruncateAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.TRUNCATE,
+        compact_threshold=0.8,
+    )
+)
+
+# Strategy 2: SLIDING_WINDOW - Keeps recent N turns
+agent_sliding = Agent(
+    name="SlidingAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.SLIDING_WINDOW,
+        keep_recent_turns=10,
+    )
+)
+
+# Strategy 3: SMART - Importance-based (default)
+agent_smart = Agent(
+    name="SmartAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.SMART,  # Default
+        compact_threshold=0.8,
+    )
+)
+
+# Strategy 4: PRUNE_TOOLS - Truncate old tool outputs
+agent_prune = Agent(
+    name="PruneAgent",
+    instructions="You are a helpful assistant with tools",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.PRUNE_TOOLS,
+        tool_output_max=5000,  # Max tokens per tool output
+    )
+)
+
+# Strategy 5: NON_DESTRUCTIVE - Only prune tools, keep all history
+agent_safe = Agent(
+    name="SafeAgent",
+    instructions="You are a helpful assistant (safety-critical)",
+    context=ContextConfig(
+        strategy=OptimizerStrategy.NON_DESTRUCTIVE,
+    )
+)
+
+if __name__ == "__main__":
+    print("=== Compaction Strategies Example ===")
+    print()
+    
+    strategies = [
+        ("TRUNCATE", OptimizerStrategy.TRUNCATE, "Remove oldest messages"),
+        ("SLIDING_WINDOW", OptimizerStrategy.SLIDING_WINDOW, "Keep last N turns"),
+        ("SMART", OptimizerStrategy.SMART, "Importance-based (default)"),
+        ("PRUNE_TOOLS", OptimizerStrategy.PRUNE_TOOLS, "Truncate tool outputs"),
+        ("NON_DESTRUCTIVE", OptimizerStrategy.NON_DESTRUCTIVE, "Only prune tools"),
+    ]
+    
+    for name, strategy, description in strategies:
+        print(f"  {name:20} - {description}")
+    
+    print()
+    print("Usage:")
+    print("  context=ContextConfig(strategy=OptimizerStrategy.SLIDING_WINDOW)")

examples/context/05_monitoring.py

@@ -0,0 +1,85 @@
+"""
+Context Monitoring Example
+
+Demonstrates monitoring context usage with logging and metrics.
+"""
+
+from praisonaiagents import Agent
+from praisonaiagents.context import ContextConfig, MonitorConfig, OptimizerStrategy
+
+# Example 1: Human-readable monitoring
+agent_human = Agent(
+    name="MonitoredAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        auto_compact=True,
+        monitor=MonitorConfig(
+            enabled=True,
+            path="./context_logs/",
+            format="human",           # Human-readable format
+            frequency="turn",         # Log on every turn
+            redact_sensitive=True,    # Redact sensitive data
+            multi_agent_files=True,   # Separate files per agent
+        )
+    )
+)
+
+# Example 2: JSON monitoring (for parsing)
+agent_json = Agent(
+    name="JSONMonitorAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        monitor=MonitorConfig(
+            enabled=True,
+            path="./context_logs/",
+            format="json",            # JSON format for parsing
+            frequency="tool_call",    # Log on tool calls
+        )
+    )
+)
+
+# Example 3: Overflow-only monitoring
+agent_overflow = Agent(
+    name="OverflowAgent",
+    instructions="You are a helpful assistant",
+    context=ContextConfig(
+        monitor=MonitorConfig(
+            enabled=True,
+            path="./context_logs/",
+            format="human",
+            frequency="overflow",     # Only log on overflow
+        )
+    )
+)
+
+if __name__ == "__main__":
+    print("=== Context Monitoring Example ===")
+    print()
+    
+    # Show monitor configuration
+    config = MonitorConfig(
+        enabled=True,
+        path="./context_logs/",
+        format="human",
+        frequency="turn",
+        redact_sensitive=True,
+    )
+    
+    print("Monitor Config:")
+    print(f"  enabled: {config.enabled}")
+    print(f"  path: {config.path}")
+    print(f"  format: {config.format}")
+    print(f"  frequency: {config.frequency}")
+    print(f"  redact_sensitive: {config.redact_sensitive}")
+    print()
+    
+    print("Frequency options:")
+    print("  'turn'      - Log on every conversation turn")
+    print("  'tool_call' - Log when tools are called")
+    print("  'manual'    - Manual logging only")
+    print("  'overflow'  - Log only when context overflows")
+    print()
+    
+    print("Format options:")
+    print("  'human' - Human-readable format")
+    print("  'json'  - JSON format for parsing")

examples/context/15_recursive_context.py

@@ -0,0 +1,274 @@
+"""
+Example 15: Recursive Context Pattern
+
+Demonstrates the RLM (Recursive Language Model) pattern for handling large context:
+- Store large documents as artifacts (not in LLM context)
+- Agent uses tools (grep/head/tail/chunk) to explore artifacts
+- Dramatically reduces token usage (~2K vs ~10K+)
+
+This pattern solves "context rot" by treating context as programmatically
+explorable data instead of dumping everything into the LLM's context window.
+"""
+
+import tempfile
+import os
+from praisonaiagents import Agent
+from praisonai.context import FileSystemArtifactStore
+from praisonai.context.queue import create_artifact_tools
+from praisonaiagents.context.artifacts import ArtifactMetadata
+
+
+def create_sample_document():
+    """Create a sample large document for testing."""
+    # Simulate a research paper / large document
+    content = """
+# Attention Is All You Need
+
+## Abstract
+
+The dominant sequence transduction models are based on complex recurrent or 
+convolutional neural networks that include an encoder and a decoder. The best
+performing models also connect the encoder and decoder through an attention
+mechanism. We propose a new simple network architecture, the Transformer,
+based solely on attention mechanisms, dispensing with recurrence and convolutions
+entirely. 
+
+## Introduction
+
+Recurrent neural networks, long short-term memory and gated recurrent neural 
+networks in particular, have been firmly established as state of the art 
+approaches in sequence modeling and transduction problems such as language 
+modeling and machine translation. 
+
+## Background
+
+The goal of reducing sequential computation also forms the foundation of the 
+Extended Neural GPU, ByteNet and ConvS2S, all of which use convolutional neural
+networks as basic building block.
+
+## Model Architecture
+
+### Encoder and Decoder Stacks
+
+The encoder is composed of a stack of N = 6 identical layers. Each layer has 
+two sub-layers. The first is a multi-head self-attention mechanism, and the 
+second is a simple, position-wise fully connected feed-forward network.
+
+### Attention
+
+An attention function can be described as mapping a query and a set of key-value
+pairs to an output, where the query, keys, values, and output are all vectors.
+
+#### Scaled Dot-Product Attention
+
+We call our particular attention "Scaled Dot-Product Attention". The input 
+consists of queries and keys of dimension dk, and values of dimension dv.
+
+The formula is: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
+
+### Multi-Head Attention
+
+Multi-head attention allows the model to jointly attend to information from
+different representation subspaces at different positions.
+
+## Results
+
+### Machine Translation
+
+On the WMT 2014 English-to-German translation task, the big transformer model
+outperforms the best previously reported models including ensembles.
+
+The model achieved a BLEU score of 28.4 on the WMT 2014 English-to-German 
+translation task, improving over the existing best results by over 2 BLEU.
+
+On the WMT 2014 English-to-French translation task, our big model achieves 
+a BLEU score of 41.8, outperforming all published single models.
+
+### Training
+
+Training took 3.5 days on 8 P100 GPUs. The big model was trained for 300,000
+steps (3.5 days) at a cost of approximately $1000 in cloud compute.
+
+## Model Variations
+
+| Model | Parameters | BLEU EN-DE | BLEU EN-FR |
+|-------|------------|------------|------------|
+| Base  | 65M        | 27.3       | 38.1       |
+| Big   | 213M       | 28.4       | 41.8       |
+
+The base model has 65 million parameters (65 × 10^6).
+The big model has 213 million parameters.
+
+## Conclusion
+
+In this work, we presented the Transformer, the first sequence transduction
+model based entirely on attention, replacing the recurrent layers most commonly
+used in encoder-decoder architectures with multi-headed self-attention.
+
+The Transformer can be trained significantly faster than architectures based
+on recurrent or convolutional layers.
+
+## Authors
+
+Ashish Vaswani (Google Brain)
+Noam Shazeer (Google Brain)
+Niki Parmar (Google Research)
+Jakob Uszkoreit (Google Research)
+Llion Jones (Google Research)
+Aidan N. Gomez (University of Toronto)
+Łukasz Kaiser (Google Brain)
+Illia Polosukhin
+
+## References
+
+[1] Neural Machine Translation by Jointly Learning to Align and Translate
+[2] Sequence to Sequence Learning with Neural Networks
+[3] Learning Phrase Representations using RNN Encoder-Decoder
+"""
+    # Repeat content to make it larger
+    full_content = (content + "\n\n") * 10  # ~35KB
+    return full_content
+
+
+def main():
+    print("=" * 70)
+    print("Recursive Context Pattern - Token Efficiency Demo")
+    print("=" * 70)
+    
+    # Track token usage across methods
+    token_stats = {
+        "traditional": {"input": 0, "output": 0},
+        "recursive": {"input": 0, "output": 0},
+    }
+    
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # 1. Create artifact store and tools
+        store = FileSystemArtifactStore(base_dir=tmpdir)
+        artifact_tools = create_artifact_tools(store=store)
+        
+        print(f"\n1. Created FileSystemArtifactStore")
+        print(f"   Location: {tmpdir}")
+        print(f"   Tools available: {[t.__name__ for t in artifact_tools]}")
+        
+        # 2. Create sample document and store as artifact
+        document = create_sample_document()
+        doc_size = len(document)
+        doc_tokens = doc_size // 4  # Rough estimate: 4 chars per token
+        
+        metadata = ArtifactMetadata(
+            agent_id="research_agent",
+            run_id="demo_run",
+            tool_name="document_loader",
+            turn_id=1
+        )
+        
+        ref = store.store(document, metadata)
+        
+        print(f"\n2. Stored document as artifact")
+        print(f"   Size: {doc_size:,} chars (~{doc_tokens:,} tokens)")
+        print(f"   Path: {ref.path}")
+        print(f"   Summary: {ref.summary[:80]}...")
+        
+        # 3. Create agent WITH artifact tools (Recursive Context approach)
+        agent = Agent(
+            instructions="""You are a research paper analyst with access to artifact exploration tools.
+When given an artifact path, use the tools to explore it efficiently:
+- artifact_grep: Search for patterns (USE THIS FIRST)
+- artifact_head: See first N lines
+- artifact_tail: See last N lines
+- artifact_chunk: Get specific line ranges
+
+Be precise and cite line numbers when finding information.""",
+            tools=artifact_tools,
+            output='silent'  # Quiet output for cleaner demo
+        )
+        
+        print(f"\n3. Created agent with {len(agent.tools)} artifact tools")
+        
+        # 4. Query using RECURSIVE CONTEXT (small prompt + tools)
+        print(f"\n4. RECURSIVE CONTEXT APPROACH")
+        print("-" * 40)
+        
+        prompt = f"""I have a research paper stored at: {ref.path}
+
+Find the BLEU score on the WMT 2014 English-to-German translation task.
+Use artifact_grep to search, then report the answer with line number."""
+
+        prompt_tokens = len(prompt) // 4
+        print(f"   Prompt size: {len(prompt)} chars (~{prompt_tokens} tokens)")
+        print(f"   Document NOT in prompt context!")
+        
+        response = agent.chat(prompt)
+        
+        # Estimate output tokens (could use actual LiteLLM response if available)
+        output_tokens = len(response) // 4
+        
+        # Calculate recursive approach tokens (prompt + tool calls + response)
+        # Tool call: ~200 tokens for tool desc + call
+        # Tool response: ~300 tokens for grep results
+        estimated_total = prompt_tokens + 200 + 300 + output_tokens
+        token_stats["recursive"]["input"] = estimated_total
+        
+        print(f"\n   Agent Response:")
+        print(f"   {response[:200]}...")
+        print(f"\n   Estimated tokens used: ~{estimated_total:,}")
+        
+        # 5. Compare with TRADITIONAL approach (full context)
+        print(f"\n5. TRADITIONAL APPROACH (for comparison)")
+        print("-" * 40)
+        
+        traditional_prompt_tokens = prompt_tokens + doc_tokens
+        token_stats["traditional"]["input"] = traditional_prompt_tokens
+        
+        print(f"   If we passed FULL document in prompt:")
+        print(f"   Prompt: {prompt_tokens} + Document: {doc_tokens} = {traditional_prompt_tokens:,} tokens")
+        
+        # 6. Show token savings
+        print(f"\n6. TOKEN COMPARISON")
+        print("=" * 40)
+        
+        savings = token_stats["traditional"]["input"] - token_stats["recursive"]["input"]
+        savings_pct = (savings / token_stats["traditional"]["input"]) * 100
+        
+        print(f"   Traditional approach: ~{token_stats['traditional']['input']:,} tokens")
+        print(f"   Recursive approach:   ~{token_stats['recursive']['input']:,} tokens")
+        print(f"   -----------------------------------")
+        print(f"   SAVINGS:             ~{savings:,} tokens ({savings_pct:.0f}%)")
+        
+        # 7. Demonstrate other artifact operations
+        print(f"\n7. ARTIFACT OPERATIONS DEMO")
+        print("-" * 40)
+        
+        # head
+        head_result = store.head(ref, lines=5)
+        print(f"   head(5 lines):")
+        for line in head_result.split('\n')[:3]:
+            print(f"      {line[:60]}")
+        
+        # grep
+        grep_results = store.grep(ref, pattern=r"BLEU.*\d+", max_matches=3)
+        print(f"\n   grep('BLEU.*\\d+'):")
+        for match in grep_results[:2]:
+            print(f"      Line {match.line_number}: {match.line_content[:50].strip()}...")
+        
+        # chunk
+        chunk_result = store.chunk(ref, start_line=50, end_line=55)
+        print(f"\n   chunk(lines 50-55):")
+        for line in chunk_result.split('\n')[:3]:
+            print(f"      {line[:60]}")
+        
+        print("\n" + "=" * 70)
+        print("CONCLUSION: Recursive Context Pattern")
+        print("=" * 70)
+        print("""
+The Recursive Context pattern (RLM) provides:
+✓ Token efficiency: ~{savings_pct:.0f}% reduction in token usage
+✓ Scalability: Handle documents of any size
+✓ Precision: Agent searches for exactly what it needs
+✓ Cost savings: Lower API costs due to fewer tokens
+✓ Better accuracy: Avoids "context rot" from large contexts
+""".format(savings_pct=savings_pct))
+
+
+if __name__ == "__main__":
+    main()

examples/embedding/agent_with_embedding.py

@@ -0,0 +1,145 @@
+"""
+Agent-centric Embedding Example
+
+This example shows how to use embeddings with PraisonAI Agents.
+Embeddings are used internally by Agent for:
+- Knowledge retrieval (RAG)
+- Memory storage and search
+- Semantic similarity
+
+You can also use the embedding API directly for custom use cases.
+"""
+
+from praisonaiagents import Agent, embedding
+
+# =============================================================================
+# Example 1: Direct Embedding API
+# =============================================================================
+
+def example_direct_embedding():
+    """Use the embedding API directly."""
+    
+    # Single text embedding
+    result = embedding("Hello, world!")
+    print(f"Single embedding: {len(result.embeddings[0])} dimensions")
+    
+    # Batch embedding
+    texts = ["AI agents are powerful", "Machine learning is fascinating"]
+    result = embedding(texts, model="text-embedding-3-small")
+    print(f"Batch embeddings: {len(result.embeddings)} vectors")
+    
+    # Custom dimensions (for models that support it)
+    result = embedding("Hello", model="text-embedding-3-large", dimensions=256)
+    print(f"Custom dimensions: {len(result.embeddings[0])} dimensions")
+    
+    return result
+
+
+# =============================================================================
+# Example 2: Agent with Knowledge (uses embeddings internally)
+# =============================================================================
+
+def example_agent_with_knowledge():
+    """Agent with knowledge uses embeddings for RAG."""
+    
+    agent = Agent(
+        instructions="You are a helpful assistant with knowledge about AI.",
+        knowledge=["AI agents can use tools to accomplish tasks."],
+        knowledge_config={
+            "embedder": {
+                "provider": "openai",
+                "config": {"model": "text-embedding-3-small"}
+            }
+        }
+    )
+    
+    # The agent uses embeddings internally for knowledge retrieval
+    response = agent.chat("What can AI agents do?")
+    print(f"Agent response: {response}")
+    
+    return agent
+
+
+# =============================================================================
+# Example 3: Agent with Memory (uses embeddings internally)
+# =============================================================================
+
+def example_agent_with_memory():
+    """Agent with memory uses embeddings for semantic search."""
+    
+    agent = Agent(
+        instructions="You are a helpful assistant that remembers conversations.",
+        memory=True,
+        memory_config={
+            "embedding_model": "text-embedding-3-small"
+        }
+    )
+    
+    # The agent uses embeddings internally for memory storage/retrieval
+    agent.chat("My name is Alice and I love Python programming.")
+    response = agent.chat("What do I love?")
+    print(f"Agent remembered: {response}")
+    
+    return agent
+
+
+# =============================================================================
+# Example 4: Custom Embedding for Similarity Search
+# =============================================================================
+
+def example_similarity_search():
+    """Use embeddings for custom similarity search."""
+    import math
+    
+    def cosine_similarity(a, b):
+        dot = sum(x * y for x, y in zip(a, b))
+        norm_a = math.sqrt(sum(x * x for x in a))
+        norm_b = math.sqrt(sum(x * x for x in b))
+        return dot / (norm_a * norm_b) if norm_a and norm_b else 0
+    
+    # Embed documents
+    documents = [
+        "Python is a programming language",
+        "JavaScript runs in browsers",
+        "AI agents can automate tasks"
+    ]
+    doc_embeddings = embedding(documents)
+    
+    # Embed query
+    query = "What language is used for web development?"
+    query_embedding = embedding(query)
+    
+    # Find most similar document
+    similarities = [
+        cosine_similarity(query_embedding.embeddings[0], doc_emb)
+        for doc_emb in doc_embeddings.embeddings
+    ]
+    
+    best_idx = similarities.index(max(similarities))
+    print(f"Most similar document: {documents[best_idx]}")
+    print(f"Similarity score: {similarities[best_idx]:.4f}")
+    
+    return similarities
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Example 1: Direct Embedding API")
+    print("=" * 60)
+    example_direct_embedding()
+    
+    print("\n" + "=" * 60)
+    print("Example 4: Custom Similarity Search")
+    print("=" * 60)
+    example_similarity_search()
+    
+    # Uncomment to run agent examples (requires API key)
+    # print("\n" + "=" * 60)
+    # print("Example 2: Agent with Knowledge")
+    # print("=" * 60)
+    # example_agent_with_knowledge()
+    
+    # print("\n" + "=" * 60)
+    # print("Example 3: Agent with Memory")
+    # print("=" * 60)
+    # example_agent_with_memory()

examples/python/audio/stt_basic.py

@@ -1,5 +1,8 @@
+# Speech-to-Text with Groq (fastest)
+# Requires: export GROQ_API_KEY=your-key
+
 from praisonaiagents import AudioAgent
 
-agent = AudioAgent(llm="openai/whisper-1")
-text = agent.listen("audio.mp3")
+agent = AudioAgent(llm="groq/whisper-large-v3")
+text = agent.listen("audio.mp3")  # Replace with your audio file
 print(text)

examples/python/audio/tts_basic.py

@@ -1,4 +1,8 @@
+# Text-to-Speech with OpenAI
+# Requires: export OPENAI_API_KEY=your-key
+
 from praisonaiagents import AudioAgent
 
 agent = AudioAgent(llm="openai/tts-1")
 agent.say("Hello world!", output="hello.mp3")
+print("Saved to hello.mp3")

examples/python/image/image_basic.py

@@ -1,3 +1,6 @@
+# Image Generation with DALL-E
+# Requires: export OPENAI_API_KEY=your-key
+
 from praisonaiagents import ImageAgent
 
 agent = ImageAgent(llm="openai/dall-e-3")

examples/python/mcp/mcp-auth-storage.py

@@ -0,0 +1,72 @@
+"""
+MCP Auth Storage Example
+
+This example demonstrates how to use the MCPAuthStorage class
+to manage OAuth tokens for remote MCP servers.
+"""
+
+from praisonaiagents.mcp import (
+    MCPAuthStorage,
+    generate_state,
+    generate_code_verifier,
+    generate_code_challenge,
+    get_redirect_url
+)
+import time
+
+# Initialize auth storage
+storage = MCPAuthStorage()
+print(f"Auth storage path: {storage.filepath}")
+
+# Generate PKCE parameters for OAuth flow
+state = generate_state()
+verifier = generate_code_verifier()
+challenge = generate_code_challenge(verifier)
+redirect_url = get_redirect_url()
+
+print("\n=== PKCE Parameters ===")
+print(f"State: {state[:30]}...")
+print(f"Verifier: {verifier[:30]}...")
+print(f"Challenge: {challenge[:30]}...")
+print(f"Redirect URL: {redirect_url}")
+
+# Store tokens (simulating OAuth callback)
+print("\n=== Storing Tokens ===")
+storage.set_tokens("example-server", {
+    "access_token": "example_access_token_12345",
+    "refresh_token": "example_refresh_token_67890",
+    "expires_at": time.time() + 3600,  # 1 hour from now
+    "scope": "read write"
+}, server_url="https://mcp.example.com/mcp")
+print("Tokens stored successfully")
+
+# Retrieve tokens
+entry = storage.get("example-server")
+print("\n=== Retrieved Entry ===")
+print(f"Server URL: {entry['server_url']}")
+print(f"Access Token: {entry['tokens']['access_token'][:20]}...")
+print(f"Expires At: {entry['tokens']['expires_at']}")
+
+# Check token expiration
+is_expired = storage.is_token_expired("example-server")
+print("\n=== Token Status ===")
+print(f"Is Expired: {is_expired}")
+
+# URL validation
+print("\n=== URL Validation ===")
+valid = storage.get_for_url("example-server", "https://mcp.example.com/mcp")
+print(f"Same URL: {valid is not None}")
+
+invalid = storage.get_for_url("example-server", "https://different.url/mcp")
+print(f"Different URL: {invalid is None}")
+
+# List all entries
+print("\n=== All Entries ===")
+all_entries = storage.all()
+for name, entry in all_entries.items():
+    print(f"  {name}: {entry.get('server_url', 'N/A')}")
+
+# Cleanup
+storage.remove("example-server")
+print("\n=== Cleanup ===")
+print("Removed example-server entry")

examples/python/mcp/remote-mcp-oauth.py

@@ -0,0 +1,32 @@
+"""
+Remote MCP Server with OAuth Authentication Example
+
+This example shows how to connect to a remote MCP server
+that requires OAuth 2.1 authentication.
+
+Prerequisites:
+1. Configure the remote server in ~/.praison/config.yaml
+2. Run: praisonai mcp auth <server-name>
+"""
+
+from praisonaiagents import Agent, MCP
+from praisonaiagents.mcp import MCPAuthStorage
+
+# Check if authenticated
+storage = MCPAuthStorage()
+entry = storage.get("my-remote-server")
+
+if not entry or not entry.get("tokens"):
+    print("Not authenticated. Run: praisonai mcp auth my-remote-server")
+    exit(1)
+
+# Create agent with remote MCP server
+agent = Agent(
+    name="Remote Assistant",
+    instructions="You help with tasks using the remote MCP server.",
+    llm="gpt-4o-mini",
+    tools=MCP("https://mcp.example.com/mcp")
+)
+
+# Start the agent
+agent.start("What tools are available?")

examples/python/ocr/ocr_advanced.py

@@ -2,13 +2,11 @@ from praisonaiagents import OCRAgent
 
 agent = OCRAgent()
 
-# Extract with page selection
-result = agent.extract("document.pdf", pages=[1, 2, 3])
+# Extract with page selection (pages is 0-indexed)
+result = agent.extract("https://arxiv.org/pdf/2201.04234", pages=[0, 1])
 for page in result.pages:
-    print(f"Page {page.index}: {page.markdown[:100]}")
+    print(f"Page {page.index}: {page.markdown[:200]}")
 
-# From image
-text = agent.read("screenshot.png")
-
-# From URL
-text = agent.read("https://example.com/doc.pdf")
+# Quick read - just get the text
+text = agent.read("https://arxiv.org/pdf/2201.04234")
+print(f"Total length: {len(text)} chars")

examples/python/ocr/ocr_basic.py

@@ -1,5 +1,9 @@
+# OCR with Mistral
+# Requires: export MISTRAL_API_KEY=your-key
+# Note: Source must be a URL (https://) or base64
+
 from praisonaiagents import OCRAgent
 
 agent = OCRAgent()
-text = agent.read("document.pdf")
-print(text)
+text = agent.read("https://arxiv.org/pdf/2201.04234")
+print(text[:500])  # First 500 chars

examples/python/ocr/ocr_mistral.py

@@ -1,5 +1,5 @@
 from praisonaiagents import OCRAgent
 
 agent = OCRAgent(llm="mistral/mistral-ocr-latest")
-text = agent.read("document.pdf")
-print(text)
+text = agent.read("https://arxiv.org/pdf/2201.04234")
+print(text[:500])  # First 500 chars

examples/python/video/video_basic.py

@@ -1,5 +1,9 @@
+# Video Generation with OpenAI Sora
+# Requires: export OPENAI_API_KEY=your-key
+# Note: Sora requires organization verification
+
 from praisonaiagents import VideoAgent
 
 agent = VideoAgent(llm="openai/sora-2")
 video = agent.generate("A cat playing with yarn")
-print(video.id)
+print(f"Video ID: {video.id}")

examples/recipe_output_modes.py

@@ -0,0 +1,93 @@
+"""
+Recipe Output Modes Example
+
+Demonstrates how to use different output modes when running recipes.
+Output modes control the verbosity and format of execution feedback.
+
+Available modes:
+- silent: No output (default, best performance)
+- status: Shows tool calls inline
+- trace: Timestamped execution trace
+- verbose: Full interactive output with panels
+- debug: Trace + metrics (tokens, cost, model)
+- json: Machine-readable JSONL events
+
+CLI Usage:
+    praisonai recipe run my-recipe --output status
+    praisonai recipe run my-recipe --output trace
+    praisonai recipe run my-recipe -v  # alias for --output verbose
+"""
+
+from praisonai.recipe import run
+
+
+def example_silent_mode():
+    """Silent mode - no output, best performance."""
+    print("=== Silent Mode (default) ===")
+    result = run(
+        "my-recipe",
+        input={"query": "Hello"},
+        options={"output": "silent"}
+    )
+    print(f"Result: {result}")
+
+
+def example_status_mode():
+    """Status mode - shows tool calls inline."""
+    print("\n=== Status Mode ===")
+    print("Shows: ▸ tool → result ✓")
+    result = run(
+        "my-recipe",
+        input={"query": "Hello"},
+        options={"output": "status"}
+    )
+    print(f"Result: {result}")
+
+
+def example_trace_mode():
+    """Trace mode - timestamped execution trace."""
+    print("\n=== Trace Mode ===")
+    print("Shows: [HH:MM:SS] ▸ tool → result [0.2s] ✓")
+    result = run(
+        "my-recipe",
+        input={"query": "Hello"},
+        options={"output": "trace"}
+    )
+    print(f"Result: {result}")
+
+
+def example_verbose_mode():
+    """Verbose mode - full interactive output with panels."""
+    print("\n=== Verbose Mode ===")
+    result = run(
+        "my-recipe",
+        input={"query": "Hello"},
+        options={"output": "verbose"}
+    )
+    print(f"Result: {result}")
+
+
+def example_backward_compat():
+    """Backward compatibility - verbose flag maps to output='verbose'."""
+    print("\n=== Backward Compatibility ===")
+    print("Using verbose=True (deprecated, use output='verbose' instead)")
+    result = run(
+        "my-recipe",
+        input={"query": "Hello"},
+        options={"verbose": True}  # Maps to output="verbose"
+    )
+    print(f"Result: {result}")
+
+
+if __name__ == "__main__":
+    print(__doc__)
+    print("Note: These examples require a recipe named 'my-recipe' to be available.")
+    print("Create one with: praisonai recipe init my-recipe")
+    print()
+    
+    # Uncomment to run examples:
+    # example_silent_mode()
+    # example_status_mode()
+    # example_trace_mode()
+    # example_verbose_mode()
+    # example_backward_compat()

examples/storage/storage_backends_example.py

@@ -0,0 +1,274 @@
+"""
+Storage Backends Example for PraisonAI Agents.
+
+Demonstrates how to use different storage backends (File, SQLite, Redis)
+for agent data persistence across all supported components.
+
+Supported Backends:
+- FileBackend: JSON file-based storage (default, zero dependencies)
+- SQLiteBackend: SQLite database storage (zero dependencies)
+- RedisBackend: Redis-based storage (requires redis package)
+
+Supported Components:
+- BaseJSONStore: Generic JSON storage
+- TrainingStorage: Agent training data
+- SessionManager: CLI session state
+- RunHistory: Recipe run history
+- MCPToolIndex: MCP tool schemas
+
+Usage:
+    python storage_backends_example.py
+"""
+
+from praisonaiagents.storage import (
+    FileBackend,
+    SQLiteBackend,
+    BaseJSONStore,
+    get_backend,
+)
+
+
+def example_file_backend():
+    """Example: Using FileBackend for JSON file storage."""
+    print("\n=== FileBackend Example ===")
+    
+    # Create file backend
+    backend = FileBackend(storage_dir="/tmp/praison_example")
+    
+    # Save data
+    backend.save("agent_state", {
+        "agent_name": "Research Assistant",
+        "messages": ["Hello", "How can I help?"],
+        "context": {"topic": "AI"}
+    })
+    
+    # Load data
+    data = backend.load("agent_state")
+    print(f"Loaded: {data}")
+    
+    # List keys
+    keys = backend.list_keys()
+    print(f"Keys: {keys}")
+    
+    # Check existence
+    print(f"Exists: {backend.exists('agent_state')}")
+    
+    # Clean up
+    backend.delete("agent_state")
+    print("Deleted agent_state")
+
+
+def example_sqlite_backend():
+    """Example: Using SQLiteBackend for database storage."""
+    print("\n=== SQLiteBackend Example ===")
+    
+    # Create SQLite backend (single file, better for concurrent access)
+    backend = SQLiteBackend(db_path="/tmp/praison_example.db")
+    
+    # Save multiple sessions
+    for i in range(3):
+        backend.save(f"session_{i}", {
+            "session_id": f"sess-{i}",
+            "messages": [f"Message {j}" for j in range(i + 1)],
+            "timestamp": "2024-01-01T00:00:00Z"
+        })
+    
+    # List all sessions
+    keys = backend.list_keys(prefix="session_")
+    print(f"Sessions: {keys}")
+    
+    # Load specific session
+    session = backend.load("session_1")
+    print(f"Session 1: {session}")
+    
+    # Clean up
+    backend.clear()
+    backend.close()
+    print("Cleaned up SQLite backend")
+
+
+def example_backend_factory():
+    """Example: Using get_backend factory function."""
+    print("\n=== Backend Factory Example ===")
+    
+    # Easy switching between backends
+    backend_type = "sqlite"  # Change to "file" or "redis" as needed
+    
+    if backend_type == "file":
+        backend = get_backend("file", storage_dir="/tmp/praison_factory")
+    elif backend_type == "sqlite":
+        backend = get_backend("sqlite", db_path="/tmp/praison_factory.db")
+    # elif backend_type == "redis":
+    #     backend = get_backend("redis", url="redis://localhost:6379")
+    
+    # Use backend
+    backend.save("config", {"model": "gpt-4", "temperature": 0.7})
+    config = backend.load("config")
+    print(f"Config: {config}")
+    
+    # Clean up
+    backend.delete("config")
+
+
+def example_with_base_json_store():
+    """Example: Using BaseJSONStore with custom backend."""
+    print("\n=== BaseJSONStore with Backend Example ===")
+    
+    # Create a custom store with SQLite backend
+    backend = SQLiteBackend(db_path="/tmp/praison_store.db")
+    
+    store = BaseJSONStore(
+        storage_path="/tmp/my_store.json",  # Used as key name
+        backend=backend,
+    )
+    
+    # Save and load data
+    store.save({"items": [1, 2, 3], "count": 3})
+    data = store.load()
+    print(f"Store data: {data}")
+    
+    # Clean up
+    store.delete()
+    backend.close()
+
+
+def example_training_storage_with_backend():
+    """Example: Using TrainingStorage with SQLite backend."""
+    print("\n=== TrainingStorage with Backend Example ===")
+    
+    try:
+        from praisonai.train.agents.storage import TrainingStorage
+        
+        # Create SQLite backend for training data
+        backend = SQLiteBackend(db_path="/tmp/training.db")
+        
+        # Use with TrainingStorage
+        storage = TrainingStorage(
+            session_id="train-example-001",
+            backend=backend,
+        )
+        
+        print(f"Training storage created with session: {storage.session_id}")
+        print("Training data will be stored in SQLite instead of JSON files")
+        
+        backend.close()
+    except ImportError:
+        print("praisonai package not installed, skipping training example")
+
+
+def example_session_manager_with_backend():
+    """Example: Using SessionManager with SQLite backend."""
+    print("\n=== SessionManager with Backend Example ===")
+    
+    try:
+        from praisonai.cli.state.sessions import SessionManager
+        
+        # Create SQLite backend for sessions
+        backend = SQLiteBackend(db_path="/tmp/sessions.db")
+        
+        # Use with SessionManager
+        manager = SessionManager(backend=backend)
+        sessions = manager.list(limit=10)
+        
+        print(f"SessionManager created with SQLite backend, {len(sessions)} sessions")
+        print("Sessions will be stored in database instead of files")
+        
+        backend.close()
+    except ImportError:
+        print("praisonai package not installed, skipping session example")
+
+
+def example_run_history_with_backend():
+    """Example: Using RunHistory with SQLite backend."""
+    print("\n=== RunHistory with Backend Example ===")
+    
+    try:
+        from praisonai.recipe.history import RunHistory
+        
+        # Create SQLite backend for run history
+        backend = SQLiteBackend(db_path="/tmp/runs.db")
+        
+        # Use with RunHistory
+        history = RunHistory(backend=backend)
+        runs = history.list()
+        
+        print(f"RunHistory created with SQLite backend, {len(runs)} runs")
+        print("Recipe runs will be stored in database instead of files")
+        
+        backend.close()
+    except ImportError:
+        print("praisonai package not installed, skipping run history example")
+
+
+def example_mcp_tool_index_with_backend():
+    """Example: Using MCPToolIndex with SQLite backend."""
+    print("\n=== MCPToolIndex with Backend Example ===")
+    
+    try:
+        from praisonai.mcp_server.tool_index import MCPToolIndex
+        
+        # Create SQLite backend for MCP tool index
+        backend = SQLiteBackend(db_path="/tmp/mcp.db")
+        
+        # Use with MCPToolIndex
+        index = MCPToolIndex(backend=backend)
+        servers = index.list_servers()
+        
+        print(f"MCPToolIndex created with SQLite backend, {len(servers)} servers")
+        print("MCP tool schemas will be stored in database instead of files")
+        
+        backend.close()
+    except ImportError:
+        print("praisonai package not installed, skipping MCP example")
+
+
+def example_redis_backend():
+    """Example: Using RedisBackend for high-speed caching."""
+    print("\n=== RedisBackend Example ===")
+    
+    try:
+        from praisonaiagents.storage import RedisBackend
+        
+        # Create Redis backend with TTL
+        backend = RedisBackend(
+            url="redis://localhost:6379",
+            prefix="example:",
+            ttl=3600  # 1 hour TTL
+        )
+        
+        # CRUD operations
+        backend.save("test_key", {"data": "value", "count": 42})
+        data = backend.load("test_key")
+        print(f"Loaded from Redis: {data}")
+        
+        # List keys
+        keys = backend.list_keys()
+        print(f"Keys: {keys}")
+        
+        # Clean up
+        backend.delete("test_key")
+        backend.close()
+        print("RedisBackend example completed")
+        
+    except ImportError:
+        print("redis package not installed: pip install redis")
+    except Exception as e:
+        print(f"Redis not available: {e}")
+
+
+if __name__ == "__main__":
+    print("PraisonAI Storage Backends Examples")
+    print("=" * 50)
+    
+    example_file_backend()
+    example_sqlite_backend()
+    example_backend_factory()
+    example_with_base_json_store()
+    example_training_storage_with_backend()
+    example_session_manager_with_backend()
+    example_run_history_with_backend()
+    example_mcp_tool_index_with_backend()
+    example_redis_backend()
+    
+    print("\n" + "=" * 50)
+    print("All examples completed!")

examples/workflow_output_modes.py

@@ -0,0 +1,106 @@
+"""
+Workflow Output Modes Example
+
+Demonstrates how to use different output modes with Workflow.
+Workflows use the same output presets as Agent for consistency (DRY approach).
+
+Available presets:
+- silent: No output (default, fastest)
+- status: Tool calls + final output inline
+- trace: Timestamped execution trace
+- verbose: Full Rich panels
+- debug: Trace + metrics
+- json: JSONL output for piping
+"""
+
+from praisonaiagents import Agent, Workflow
+
+
+def main():
+    # Create a simple agent
+    researcher = Agent(
+        name="Researcher",
+        role="Research Analyst",
+        goal="Research and provide information",
+        instructions="You are a research analyst. Provide concise, factual information."
+    )
+    
+    # Example 1: Status output (recommended for CLI)
+    print("=" * 50)
+    print("Example 1: Status Output (output='status')")
+    print("=" * 50)
+    
+    workflow_status = Workflow(
+        name="Status Example",
+        steps=[researcher],
+        output="status"  # Shows: ▸ tool → result ✓
+    )
+    # result = workflow_status.start("What is Python?")
+    print("Workflow created with output='status'")
+    print(f"  actions_trace: {getattr(workflow_status._output_config, 'actions_trace', False)}")
+    print()
+    
+    # Example 2: Trace output (with timestamps)
+    print("=" * 50)
+    print("Example 2: Trace Output (output='trace')")
+    print("=" * 50)
+    
+    workflow_trace = Workflow(
+        name="Trace Example",
+        steps=[researcher],
+        output="trace"  # Shows: [HH:MM:SS] ▸ tool → result [0.2s] ✓
+    )
+    print("Workflow created with output='trace'")
+    print(f"  status_trace: {getattr(workflow_trace._output_config, 'status_trace', False)}")
+    print()
+    
+    # Example 3: Verbose output (Rich panels)
+    print("=" * 50)
+    print("Example 3: Verbose Output (output='verbose')")
+    print("=" * 50)
+    
+    workflow_verbose = Workflow(
+        name="Verbose Example",
+        steps=[researcher],
+        output="verbose"  # Full interactive display
+    )
+    print("Workflow created with output='verbose'")
+    print(f"  verbose: {getattr(workflow_verbose._output_config, 'verbose', False)}")
+    print()
+    
+    # Example 4: Silent output (default, fastest)
+    print("=" * 50)
+    print("Example 4: Silent Output (output='silent' or default)")
+    print("=" * 50)
+    
+    workflow_silent = Workflow(
+        name="Silent Example",
+        steps=[researcher],
+        output="silent"  # No output - zero overhead
+    )
+    print("Workflow created with output='silent'")
+    print(f"  verbose: {getattr(workflow_silent._output_config, 'verbose', False)}")
+    print(f"  actions_trace: {getattr(workflow_silent._output_config, 'actions_trace', False)}")
+    print()
+    
+    # Example 5: JSON output (for piping)
+    print("=" * 50)
+    print("Example 5: JSON Output (output='json')")
+    print("=" * 50)
+    
+    workflow_json = Workflow(
+        name="JSON Example",
+        steps=[researcher],
+        output="json"  # JSONL output for scripting
+    )
+    print("Workflow created with output='json'")
+    print(f"  json_output: {getattr(workflow_json._output_config, 'json_output', False)}")
+    print()
+    
+    print("=" * 50)
+    print("All examples completed successfully!")
+    print("=" * 50)
+
+
+if __name__ == "__main__":
+    main()

src/praisonai-agents/praisonaiagents/__init__.py

@@ -56,6 +56,15 @@ from . import memory
 from . import workflows
 # Note: knowledge and mcp are lazy-loaded via __getattr__ due to heavy deps
 
+# Embedding API - explicit import shadows subpackage, enabling:
+#   from praisonaiagents import embedding, embeddings, EmbeddingResult, get_dimensions
+# Note: litellm is still lazy-loaded INSIDE the functions (no performance impact)
+from .embedding.embed import embedding, aembedding
+from .embedding.result import EmbeddingResult
+from .embedding.dimensions import get_dimensions
+embeddings = embedding  # Plural alias (OpenAI style: client.embeddings.create)
+aembeddings = aembedding  # Plural alias for async
+
 # Workflows - lightweight module
 from .workflows import (
     Workflow, WorkflowStep, WorkflowContext, StepResult,
@@ -392,6 +401,33 @@ def __getattr__(name):
         _lazy_cache[name] = Memory
         return Memory
     
+    # Embedding module (unified embedding API)
+    # Note: Must import from .embed submodule to get functions, not the module
+    elif name == "embedding":
+        from praisonaiagents.embedding.embed import embedding
+        _lazy_cache[name] = embedding
+        return embedding
+    elif name == "aembedding":
+        from praisonaiagents.embedding.embed import aembedding
+        _lazy_cache[name] = aembedding
+        return aembedding
+    elif name == "embed":
+        from praisonaiagents.embedding.embed import embed
+        _lazy_cache[name] = embed
+        return embed
+    elif name == "aembed":
+        from praisonaiagents.embedding.embed import aembed
+        _lazy_cache[name] = aembed
+        return aembed
+    elif name == "EmbeddingResult":
+        from praisonaiagents.embedding.result import EmbeddingResult
+        _lazy_cache[name] = EmbeddingResult
+        return EmbeddingResult
+    elif name == "get_dimensions":
+        from praisonaiagents.embedding.dimensions import get_dimensions
+        _lazy_cache[name] = get_dimensions
+        return get_dimensions
+    
     # Planning mode support (lazy)
     elif name in ("Plan", "PlanStep", "TodoList", "TodoItem", "PlanStorage", 
                   "PlanningAgent", "ApprovalCallback", "READ_ONLY_TOOLS", 
@@ -498,7 +534,15 @@ def __getattr__(name):
         _lazy_cache[name] = result
         return result
     
-    # Context management config (already exists)
+    # Context management config
+    elif name == "ContextConfig":
+        from .context.models import ContextConfig
+        _lazy_cache[name] = ContextConfig
+        return ContextConfig
+    elif name == "OptimizerStrategy":
+        from .context.models import OptimizerStrategy
+        _lazy_cache[name] = OptimizerStrategy
+        return OptimizerStrategy
     elif name == "ManagerConfig":
         from .context.manager import ManagerConfig
         _lazy_cache[name] = ManagerConfig
@@ -661,6 +705,15 @@ __all__ = [
     'tool',
     'Tools',
     
+    # Embedding API - simplified imports
+    # Usage: from praisonaiagents import embedding, EmbeddingResult
+    'embedding',
+    'embeddings',  # Plural alias (OpenAI style)
+    'aembedding',
+    'aembeddings',  # Plural alias for async
+    'EmbeddingResult',
+    'get_dimensions',
+    
     # Sub-packages for organized imports
     # Usage: import praisonaiagents as pa; pa.config.MemoryConfig
     'config',

src/praisonai-agents/praisonaiagents/agent/__init__.py

@@ -39,6 +39,7 @@ from .protocols import (
     ToolAwareAgentProtocol,
     MemoryAwareAgentProtocol,
     FullAgentProtocol,
+    ContextEngineerProtocol,
 )
 
 __all__ = [
@@ -87,4 +88,5 @@ __all__ = [
     'ToolAwareAgentProtocol',
     'MemoryAwareAgentProtocol',
     'FullAgentProtocol',
+    'ContextEngineerProtocol',
 ]

src/praisonai-agents/praisonaiagents/agent/agent.py

@@ -1220,12 +1220,38 @@ Your Goal: {self.goal}
                 agent_name=self.name or "Agent",
             )
         elif isinstance(self._context_param, ManagerConfig):
-            # Use provided config
+            # Use provided ManagerConfig
             self._context_manager = ContextManager(
                 model=self.llm if isinstance(self.llm, str) else "gpt-4o-mini",
                 config=self._context_param,
                 agent_name=self.name or "Agent",
             )
+        elif hasattr(self._context_param, 'auto_compact') and hasattr(self._context_param, 'tool_output_max'):
+            # ContextConfig from YAML - convert to ManagerConfig
+            try:
+                from ..context.models import ContextConfig as _ContextConfig
+                if isinstance(self._context_param, _ContextConfig):
+                    # Build ManagerConfig from ContextConfig fields
+                    manager_config = ManagerConfig(
+                        auto_compact=self._context_param.auto_compact,
+                        compact_threshold=self._context_param.compact_threshold,
+                        strategy=self._context_param.strategy,
+                        output_reserve=self._context_param.output_reserve,
+                        default_tool_output_max=self._context_param.tool_output_max,  # Map field name
+                        protected_tools=list(self._context_param.protected_tools),
+                        keep_recent_turns=self._context_param.keep_recent_turns,
+                        monitor_enabled=self._context_param.monitor.enabled if self._context_param.monitor else False,
+                    )
+                    self._context_manager = ContextManager(
+                        model=self.llm if isinstance(self.llm, str) else "gpt-4o-mini",
+                        config=manager_config,
+                        agent_name=self.name or "Agent",
+                    )
+                else:
+                    self._context_manager = None
+            except Exception as e:
+                logging.debug(f"ContextConfig conversion failed: {e}")
+                self._context_manager = None
         elif hasattr(self._context_param, 'process'):
             # Already a ContextManager instance
             self._context_manager = self._context_param
@@ -2676,7 +2702,98 @@ Your Goal: {self.goal}"""
             self._system_prompt_cache[cache_key] = system_prompt
         return system_prompt
 
-    def _build_messages(self, prompt, temperature=1.0, output_json=None, output_pydantic=None, tools=None):
+    def _build_response_format(self, schema_model):
+        """Build response_format dict for native structured output.
+        
+        Args:
+            schema_model: Pydantic model or dict schema
+            
+        Returns:
+            Dict suitable for response_format parameter, or None if not applicable
+        """
+        if not schema_model:
+            return None
+        
+        def _add_additional_properties_false(schema):
+            """Recursively add additionalProperties: false and required array to all object schemas."""
+            if isinstance(schema, dict):
+                if schema.get('type') == 'object':
+                    schema['additionalProperties'] = False
+                    # Add required array with all property keys (OpenAI strict mode requirement)
+                    if 'properties' in schema:
+                        schema['required'] = list(schema['properties'].keys())
+                # Recurse into properties
+                if 'properties' in schema:
+                    for prop in schema['properties'].values():
+                        _add_additional_properties_false(prop)
+                # Recurse into items (for arrays)
+                if 'items' in schema:
+                    _add_additional_properties_false(schema['items'])
+                # Recurse into $defs
+                if '$defs' in schema:
+                    for def_schema in schema['$defs'].values():
+                        _add_additional_properties_false(def_schema)
+            return schema
+        
+        def _wrap_array_in_object(schema):
+            """Wrap array schema in object since OpenAI requires root type to be object."""
+            if schema.get('type') == 'array':
+                return {
+                    'type': 'object',
+                    'properties': {
+                        'items': schema
+                    },
+                    'required': ['items'],
+                    'additionalProperties': False
+                }
+            return schema
+        
+        # Handle Pydantic model
+        if hasattr(schema_model, 'model_json_schema'):
+            schema = schema_model.model_json_schema()
+            schema = _add_additional_properties_false(schema)
+            schema = _wrap_array_in_object(schema)
+            name = getattr(schema_model, '__name__', 'response')
+            return {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": name,
+                    "schema": schema,
+                    "strict": True
+                }
+            }
+        
+        # Handle dict schema (inline JSON schema from YAML)
+        if isinstance(schema_model, dict):
+            schema = schema_model.copy()
+            schema = _add_additional_properties_false(schema)
+            schema = _wrap_array_in_object(schema)
+            return {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "structured_output",
+                    "schema": schema,
+                    "strict": True
+                }
+            }
+        
+        return None
+
+    def _supports_native_structured_output(self):
+        """Check if current model supports native structured output via response_format.
+        
+        Auto-detects based on model capabilities using LiteLLM.
+        
+        Returns:
+            bool: True if model supports response_format with json_schema
+        """
+        try:
+            from ..llm.model_capabilities import supports_structured_outputs
+            return supports_structured_outputs(self.llm)
+        except Exception:
+            return False
+
+    def _build_messages(self, prompt, temperature=1.0, output_json=None, output_pydantic=None, tools=None, use_native_format=False):
         """Build messages list for chat completion.
         
         Args:
@@ -2685,32 +2802,36 @@ Your Goal: {self.goal}"""
             output_json: Optional Pydantic model for JSON output
             output_pydantic: Optional Pydantic model for JSON output (alias)
             tools: Optional list of tools to use (defaults to self.tools)
+            use_native_format: If True, skip text injection (native response_format will be used)
             
         Returns:
-            tuple: (messages list, original prompt)
+            Tuple of (messages list, original prompt)
         """
-        # Build system prompt using the helper method
-        system_prompt = self._build_system_prompt(tools)
+        messages = []
+        original_prompt = None
         
         # Use openai_client's build_messages method if available
         if self._openai_client is not None:
             messages, original_prompt = self._openai_client.build_messages(
                 prompt=prompt,
-                system_prompt=system_prompt,
+                system_prompt=self._build_system_prompt(
+                    tools=tools,
+                ),
                 chat_history=self.chat_history,
-                output_json=output_json,
-                output_pydantic=output_pydantic
+                output_json=None if use_native_format else output_json,
+                output_pydantic=None if use_native_format else output_pydantic
             )
         else:
-            # Fallback implementation for when OpenAI client is not available
-            messages = []
-            
-            # Add system message if provided
+            # Build messages manually
+            system_prompt = self._build_system_prompt(
+                tools=tools,
+            )
             if system_prompt:
                 messages.append({"role": "system", "content": system_prompt})
             
             # Add chat history
-            messages.extend(self.chat_history)
+            if self.chat_history:
+                messages.extend(self.chat_history)
             
             # Add user prompt
             if isinstance(prompt, list):
@@ -2720,11 +2841,19 @@ Your Goal: {self.goal}"""
                 messages.append({"role": "user", "content": str(prompt)})
                 original_prompt = str(prompt)
             
-            # Add JSON format instruction if needed
-            if output_json or output_pydantic:
-                model = output_pydantic or output_json
-                json_instruction = f"\nPlease respond with valid JSON matching this schema: {model.model_json_schema()}"
-                messages[-1]["content"] += json_instruction
+            # Add JSON format instruction if needed (only when not using native format)
+            if not use_native_format and (output_json or output_pydantic):
+                schema_model = output_pydantic or output_json
+                # Handle Pydantic model
+                if hasattr(schema_model, 'model_json_schema'):
+                    import json
+                    json_instruction = f"\nPlease respond with valid JSON matching this schema: {json.dumps(schema_model.model_json_schema())}"
+                    messages[-1]["content"] += json_instruction
+                # Handle inline dict schema (Option A from YAML)
+                elif isinstance(schema_model, dict):
+                    import json
+                    json_instruction = f"\nPlease respond with valid JSON matching this schema: {json.dumps(schema_model)}"
+                    messages[-1]["content"] += json_instruction
         
         return messages, original_prompt
 
@@ -2912,11 +3041,54 @@ Your Goal: {self.goal}"""
         return self._execute_tool_with_context(function_name, arguments, state)
     
     def _execute_tool_with_context(self, function_name, arguments, state):
-        """Execute tool within injection context."""
+        """Execute tool within injection context, with optional output truncation."""
         from ..tools.injected import with_injection_context
         
         with with_injection_context(state):
-            return self._execute_tool_impl(function_name, arguments)
+            result = self._execute_tool_impl(function_name, arguments)
+        
+        # Apply context-aware truncation if context management is enabled
+        # This prevents tool outputs (e.g., search results) from exploding the context
+        if self.context_manager and result:
+            try:
+                result_str = str(result)
+                truncated = self._truncate_tool_output(function_name, result_str)
+                if len(truncated) < len(result_str):
+                    logging.debug(f"Truncated {function_name} output from {len(result_str)} to {len(truncated)} chars")
+                    # For dicts, truncate large string fields (e.g., raw_content from search)
+                    if isinstance(result, dict):
+                        result = self._truncate_dict_fields(result, function_name)
+                    else:
+                        return truncated
+            except Exception as e:
+                logging.debug(f"Tool truncation skipped: {e}")
+        
+        return result
+    
+    def _truncate_dict_fields(self, data: dict, tool_name: str, max_field_chars: int = None) -> dict:
+        """Truncate large string fields in a dict to prevent context overflow."""
+        if max_field_chars is None:
+            # Use tool budget from context manager (default 5000 tokens * 4 chars/token = 20000 chars)
+            max_tokens = self.context_manager.get_tool_budget(tool_name) if self.context_manager else 5000
+            max_field_chars = max_tokens * 4
+        
+        result = {}
+        for key, value in data.items():
+            if isinstance(value, str) and len(value) > max_field_chars:
+                # Truncate large string fields
+                result[key] = value[:max_field_chars] + "...[truncated]"
+                logging.debug(f"Truncated field '{key}' from {len(value)} to {max_field_chars} chars")
+            elif isinstance(value, dict):
+                result[key] = self._truncate_dict_fields(value, tool_name, max_field_chars)
+            elif isinstance(value, list):
+                result[key] = [
+                    self._truncate_dict_fields(item, tool_name, max_field_chars) if isinstance(item, dict)
+                    else (item[:max_field_chars] + "...[truncated]" if isinstance(item, str) and len(item) > max_field_chars else item)
+                    for item in value
+                ]
+            else:
+                result[key] = value
+        return result
     
     def _execute_tool_impl(self, function_name, arguments):
         """Internal tool execution implementation."""
@@ -3254,7 +3426,7 @@ Your Goal: {self.goal}"""
             reasoning_steps=reasoning_steps
         )
 
-    def _chat_completion(self, messages, temperature=1.0, tools=None, stream=True, reasoning_steps=False, task_name=None, task_description=None, task_id=None):
+    def _chat_completion(self, messages, temperature=1.0, tools=None, stream=True, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None):
         start_time = time.time()
         logging.debug(f"{self.name} sending messages to LLM: {messages}")
 
@@ -3341,19 +3513,24 @@ Your Goal: {self.goal}"""
                 if self._openai_client is None:
                     raise ValueError("OpenAI client is not initialized. Please provide OPENAI_API_KEY or use a custom LLM provider.")
                 
-                final_response = self._openai_client.chat_completion_with_tools(
-                    messages=messages,
-                    model=self.llm,
-                    temperature=temperature,
-                    tools=formatted_tools,  # Already formatted for OpenAI
-                    execute_tool_fn=self.execute_tool,
-                    stream=stream,
-                    console=self.console if (self.verbose or stream) else None,
-                    display_fn=self.display_generating if self.verbose else None,
-                    reasoning_steps=reasoning_steps,
-                    verbose=self.verbose,
-                    max_iterations=10
-                )
+                # Build kwargs including response_format if provided
+                chat_kwargs = {
+                    "messages": messages,
+                    "model": self.llm,
+                    "temperature": temperature,
+                    "tools": formatted_tools,  # Already formatted for OpenAI
+                    "execute_tool_fn": self.execute_tool,
+                    "stream": stream,
+                    "console": self.console if (self.verbose or stream) else None,
+                    "display_fn": self.display_generating if self.verbose else None,
+                    "reasoning_steps": reasoning_steps,
+                    "verbose": self.verbose,
+                    "max_iterations": 10
+                }
+                if response_format:
+                    chat_kwargs["response_format"] = response_format
+                
+                final_response = self._openai_client.chat_completion_with_tools(**chat_kwargs)
 
             return final_response
 
@@ -3852,8 +4029,23 @@ Your Goal: {self.goal}"""
                 display_error(f"Error in LLM chat: {e}")
                 return None
         else:
+            # Determine if we should use native structured output
+            schema_model = output_pydantic or output_json
+            use_native_format = False
+            response_format = None
+            
+            if schema_model and self._supports_native_structured_output():
+                # Model supports native structured output - build response_format
+                response_format = self._build_response_format(schema_model)
+                if response_format:
+                    use_native_format = True
+                    logging.debug(f"Agent {self.name} using native structured output with response_format")
+            
             # Use the new _build_messages helper method
-            messages, original_prompt = self._build_messages(prompt, temperature, output_json, output_pydantic)
+            messages, original_prompt = self._build_messages(
+                prompt, temperature, output_json, output_pydantic,
+                use_native_format=use_native_format
+            )
             
             # Store chat history length for potential rollback
             chat_history_length = len(self.chat_history)
@@ -3909,7 +4101,7 @@ Your Goal: {self.goal}"""
                                     agent_tools=agent_tools
                                 )
 
-                        response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps, stream=stream, task_name=task_name, task_description=task_description, task_id=task_id)
+                        response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps, stream=stream, task_name=task_name, task_description=task_description, task_id=task_id, response_format=response_format)
                         if not response:
                             # Rollback chat history on response failure
                             self.chat_history = self.chat_history[:chat_history_length]
@@ -4572,6 +4764,11 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
         if 'stream' not in kwargs:
             kwargs['stream'] = False
         
+        # Substitute dynamic variables ({{today}}, {{now}}, {{uuid}}, etc.)
+        if prompt and "{{" in prompt:
+            from praisonaiagents.utils.variables import substitute_variables
+            prompt = substitute_variables(prompt, {})
+        
         # Load history context
         self._load_history_context()
         
@@ -4759,6 +4956,11 @@ Write the complete compiled report:"""
         if prompt is None:
             prompt = self.instructions or "Hello"
         
+        # Substitute dynamic variables ({{today}}, {{now}}, {{uuid}}, etc.)
+        if prompt and "{{" in prompt:
+            from praisonaiagents.utils.variables import substitute_variables
+            prompt = substitute_variables(prompt, {})
+        
         # Load history from past sessions
         self._load_history_context()
         

src/praisonai-agents/praisonaiagents/agent/context_agent.py

@@ -13,17 +13,88 @@ This class implements proper Context Engineering principles following the PRD te
 
 import os
 import json
-import asyncio
-import glob
-import subprocess
-import ast
 import logging
 from datetime import datetime
 from pathlib import Path
-from typing import Optional, Any, Dict, Union, List
+from typing import Optional, Any, Dict, Union, List, TYPE_CHECKING
+
+# Lazy imports for performance - these are only loaded when needed
+_subprocess = None
+_glob = None
+_ast = None
+_asyncio = None
+
+def _get_subprocess():
+    """Lazy import subprocess to avoid import-time overhead."""
+    global _subprocess
+    if _subprocess is None:
+        import subprocess
+        _subprocess = subprocess
+    return _subprocess
+
+def _get_glob():
+    """Lazy import glob to avoid import-time overhead."""
+    global _glob
+    if _glob is None:
+        import glob
+        _glob = glob
+    return _glob
+
+def _get_ast():
+    """Lazy import ast to avoid import-time overhead."""
+    global _ast
+    if _ast is None:
+        import ast
+        _ast = ast
+    return _ast
+
+def _get_asyncio():
+    """Lazy import asyncio to avoid import-time overhead."""
+    global _asyncio
+    if _asyncio is None:
+        import asyncio
+        _asyncio = asyncio
+    return _asyncio
+
+async def _async_subprocess_run(cmd: list, timeout: int = 60) -> tuple:
+    """
+    Run subprocess asynchronously (non-blocking).
+    
+    Args:
+        cmd: Command list to execute
+        timeout: Timeout in seconds
+        
+    Returns:
+        Tuple of (stdout, stderr, returncode)
+    """
+    asyncio = _get_asyncio()
+    try:
+        process = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await asyncio.wait_for(
+            process.communicate(),
+            timeout=timeout
+        )
+        return (
+            stdout.decode() if stdout else "",
+            stderr.decode() if stderr else "",
+            process.returncode
+        )
+    except asyncio.TimeoutError:
+        if process:
+            process.kill()
+            await process.wait()
+        raise
+    except FileNotFoundError:
+        raise
+
 from ..agent.agent import Agent
 from ..task import Task
 
+
 class ContextAgent(Agent):
     """
     Advanced Context Engineering Agent - Comprehensive context generation for AI coding assistants.
@@ -69,6 +140,7 @@ class ContextAgent(Agent):
         deeply and generate actionable implementation guidance."""
         
         # Initialize parent Agent
+        # Note: self.verbose is set by parent Agent based on output= parameter
         super().__init__(
             name=default_name,
             role=default_role,
@@ -102,6 +174,34 @@ class ContextAgent(Agent):
         if self.auto_analyze and self.project_path:
             self._perform_context_engineering_analysis()
 
+    def _log(self, message: str, level: str = "info"):
+        """
+        Configurable logging that respects verbose flag.
+        
+        Args:
+            message: Message to log
+            level: Log level (info, debug, warning, error)
+        """
+        if level == "debug":
+            if hasattr(self, 'logger'):
+                self.logger.debug(message)
+            if self.debug_mode:
+                print(f"🐛 {message}")
+        elif level == "warning":
+            if hasattr(self, 'logger'):
+                self.logger.warning(message)
+            if self.verbose:
+                print(f"⚠️ {message}")
+        elif level == "error":
+            if hasattr(self, 'logger'):
+                self.logger.error(message)
+            print(f"❌ {message}")  # Always show errors
+        else:  # info
+            if hasattr(self, 'logger'):
+                self.logger.info(message)
+            if self.verbose:
+                print(message)
+
     def setup_logging(self):
         """Setup comprehensive logging based on debug mode."""
         try:
@@ -474,7 +574,7 @@ codebase style and architecture following PRD template principles."""
         """Run gitingest analysis on the codebase."""
         try:
             # Try to run gitingest command
-            result = subprocess.run(
+            result = _get_subprocess().run(
                 ["gitingest", project_path, "--output", "-"],
                 capture_output=True,
                 text=True,
@@ -487,7 +587,7 @@ codebase style and architecture following PRD template principles."""
                 print(f"  ⚠️ Gitingest command failed: {result.stderr}")
                 return None
                 
-        except subprocess.TimeoutExpired:
+        except _get_subprocess().TimeoutExpired:
             print("  ⚠️ Gitingest analysis timed out")
             return None
         except FileNotFoundError:
@@ -1030,7 +1130,7 @@ detailed, actionable implementation steps."""
     def _get_filtered_files(self, project_path: str, pattern: str, max_files: int) -> List[str]:
         """Get filtered files excluding unwanted directories following PRD principles."""
         try:
-            files = glob.glob(os.path.join(project_path, "**", pattern), recursive=True)
+            files = _get_glob().glob(os.path.join(project_path, "**", pattern), recursive=True)
             
             # Filter out unwanted files following PRD methodology
             filtered = []
@@ -1403,7 +1503,7 @@ Return in the exact format specified in your instructions."""
         """Get basic repository file structure."""
         try:
             # Try using gitingest with structure-only flag first
-            result = subprocess.run(
+            result = _get_subprocess().run(
                 ["gitingest", github_url, "--tree-only", "--output", "-"],
                 capture_output=True,
                 text=True,
@@ -1421,7 +1521,7 @@ Return in the exact format specified in your instructions."""
                 print(f"    ⚠️ Gitingest tree failed, using GitHub API...")
                 return self._get_structure_fallback(github_url)
                 
-        except (subprocess.TimeoutExpired, FileNotFoundError) as e:
+        except (_get_subprocess().TimeoutExpired, FileNotFoundError) as e:
             print(f"    ⚠️ Gitingest not available: {e}, using GitHub API...")
             return self._get_structure_fallback(github_url)
 
@@ -1626,7 +1726,7 @@ Focus on creating clear, structured metadata that will help with intelligent fil
                 print("    ⚠️ gitingest Python package not available, using command line...")
                 
                 # Fallback to command line gitingest for actual content
-                result = subprocess.run(
+                result = _get_subprocess().run(
                     ["gitingest", github_url, "--output", "-"],
                     capture_output=True,
                     text=True,
@@ -1946,7 +2046,7 @@ Since these files were pre-selected for relevance, explain how each contributes
             except ImportError:
                 print(f"    ⚠️ gitingest Python package not available, trying command line...")
                 # Fallback to command line gitingest
-                result = subprocess.run(
+                result = _get_subprocess().run(
                     ["gitingest", github_url, "--output", "-"],
                     capture_output=True,
                     text=True,
@@ -1959,7 +2059,7 @@ Since these files were pre-selected for relevance, explain how each contributes
                     print(f"    ⚠️ Gitingest command failed: {result.stderr}")
                     return None
                     
-        except subprocess.TimeoutExpired:
+        except _get_subprocess().TimeoutExpired:
             print("    ⚠️ Gitingest analysis timed out")
             return None
         except FileNotFoundError:
@@ -2282,6 +2382,105 @@ This PRP provides:
                 "success": False
             }
 
+    # ============================================================================
+    # P1: Protocol-Compatible Aliases (ContextEngineerProtocol compliance)
+    # ============================================================================
+    
+    def analyze_codebase(self, project_path: str) -> Dict[str, Any]:
+        """
+        Protocol-compatible alias for analyze_codebase_with_gitingest.
+        
+        Args:
+            project_path: Path to the project directory
+            
+        Returns:
+            Analysis results including patterns, architecture, conventions
+        """
+        return self.analyze_codebase_with_gitingest(project_path)
+    
+    def generate_prp(
+        self, 
+        feature_request: str, 
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Protocol-compatible alias for generate_comprehensive_prp.
+        
+        Args:
+            feature_request: Description of the feature to implement
+            context_analysis: Optional codebase analysis to include
+            
+        Returns:
+            Comprehensive PRP document as string
+        """
+        return self.generate_comprehensive_prp(feature_request, context_analysis)
+    
+    def create_implementation_blueprint(
+        self,
+        feature_request: str,
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Protocol-compatible alias for build_implementation_blueprint.
+        
+        Args:
+            feature_request: Description of the feature
+            context_analysis: Optional codebase analysis
+            
+        Returns:
+            Blueprint with implementation steps, files to modify, etc.
+        """
+        return self.build_implementation_blueprint(feature_request, context_analysis)
+
+    # ============================================================================
+    # P3: Async Methods (ContextEngineerProtocol compliance)
+    # ============================================================================
+    
+    async def aanalyze_codebase(self, project_path: str) -> Dict[str, Any]:
+        """
+        Async version of analyze_codebase.
+        
+        Uses async subprocess for non-blocking execution.
+        """
+        asyncio = _get_asyncio()
+        # Run sync method in executor for now (subprocess calls are blocking)
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            self.analyze_codebase_with_gitingest,
+            project_path
+        )
+    
+    async def agenerate_prp(
+        self,
+        feature_request: str,
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Async version of generate_prp.
+        """
+        asyncio = _get_asyncio()
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            lambda: self.generate_comprehensive_prp(feature_request, context_analysis)
+        )
+    
+    async def acreate_implementation_blueprint(
+        self,
+        feature_request: str,
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Async version of create_implementation_blueprint.
+        """
+        asyncio = _get_asyncio()
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            lambda: self.build_implementation_blueprint(feature_request, context_analysis)
+        )
+
 
 def create_context_agent(llm: Optional[Union[str, Any]] = None, **kwargs) -> ContextAgent:
     """

src/praisonai-agents/praisonaiagents/agent/deep_research_agent.py

@@ -487,22 +487,85 @@ You are a professional research analyst. When conducting research:
             raw_response=response
         )
     
-    def _parse_gemini_response(self, interaction: Any) -> DeepResearchResponse:
-        """Parse Gemini Deep Research API response."""
+    def _parse_gemini_response(
+        self, 
+        interaction: Any, 
+        fallback_text: str = "",
+        fallback_reasoning: Optional[List[ReasoningStep]] = None
+    ) -> DeepResearchResponse:
+        """Parse Gemini Deep Research API response.
+        
+        Args:
+            interaction: The Gemini interaction object
+            fallback_text: Accumulated text from streaming (used if outputs parsing fails)
+            fallback_reasoning: Reasoning steps collected during streaming
+        """
         report = ""
         citations = []
-        reasoning_steps = []
+        reasoning_steps = fallback_reasoning or []
         
-        # Get the final output
+        # Try multiple attribute paths for Gemini output structure
+        # Path 1: Direct outputs with text attribute
         if hasattr(interaction, 'outputs') and interaction.outputs:
             last_output = interaction.outputs[-1]
-            if hasattr(last_output, 'text'):
+            if hasattr(last_output, 'text') and last_output.text:
                 report = last_output.text
             elif hasattr(last_output, 'content'):
-                report = str(last_output.content)
+                # Gemini nested structure: content.parts[0].text
+                content = last_output.content
+                if hasattr(content, 'parts') and content.parts:
+                    first_part = content.parts[0]
+                    if hasattr(first_part, 'text'):
+                        report = first_part.text
+                    else:
+                        report = str(first_part)
+                elif hasattr(content, 'text'):
+                    report = content.text
+                else:
+                    report = str(content)
         
-        # Gemini doesn't provide structured citations in the same way
-        # but we can try to extract them from annotations if available
+        # Path 2: Direct result attribute
+        if not report and hasattr(interaction, 'result'):
+            result = interaction.result
+            if hasattr(result, 'text'):
+                report = result.text
+            else:
+                report = str(result)
+        
+        # Path 3: Response attribute  
+        if not report and hasattr(interaction, 'response'):
+            resp = interaction.response
+            if hasattr(resp, 'text'):
+                report = resp.text
+            elif hasattr(resp, 'content'):
+                report = str(resp.content)
+        
+        # Path 4: Fallback to streamed content (critical fix)
+        if not report and fallback_text:
+            report = fallback_text
+            if self.verbose:
+                self.logger.debug("Using fallback streamed text for report")
+        
+        # Try to extract citations from grounding metadata
+        if hasattr(interaction, 'outputs') and interaction.outputs:
+            for output in interaction.outputs:
+                if hasattr(output, 'grounding_metadata'):
+                    metadata = output.grounding_metadata
+                    if hasattr(metadata, 'grounding_chunks'):
+                        for chunk in metadata.grounding_chunks:
+                            if hasattr(chunk, 'web') and chunk.web:
+                                citations.append(Citation(
+                                    title=getattr(chunk.web, 'title', ''),
+                                    url=getattr(chunk.web, 'uri', ''),
+                                ))
+        
+        # Log warning if report is empty
+        if not report:
+            self.logger.warning(
+                "Gemini response parsing returned empty report. "
+                f"Interaction ID: {getattr(interaction, 'id', 'unknown')}, "
+                f"Status: {getattr(interaction, 'status', 'unknown')}"
+            )
         
         return DeepResearchResponse(
             report=report,
@@ -878,10 +941,32 @@ You are a professional research analyst. When conducting research:
                         print("\n\n" + "=" * 60)
                         print("✅ Research Complete")
                         print("=" * 60 + "\n")
-                    # Get final interaction for full response
+                    
+                    # Poll until interaction status is actually 'completed'
+                    # (workaround for timing issue where GET returns stale status)
                     if interaction_id:
-                        final_interaction = self.gemini_client.interactions.get(interaction_id)
-                        return self._parse_gemini_response(final_interaction)
+                        max_poll_attempts = 30  # 30 seconds max
+                        for attempt in range(max_poll_attempts):
+                            final_interaction = self.gemini_client.interactions.get(interaction_id)
+                            if final_interaction.status == "completed":
+                                return self._parse_gemini_response(
+                                    final_interaction, 
+                                    fallback_text=final_text,
+                                    fallback_reasoning=reasoning_steps
+                                )
+                            elif final_interaction.status in ["failed", "cancelled"]:
+                                raise RuntimeError(f"Research {final_interaction.status}")
+                            time.sleep(1)
+                        
+                        # If still not completed, use fallback
+                        if self.verbose:
+                            self.logger.warning("Interaction not completed after polling, using streamed content")
+                        return DeepResearchResponse(
+                            report=final_text,
+                            reasoning_steps=reasoning_steps,
+                            provider="gemini",
+                            interaction_id=interaction_id
+                        )
                 
                 elif chunk.event_type == "error":
                     error_msg = getattr(chunk, 'error', 'Unknown streaming error')
@@ -949,7 +1034,11 @@ You are a professional research analyst. When conducting research:
                         if self.verbose:
                             print("\n\n✅ Research Complete (resumed)")
                         final_interaction = self.gemini_client.interactions.get(interaction_id)
-                        return self._parse_gemini_response(final_interaction)
+                        return self._parse_gemini_response(
+                            final_interaction,
+                            fallback_text=accumulated_text,
+                            fallback_reasoning=reasoning_steps
+                        )
                 
             except Exception as e:
                 retry_count += 1
@@ -963,7 +1052,11 @@ You are a professional research analyst. When conducting research:
         while True:
             interaction = self.gemini_client.interactions.get(interaction_id)
             if interaction.status == "completed":
-                return self._parse_gemini_response(interaction)
+                return self._parse_gemini_response(
+                    interaction,
+                    fallback_text=accumulated_text,
+                    fallback_reasoning=reasoning_steps
+                )
             elif interaction.status in ["failed", "cancelled"]:
                 raise RuntimeError(f"Research {interaction.status}")
             time.sleep(self.poll_interval)

src/praisonai-agents/praisonaiagents/agent/protocols.py

@@ -142,10 +142,97 @@ class FullAgentProtocol(RunnableAgentProtocol, ToolAwareAgentProtocol, MemoryAwa
     pass
 
 
+@runtime_checkable
+class ContextEngineerProtocol(Protocol):
+    """
+    Protocol for Context Engineering agents.
+    
+    Defines the interface for agents that perform codebase analysis,
+    PRP (Product Requirements Prompt) generation, and implementation planning.
+    
+    This protocol enables:
+    - Lightweight core with protocol definition only
+    - Heavy implementations in wrapper/tools
+    - Easy mocking for tests
+    
+    Example:
+        ```python
+        class MockContextEngineer:
+            def analyze_codebase(self, path: str) -> Dict[str, Any]:
+                return {"project": path, "patterns": []}
+            
+            def generate_prp(self, request: str, analysis: Dict = None) -> str:
+                return f"PRP for: {request}"
+            
+            async def aanalyze_codebase(self, path: str) -> Dict[str, Any]:
+                return await self.analyze_codebase(path)
+        ```
+    """
+    
+    def analyze_codebase(self, project_path: str) -> Dict[str, Any]:
+        """
+        Analyze a codebase and extract patterns, structure, and conventions.
+        
+        Args:
+            project_path: Path to the project directory
+            
+        Returns:
+            Analysis results including patterns, architecture, conventions
+        """
+        ...
+    
+    def generate_prp(
+        self, 
+        feature_request: str, 
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Generate a Product Requirements Prompt (PRP) for a feature.
+        
+        Args:
+            feature_request: Description of the feature to implement
+            context_analysis: Optional codebase analysis to include
+            
+        Returns:
+            Comprehensive PRP document as string
+        """
+        ...
+    
+    def create_implementation_blueprint(
+        self,
+        feature_request: str,
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Create a step-by-step implementation blueprint.
+        
+        Args:
+            feature_request: Description of the feature
+            context_analysis: Optional codebase analysis
+            
+        Returns:
+            Blueprint with implementation steps, files to modify, etc.
+        """
+        ...
+    
+    async def aanalyze_codebase(self, project_path: str) -> Dict[str, Any]:
+        """Async version of analyze_codebase."""
+        ...
+    
+    async def agenerate_prp(
+        self,
+        feature_request: str,
+        context_analysis: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """Async version of generate_prp."""
+        ...
+
+
 __all__ = [
     'AgentProtocol',
     'RunnableAgentProtocol', 
     'ToolAwareAgentProtocol',
     'MemoryAwareAgentProtocol',
     'FullAgentProtocol',
+    'ContextEngineerProtocol',
 ]

src/praisonai-agents/praisonaiagents/config/presets.py

@@ -334,16 +334,11 @@ WORKFLOW_STEP_EXECUTION_PRESETS: Dict[str, Dict[str, Any]] = {
 
 
 # =============================================================================
-# Workflow Output Presets
+# Workflow Output Presets - DEPRECATED, use OUTPUT_PRESETS instead (DRY)
 # =============================================================================
-
-WORKFLOW_OUTPUT_PRESETS: Dict[str, Dict[str, Any]] = {
-    "silent": {"verbose": False, "stream": False},
-    "minimal": {"verbose": False, "stream": True},
-    "normal": {"verbose": False, "stream": True},
-    "verbose": {"verbose": True, "stream": True},
-    "debug": {"verbose": True, "stream": True},
-}
+# Workflows now use the same OUTPUT_PRESETS as Agent for consistency.
+# This alias is kept for backward compatibility only.
+WORKFLOW_OUTPUT_PRESETS = OUTPUT_PRESETS
 
 
 # =============================================================================

src/praisonai-agents/praisonaiagents/context/__init__.py

@@ -269,4 +269,14 @@ def __getattr__(name: str):
         from . import artifacts
         return getattr(artifacts, name)
     
+    # Session Context Tracking (Agno pattern)
+    if name in ("SessionContextTracker", "SessionState"):
+        from . import session_tracker
+        return getattr(session_tracker, name)
+    
+    # Context Aggregation (CrewAI pattern)
+    if name in ("ContextAggregator", "AggregatedContext", "create_aggregator_from_config"):
+        from . import aggregator
+        return getattr(aggregator, name)
+    
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai-agents/praisonaiagents/context/aggregator.py

@@ -0,0 +1,319 @@
+"""
+Context Aggregator for PraisonAI Agents.
+
+Aggregates context from multiple memory sources concurrently.
+Inspired by CrewAI's ContextualMemory pattern.
+"""
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AggregatedContext:
+    """Result of context aggregation."""
+    context: str = ""
+    sources_used: List[str] = field(default_factory=list)
+    tokens_used: int = 0
+    fetch_times: Dict[str, float] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "context": self.context[:500] + "..." if len(self.context) > 500 else self.context,
+            "sources_used": self.sources_used,
+            "tokens_used": self.tokens_used,
+            "fetch_times": self.fetch_times,
+        }
+
+
+class ContextAggregator:
+    """
+    Aggregates context from multiple memory sources concurrently.
+    
+    Implements the CrewAI ContextualMemory pattern for PraisonAI.
+    Fetches from memory, knowledge, and RAG in parallel and merges results.
+    
+    Example:
+        aggregator = ContextAggregator()
+        
+        # Register sources
+        aggregator.register_source("memory", memory.search)
+        aggregator.register_source("knowledge", knowledge.search)
+        aggregator.register_source("rag", rag.retrieve)
+        
+        # Aggregate context
+        result = await aggregator.aggregate("user query", max_tokens=4000)
+        print(result.context)
+    """
+    
+    def __init__(
+        self,
+        max_tokens: int = 4000,
+        separator: str = "\n\n---\n\n",
+        include_source_labels: bool = True,
+    ):
+        """
+        Initialize context aggregator.
+        
+        Args:
+            max_tokens: Maximum tokens for aggregated context
+            separator: Separator between source contexts
+            include_source_labels: Whether to label each source section
+        """
+        self.max_tokens = max_tokens
+        self.separator = separator
+        self.include_source_labels = include_source_labels
+        
+        # Registered context sources: name -> (fetch_fn, priority)
+        self._sources: Dict[str, Tuple[Callable, int]] = {}
+    
+    def _estimate_tokens(self, text: str) -> int:
+        """Estimate token count (4 chars per token heuristic)."""
+        return len(text) // 4 + 1
+    
+    def register_source(
+        self,
+        name: str,
+        fetch_fn: Callable[[str], Union[str, List[Dict[str, Any]]]],
+        priority: int = 50,
+    ) -> None:
+        """
+        Register a context source.
+        
+        Args:
+            name: Source name (e.g., "memory", "knowledge", "rag")
+            fetch_fn: Function that takes query and returns context
+            priority: Priority (lower = first, 0-100)
+        """
+        self._sources[name] = (fetch_fn, priority)
+    
+    def unregister_source(self, name: str) -> None:
+        """Remove a registered source."""
+        self._sources.pop(name, None)
+    
+    @property
+    def sources(self) -> List[str]:
+        """Get list of registered source names."""
+        return list(self._sources.keys())
+    
+    async def _fetch_source(
+        self,
+        name: str,
+        query: str,
+    ) -> Tuple[str, Optional[str], float]:
+        """
+        Fetch context from a single source.
+        
+        Returns: (context_str, source_name, fetch_time)
+        """
+        import time
+        
+        fetch_fn, _ = self._sources.get(name, (None, 0))
+        if not fetch_fn:
+            return "", None, 0.0
+        
+        start = time.time()
+        try:
+            # Call the fetch function
+            if asyncio.iscoroutinefunction(fetch_fn):
+                result = await fetch_fn(query)
+            else:
+                # Run sync function in executor
+                loop = asyncio.get_event_loop()
+                result = await loop.run_in_executor(None, fetch_fn, query)
+            
+            # Convert result to string
+            if isinstance(result, str):
+                context = result
+            elif isinstance(result, list):
+                # Handle list of results (e.g., from memory search)
+                parts = []
+                for item in result:
+                    if isinstance(item, dict):
+                        text = item.get("text") or item.get("memory") or item.get("content", "")
+                        if text:
+                            parts.append(str(text))
+                    elif isinstance(item, str):
+                        parts.append(item)
+                context = "\n".join(parts)
+            else:
+                context = str(result) if result else ""
+            
+            elapsed = time.time() - start
+            return context, name, elapsed
+            
+        except Exception as e:
+            logger.warning(f"Failed to fetch from {name}: {e}")
+            return "", None, time.time() - start
+    
+    async def aggregate(
+        self,
+        query: str,
+        sources: Optional[List[str]] = None,
+        max_tokens: Optional[int] = None,
+    ) -> AggregatedContext:
+        """
+        Aggregate context from multiple sources concurrently.
+        
+        Args:
+            query: Search query
+            sources: Specific sources to use (default: all registered)
+            max_tokens: Override max tokens
+            
+        Returns:
+            AggregatedContext with merged results
+        """
+        max_tokens = max_tokens or self.max_tokens
+        sources_to_use = sources or list(self._sources.keys())
+        
+        if not sources_to_use:
+            return AggregatedContext()
+        
+        # Filter to registered sources only
+        sources_to_use = [s for s in sources_to_use if s in self._sources]
+        
+        if not sources_to_use:
+            return AggregatedContext()
+        
+        # Fetch from all sources concurrently
+        tasks = [
+            self._fetch_source(name, query)
+            for name in sources_to_use
+        ]
+        
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        
+        # Process results
+        contexts: List[Tuple[str, str, float, int]] = []  # (context, name, time, priority)
+        fetch_times: Dict[str, float] = {}
+        
+        for result in results:
+            if isinstance(result, Exception):
+                logger.warning(f"Source fetch failed: {result}")
+                continue
+            
+            context, name, elapsed = result
+            if context and name:
+                _, priority = self._sources.get(name, (None, 50))
+                contexts.append((context, name, elapsed, priority))
+                fetch_times[name] = elapsed
+        
+        # Sort by priority (lower = first)
+        contexts.sort(key=lambda x: x[3])
+        
+        # Merge and truncate to fit token budget
+        merged_parts = []
+        total_tokens = 0
+        sources_used = []
+        separator_tokens = self._estimate_tokens(self.separator)
+        
+        for context, name, _, _ in contexts:
+            context_tokens = self._estimate_tokens(context)
+            
+            # Check if we can fit this context
+            if total_tokens + context_tokens + separator_tokens > max_tokens:
+                # Try to fit partial
+                remaining = max_tokens - total_tokens - separator_tokens
+                if remaining > 100:
+                    # Truncate context
+                    chars_to_keep = int(remaining * 4 * 0.9)
+                    truncated = context[:chars_to_keep] + "..."
+                    
+                    if self.include_source_labels:
+                        merged_parts.append(f"**{name}**:\n{truncated}")
+                    else:
+                        merged_parts.append(truncated)
+                    
+                    sources_used.append(name)
+                    total_tokens += self._estimate_tokens(truncated)
+                break
+            
+            if self.include_source_labels:
+                merged_parts.append(f"**{name}**:\n{context}")
+            else:
+                merged_parts.append(context)
+            
+            sources_used.append(name)
+            total_tokens += context_tokens + separator_tokens
+        
+        merged_context = self.separator.join(merged_parts)
+        
+        return AggregatedContext(
+            context=merged_context,
+            sources_used=sources_used,
+            tokens_used=total_tokens,
+            fetch_times=fetch_times,
+        )
+    
+    def aggregate_sync(
+        self,
+        query: str,
+        sources: Optional[List[str]] = None,
+        max_tokens: Optional[int] = None,
+    ) -> AggregatedContext:
+        """
+        Synchronous version of aggregate.
+        
+        Args:
+            query: Search query
+            sources: Specific sources to use
+            max_tokens: Override max tokens
+            
+        Returns:
+            AggregatedContext with merged results
+        """
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Create a new loop for sync execution
+                import concurrent.futures
+                with concurrent.futures.ThreadPoolExecutor() as pool:
+                    future = pool.submit(
+                        asyncio.run,
+                        self.aggregate(query, sources, max_tokens)
+                    )
+                    return future.result()
+            else:
+                return loop.run_until_complete(
+                    self.aggregate(query, sources, max_tokens)
+                )
+        except RuntimeError:
+            # No event loop, create one
+            return asyncio.run(self.aggregate(query, sources, max_tokens))
+
+
+def create_aggregator_from_config(
+    memory = None,
+    knowledge = None,
+    rag = None,
+    max_tokens: int = 4000,
+) -> ContextAggregator:
+    """
+    Create a ContextAggregator with standard sources.
+    
+    Args:
+        memory: Memory instance with search() method
+        knowledge: Knowledge instance with search() method
+        rag: RAG instance with retrieve() method
+        max_tokens: Maximum tokens for aggregated context
+        
+    Returns:
+        Configured ContextAggregator
+    """
+    aggregator = ContextAggregator(max_tokens=max_tokens)
+    
+    if memory and hasattr(memory, 'search_short_term'):
+        aggregator.register_source("memory", memory.search_short_term, priority=10)
+    
+    if knowledge and hasattr(knowledge, 'search'):
+        aggregator.register_source("knowledge", knowledge.search, priority=20)
+    
+    if rag and hasattr(rag, 'retrieve'):
+        aggregator.register_source("rag", rag.retrieve, priority=30)
+    
+    return aggregator

src/praisonai-agents/praisonaiagents/context/fast/__init__.py

@@ -9,6 +9,12 @@ Key features:
 - Limited serial turns (max 4) for fast response
 - Restricted tool set (grep, glob, read) for safety
 - Returns files + line ranges (no summarization)
+
+Optional optimizations (no performance impact when disabled):
+- Ripgrep backend: Use `search_backend='ripgrep'` for faster grep
+- Async file I/O: Uses aiofiles if installed
+- Incremental indexing: Use `enable_indexing=True`
+- Context compression: Use `compression='truncate'` or `compression='smart'`
 """
 
 # Lazy imports to avoid impacting package load time
@@ -29,6 +35,28 @@ def __getattr__(name):
     elif name == "FastContextAgent":
         from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
         return FastContextAgent
+    # New optimization exports
+    elif name == "get_search_backend":
+        from praisonaiagents.context.fast.search_backends import get_search_backend
+        return get_search_backend
+    elif name == "PythonSearchBackend":
+        from praisonaiagents.context.fast.search_backends import PythonSearchBackend
+        return PythonSearchBackend
+    elif name == "RipgrepBackend":
+        from praisonaiagents.context.fast.search_backends import RipgrepBackend
+        return RipgrepBackend
+    elif name == "FileIndex":
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        return FileIndex
+    elif name == "get_compressor":
+        from praisonaiagents.context.fast.compressor import get_compressor
+        return get_compressor
+    elif name == "TruncateCompressor":
+        from praisonaiagents.context.fast.compressor import TruncateCompressor
+        return TruncateCompressor
+    elif name == "SmartCompressor":
+        from praisonaiagents.context.fast.compressor import SmartCompressor
+        return SmartCompressor
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 __all__ = [
@@ -37,4 +65,13 @@ __all__ = [
     "FileMatch",
     "LineRange",
     "FastContextAgent",
+    # Optimization exports
+    "get_search_backend",
+    "PythonSearchBackend", 
+    "RipgrepBackend",
+    "FileIndex",
+    "get_compressor",
+    "TruncateCompressor",
+    "SmartCompressor",
 ]
+

src/praisonai-agents/praisonaiagents/context/fast/async_file_ops.py

@@ -0,0 +1,159 @@
+"""
+Async file operations for FastContext.
+
+Provides optional aiofiles integration with graceful fallback:
+- async_read_file: Async file reading (uses aiofiles if available)
+- async_read_lines: Async line reading with range support
+
+Design principles:
+- Lazy import of aiofiles (optional dependency)
+- Graceful fallback to sync operations in executor
+- No import-time overhead
+"""
+
+import asyncio
+import logging
+from typing import Optional, List
+
+logger = logging.getLogger(__name__)
+
+# Lazy availability check
+_AIOFILES_AVAILABLE: Optional[bool] = None
+
+
+def _check_aiofiles() -> bool:
+    """Lazy check for aiofiles availability.
+    
+    Returns:
+        True if aiofiles is installed and importable
+    """
+    global _AIOFILES_AVAILABLE
+    if _AIOFILES_AVAILABLE is None:
+        try:
+            import aiofiles  # noqa: F401
+            _AIOFILES_AVAILABLE = True
+            logger.debug("aiofiles available")
+        except ImportError:
+            _AIOFILES_AVAILABLE = False
+            logger.debug("aiofiles not available, using sync fallback")
+    return _AIOFILES_AVAILABLE
+
+
+def _sync_read_file(
+    path: str,
+    encoding: str = 'utf-8',
+    errors: str = 'ignore'
+) -> str:
+    """Sync file read for fallback."""
+    with open(path, 'r', encoding=encoding, errors=errors) as f:
+        return f.read()
+
+
+def _sync_read_lines(
+    path: str,
+    start_line: Optional[int] = None,
+    end_line: Optional[int] = None,
+    encoding: str = 'utf-8',
+    errors: str = 'ignore'
+) -> List[str]:
+    """Sync line read for fallback."""
+    with open(path, 'r', encoding=encoding, errors=errors) as f:
+        lines = f.readlines()
+    
+    if start_line is not None and end_line is not None:
+        # Convert to 0-indexed
+        start_idx = max(0, start_line - 1)
+        end_idx = min(len(lines), end_line)
+        return [line.rstrip('\n\r') for line in lines[start_idx:end_idx]]
+    elif start_line is not None:
+        start_idx = max(0, start_line - 1)
+        return [line.rstrip('\n\r') for line in lines[start_idx:]]
+    elif end_line is not None:
+        end_idx = min(len(lines), end_line)
+        return [line.rstrip('\n\r') for line in lines[:end_idx]]
+    else:
+        return [line.rstrip('\n\r') for line in lines]
+
+
+async def async_read_file(
+    path: str,
+    encoding: str = 'utf-8',
+    errors: str = 'ignore'
+) -> str:
+    """Read file asynchronously.
+    
+    Uses aiofiles if available, otherwise runs sync read in executor.
+    
+    Args:
+        path: File path to read
+        encoding: Text encoding
+        errors: Error handling ('ignore', 'strict', 'replace')
+        
+    Returns:
+        File content as string
+    """
+    if _check_aiofiles():
+        import aiofiles
+        async with aiofiles.open(path, 'r', encoding=encoding, errors=errors) as f:
+            return await f.read()
+    else:
+        # Fallback: run sync read in thread executor
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None, 
+            lambda: _sync_read_file(path, encoding, errors)
+        )
+
+
+async def async_read_lines(
+    path: str,
+    start_line: Optional[int] = None,
+    end_line: Optional[int] = None,
+    encoding: str = 'utf-8',
+    errors: str = 'ignore'
+) -> List[str]:
+    """Read specific lines from file asynchronously.
+    
+    Args:
+        path: File path to read
+        start_line: First line to read (1-indexed, inclusive)
+        end_line: Last line to read (1-indexed, inclusive)
+        encoding: Text encoding
+        errors: Error handling
+        
+    Returns:
+        List of lines (without line endings)
+    """
+    if _check_aiofiles():
+        import aiofiles
+        async with aiofiles.open(path, 'r', encoding=encoding, errors=errors) as f:
+            lines = await f.readlines()
+        
+        if start_line is not None and end_line is not None:
+            start_idx = max(0, start_line - 1)
+            end_idx = min(len(lines), end_line)
+            return [line.rstrip('\n\r') for line in lines[start_idx:end_idx]]
+        elif start_line is not None:
+            start_idx = max(0, start_line - 1)
+            return [line.rstrip('\n\r') for line in lines[start_idx:]]
+        elif end_line is not None:
+            end_idx = min(len(lines), end_line)
+            return [line.rstrip('\n\r') for line in lines[:end_idx]]
+        else:
+            return [line.rstrip('\n\r') for line in lines]
+    else:
+        # Fallback
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            lambda: _sync_read_lines(path, start_line, end_line, encoding, errors)
+        )
+
+
+def is_aiofiles_available() -> bool:
+    """Check if aiofiles is available for async operations.
+    
+    Returns:
+        True if aiofiles is installed
+    """
+    return _check_aiofiles()

src/praisonai-agents/praisonaiagents/context/fast/compressor.py

@@ -0,0 +1,254 @@
+"""
+Context compressor for FastContext.
+
+Provides token-aware context compression:
+- TruncateCompressor: Simple token-based truncation
+- SmartCompressor: Preserves important content
+- LLMLinguaCompressor: Optional LLMLingua integration (wrapper only)
+
+Design principles:
+- No external dependencies for basic compression
+- Lazy import for optional LLMLingua
+- Preserves most relevant content
+"""
+
+import logging
+from typing import Protocol, Optional, List
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+# Lazy availability check for llmlingua
+_LLMLINGUA_AVAILABLE: Optional[bool] = None
+
+
+def _check_llmlingua() -> bool:
+    """Lazy check for llmlingua availability."""
+    global _LLMLINGUA_AVAILABLE
+    if _LLMLINGUA_AVAILABLE is None:
+        try:
+            import llmlingua  # noqa: F401
+            _LLMLINGUA_AVAILABLE = True
+        except ImportError:
+            _LLMLINGUA_AVAILABLE = False
+    return _LLMLINGUA_AVAILABLE
+
+
+class ContextCompressor(Protocol):
+    """Protocol for context compression."""
+    
+    def compress(self, text: str, max_tokens: int) -> str:
+        """Compress text to fit within token budget.
+        
+        Args:
+            text: Text to compress
+            max_tokens: Maximum tokens allowed
+            
+        Returns:
+            Compressed text
+        """
+        ...
+
+
+@dataclass
+class CompressionConfig:
+    """Configuration for context compression."""
+    chars_per_token: float = 4.0  # Conservative estimate
+    preserve_start_lines: int = 5  # Preserve first N lines
+    preserve_end_lines: int = 3   # Preserve last N lines
+    truncation_marker: str = "\n\n... [content truncated] ...\n\n"
+
+
+class TruncateCompressor:
+    """Simple token-aware truncation compressor.
+    
+    Preserves start and end of content, truncates middle.
+    """
+    
+    def __init__(self, config: Optional[CompressionConfig] = None):
+        self.config = config or CompressionConfig()
+    
+    def estimate_tokens(self, text: str) -> int:
+        """Estimate token count."""
+        return int(len(text) / self.config.chars_per_token)
+    
+    def compress(self, text: str, max_tokens: int) -> str:
+        """Compress text by truncating middle.
+        
+        Args:
+            text: Text to compress
+            max_tokens: Maximum tokens allowed
+            
+        Returns:
+            Compressed text
+        """
+        current_tokens = self.estimate_tokens(text)
+        
+        if current_tokens <= max_tokens:
+            return text
+        
+        # Calculate how much to keep
+        max_chars = int(max_tokens * self.config.chars_per_token)
+        marker_len = len(self.config.truncation_marker)
+        available_chars = max_chars - marker_len
+        
+        if available_chars <= 0:
+            return text[:max_chars]
+        
+        # Split into lines for smarter truncation
+        lines = text.split('\n')
+        
+        # Preserve start and end lines
+        start_lines = lines[:self.config.preserve_start_lines]
+        end_lines = lines[-self.config.preserve_end_lines:] if self.config.preserve_end_lines > 0 else []
+        
+        start_text = '\n'.join(start_lines)
+        end_text = '\n'.join(end_lines)
+        
+        # Calculate remaining budget for middle
+        preserved_len = len(start_text) + len(end_text)
+        middle_budget = available_chars - preserved_len
+        
+        if middle_budget <= 0:
+            # Not enough room, just truncate at char level
+            half = available_chars // 2
+            return text[:half] + self.config.truncation_marker + text[-half:]
+        
+        # Take some middle content
+        middle_lines = lines[self.config.preserve_start_lines:-self.config.preserve_end_lines] if self.config.preserve_end_lines > 0 else lines[self.config.preserve_start_lines:]
+        middle_text = '\n'.join(middle_lines)
+        
+        if len(middle_text) <= middle_budget:
+            return text  # Actually fits
+        
+        # Truncate middle
+        middle_truncated = middle_text[:middle_budget // 2]
+        
+        return start_text + '\n' + middle_truncated + self.config.truncation_marker + end_text
+
+
+class SmartCompressor:
+    """Smart compressor that preserves important content.
+    
+    Uses heuristics to identify important lines (definitions, headers, etc.)
+    """
+    
+    IMPORTANT_PATTERNS = [
+        'def ', 'class ', 'async def ',  # Function/class definitions
+        'import ', 'from ',  # Imports
+        '# ', '"""', "'''",  # Comments/docstrings
+        'if __name__', 'return ',  # Key statements
+    ]
+    
+    def __init__(self, config: Optional[CompressionConfig] = None):
+        self.config = config or CompressionConfig()
+        self._fallback = TruncateCompressor(config)
+    
+    def estimate_tokens(self, text: str) -> int:
+        """Estimate token count."""
+        return int(len(text) / self.config.chars_per_token)
+    
+    def _score_line(self, line: str) -> int:
+        """Score a line's importance."""
+        score = 0
+        stripped = line.strip()
+        
+        for pattern in self.IMPORTANT_PATTERNS:
+            if stripped.startswith(pattern):
+                score += 10
+        
+        # Penalize empty lines
+        if not stripped:
+            score -= 5
+        
+        return score
+    
+    def compress(self, text: str, max_tokens: int) -> str:
+        """Compress text preserving important lines.
+        
+        Args:
+            text: Text to compress
+            max_tokens: Maximum tokens allowed
+            
+        Returns:
+            Compressed text
+        """
+        current_tokens = self.estimate_tokens(text)
+        
+        if current_tokens <= max_tokens:
+            return text
+        
+        lines = text.split('\n')
+        
+        if len(lines) <= 10:
+            # Too few lines for smart compression
+            return self._fallback.compress(text, max_tokens)
+        
+        # Score each line
+        scored_lines = [(i, line, self._score_line(line)) for i, line in enumerate(lines)]
+        
+        # Sort by score (keep highest)
+        sorted_by_score = sorted(scored_lines, key=lambda x: x[2], reverse=True)
+        
+        # Select lines until we hit token budget
+        max_chars = int(max_tokens * self.config.chars_per_token)
+        selected_indices = set()
+        current_chars = len(self.config.truncation_marker)
+        
+        for idx, line, score in sorted_by_score:
+            line_chars = len(line) + 1  # +1 for newline
+            if current_chars + line_chars <= max_chars:
+                selected_indices.add(idx)
+                current_chars += line_chars
+        
+        # Rebuild text in original order
+        result_lines = []
+        last_included = -1
+        
+        for i in range(len(lines)):
+            if i in selected_indices:
+                if last_included >= 0 and i - last_included > 1:
+                    result_lines.append("...")
+                result_lines.append(lines[i])
+                last_included = i
+        
+        return '\n'.join(result_lines)
+
+
+def get_compressor(compressor_type: str = "truncate") -> ContextCompressor:
+    """Get a context compressor by type.
+    
+    Args:
+        compressor_type: One of "truncate", "smart", "llmlingua"
+        
+    Returns:
+        ContextCompressor instance
+        
+    Raises:
+        ValueError: If compressor_type is invalid
+    """
+    if compressor_type == "truncate":
+        return TruncateCompressor()
+    
+    elif compressor_type == "smart":
+        return SmartCompressor()
+    
+    elif compressor_type == "llmlingua":
+        if not _check_llmlingua():
+            logger.warning("llmlingua not available, using smart compressor")
+            return SmartCompressor()
+        # LLMLingua integration would go here
+        # For now, fallback to smart
+        return SmartCompressor()
+    
+    else:
+        raise ValueError(f"Invalid compressor type: {compressor_type}. Must be 'truncate', 'smart', or 'llmlingua'")
+
+
+def is_llmlingua_available() -> bool:
+    """Check if llmlingua is available.
+    
+    Returns:
+        True if llmlingua is installed
+    """
+    return _check_llmlingua()

src/praisonai-agents/praisonaiagents/context/fast/fast_context.py

@@ -58,6 +58,7 @@ class FastContext:
     ENV_TIMEOUT = "FAST_CONTEXT_TIMEOUT"
     ENV_CACHE_ENABLED = "FAST_CONTEXT_CACHE"
     ENV_CACHE_TTL = "FAST_CONTEXT_CACHE_TTL"
+    ENV_SEARCH_BACKEND = "FAST_CONTEXT_BACKEND"
     
     def __init__(
         self,
@@ -68,7 +69,12 @@ class FastContext:
         timeout: Optional[float] = None,
         cache_enabled: Optional[bool] = None,
         cache_ttl: Optional[int] = None,
-        verbose: bool = False
+        verbose: bool = False,
+        # New optimization parameters
+        search_backend: Optional[str] = None,
+        enable_indexing: bool = False,
+        index_path: Optional[str] = None,
+        compression: Optional[str] = None
     ):
         """Initialize FastContext.
         
@@ -82,6 +88,7 @@ class FastContext:
             FAST_CONTEXT_TIMEOUT: Timeout in seconds (default: 30.0)
             FAST_CONTEXT_CACHE: Enable caching (default: true)
             FAST_CONTEXT_CACHE_TTL: Cache TTL in seconds (default: 300)
+            FAST_CONTEXT_BACKEND: Search backend (default: auto)
         
         Args:
             workspace_path: Root directory for searches (defaults to cwd)
@@ -92,6 +99,10 @@ class FastContext:
             cache_enabled: Whether to cache search results
             cache_ttl: Cache time-to-live in seconds
             verbose: If True, print debug information
+            search_backend: Search backend ("auto", "python", "ripgrep")
+            enable_indexing: Enable incremental file indexing
+            index_path: Custom path for index file
+            compression: Compression strategy ("truncate", "smart", None)
         """
         self.workspace_path = os.path.abspath(workspace_path or os.getcwd())
         
@@ -108,11 +119,20 @@ class FastContext:
         
         self.verbose = verbose
         
+        # Optimization settings (all optional with graceful fallback)
+        self.search_backend = search_backend or os.environ.get(self.ENV_SEARCH_BACKEND, "auto")
+        self.enable_indexing = enable_indexing
+        self.index_path = index_path
+        self.compression = compression
+        
         # Cache storage
         self._cache: Dict[str, Dict[str, Any]] = {}
         
-        # Agent instance (lazy loaded)
+        # Lazy-loaded components
         self._agent: Optional[FastContextAgent] = None
+        self._backend = None  # SearchBackend instance
+        self._index = None    # FileIndex instance
+        self._compressor = None  # ContextCompressor instance
     
     def _get_agent(self) -> FastContextAgent:
         """Get or create FastContextAgent instance."""
@@ -127,6 +147,31 @@ class FastContext:
             )
         return self._agent
     
+    def _get_backend(self):
+        """Get or create SearchBackend instance (lazy)."""
+        if self._backend is None:
+            from praisonaiagents.context.fast.search_backends import get_search_backend
+            self._backend = get_search_backend(self.search_backend, workspace_path=self.workspace_path)
+        return self._backend
+    
+    def _get_index(self):
+        """Get or create FileIndex instance (lazy)."""
+        if not self.enable_indexing:
+            return None
+        if self._index is None:
+            from praisonaiagents.context.fast.index_manager import FileIndex
+            self._index = FileIndex.load_or_create(self.workspace_path)
+        return self._index
+    
+    def _get_compressor(self):
+        """Get or create ContextCompressor instance (lazy)."""
+        if not self.compression:
+            return None
+        if self._compressor is None:
+            from praisonaiagents.context.fast.compressor import get_compressor
+            self._compressor = get_compressor(self.compression)
+        return self._compressor
+    
     def _get_cache_key(self, query: str, **kwargs) -> str:
         """Generate cache key for a query."""
         key_data = {

src/praisonai-agents/praisonaiagents/context/fast/index_manager.py

@@ -0,0 +1,279 @@
+"""
+Index manager for FastContext incremental indexing.
+
+Provides file indexing for faster repeat searches:
+- FileIndex: Lightweight file index with mtime tracking
+- Incremental updates: Only rescan changed files
+- Optional file watcher integration (watchfiles)
+
+Design principles:
+- Index stored as JSON for simplicity
+- Lazy loading of optional dependencies
+- No performance impact when disabled
+"""
+
+import os
+import json
+import hashlib
+import logging
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Dict, Optional, Set, List, Any
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+# Lazy availability check for watchfiles
+_WATCHFILES_AVAILABLE: Optional[bool] = None
+
+
+def _check_watchfiles() -> bool:
+    """Lazy check for watchfiles availability."""
+    global _WATCHFILES_AVAILABLE
+    if _WATCHFILES_AVAILABLE is None:
+        try:
+            import watchfiles  # noqa: F401
+            _WATCHFILES_AVAILABLE = True
+        except ImportError:
+            _WATCHFILES_AVAILABLE = False
+    return _WATCHFILES_AVAILABLE
+
+
+@dataclass
+class FileEntry:
+    """Entry for a single file in the index."""
+    path: str
+    mtime: float  # Modification time
+    size: int
+    content_hash: Optional[str] = None  # Optional for large files
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "FileEntry":
+        return cls(**data)
+
+
+@dataclass  
+class FileIndex:
+    """Lightweight file index for incremental search.
+    
+    Tracks file modification times to avoid rescanning
+    unchanged files on repeat searches.
+    
+    Attributes:
+        workspace_path: Root directory being indexed
+        index_version: Version for format compatibility
+    """
+    
+    INDEX_VERSION = 1
+    INDEX_FILENAME = ".fast_context_index.json"
+    
+    workspace_path: str
+    entries: Dict[str, FileEntry] = field(default_factory=dict)
+    index_version: int = INDEX_VERSION
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    
+    @property
+    def index_path(self) -> str:
+        """Path to the index file."""
+        return os.path.join(self.workspace_path, self.INDEX_FILENAME)
+    
+    def needs_rescan(self, filepath: str) -> bool:
+        """Check if file needs rescan based on mtime.
+        
+        Args:
+            filepath: Absolute path to file
+            
+        Returns:
+            True if file needs rescan (new, modified, or not indexed)
+        """
+        rel_path = os.path.relpath(filepath, self.workspace_path)
+        
+        if rel_path not in self.entries:
+            return True
+        
+        entry = self.entries[rel_path]
+        
+        try:
+            stat = os.stat(filepath)
+            return stat.st_mtime != entry.mtime or stat.st_size != entry.size
+        except OSError:
+            return True
+    
+    def update(self, filepath: str, content_hash: Optional[str] = None) -> None:
+        """Update index entry for file.
+        
+        Args:
+            filepath: Absolute path to file
+            content_hash: Optional hash of file content
+        """
+        rel_path = os.path.relpath(filepath, self.workspace_path)
+        
+        try:
+            stat = os.stat(filepath)
+            self.entries[rel_path] = FileEntry(
+                path=rel_path,
+                mtime=stat.st_mtime,
+                size=stat.st_size,
+                content_hash=content_hash
+            )
+            self.updated_at = datetime.now().isoformat()
+        except OSError as e:
+            logger.debug(f"Failed to update index for {filepath}: {e}")
+    
+    def remove(self, filepath: str) -> None:
+        """Remove file from index.
+        
+        Args:
+            filepath: Absolute or relative path
+        """
+        rel_path = os.path.relpath(filepath, self.workspace_path) if os.path.isabs(filepath) else filepath
+        self.entries.pop(rel_path, None)
+        self.updated_at = datetime.now().isoformat()
+    
+    def get_changed_files(self, filepaths: List[str]) -> List[str]:
+        """Get list of files that need rescan.
+        
+        Args:
+            filepaths: List of absolute file paths
+            
+        Returns:
+            List of files that need rescan
+        """
+        return [fp for fp in filepaths if self.needs_rescan(fp)]
+    
+    def save(self) -> bool:
+        """Save index to disk.
+        
+        Returns:
+            True if saved successfully
+        """
+        try:
+            data = {
+                "version": self.index_version,
+                "workspace": self.workspace_path,
+                "created_at": self.created_at,
+                "updated_at": self.updated_at,
+                "entries": {k: v.to_dict() for k, v in self.entries.items()}
+            }
+            
+            with open(self.index_path, 'w', encoding='utf-8') as f:
+                json.dump(data, f, indent=2)
+            
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to save index: {e}")
+            return False
+    
+    @classmethod
+    def load(cls, workspace_path: str) -> Optional["FileIndex"]:
+        """Load index from disk.
+        
+        Args:
+            workspace_path: Root directory
+            
+        Returns:
+            FileIndex if loaded, None if not found or invalid
+        """
+        index_path = os.path.join(workspace_path, cls.INDEX_FILENAME)
+        
+        if not os.path.exists(index_path):
+            return None
+        
+        try:
+            with open(index_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+            
+            # Version check
+            if data.get("version") != cls.INDEX_VERSION:
+                logger.info("Index version mismatch, will rebuild")
+                return None
+            
+            entries = {
+                k: FileEntry.from_dict(v) 
+                for k, v in data.get("entries", {}).items()
+            }
+            
+            return cls(
+                workspace_path=workspace_path,
+                entries=entries,
+                index_version=data.get("version", cls.INDEX_VERSION),
+                created_at=data.get("created_at", datetime.now().isoformat()),
+                updated_at=data.get("updated_at", datetime.now().isoformat())
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load index: {e}")
+            return None
+    
+    @classmethod
+    def load_or_create(cls, workspace_path: str) -> "FileIndex":
+        """Load existing index or create new one.
+        
+        Args:
+            workspace_path: Root directory
+            
+        Returns:
+            FileIndex instance
+        """
+        index = cls.load(workspace_path)
+        if index is None:
+            index = cls(workspace_path=workspace_path)
+        return index
+    
+    def clear(self) -> None:
+        """Clear all index entries."""
+        self.entries.clear()
+        self.updated_at = datetime.now().isoformat()
+    
+    def __len__(self) -> int:
+        return len(self.entries)
+
+
+class IndexWatcher:
+    """Optional file watcher for real-time index updates.
+    
+    Uses watchfiles if available, otherwise disabled.
+    """
+    
+    def __init__(self, index: FileIndex):
+        self.index = index
+        self._watching = False
+    
+    def is_available(self) -> bool:
+        """Check if file watching is available."""
+        return _check_watchfiles()
+    
+    async def start(self) -> None:
+        """Start watching for file changes."""
+        if not self.is_available():
+            logger.warning("watchfiles not available, file watching disabled")
+            return
+        
+        if self._watching:
+            return
+        
+        import watchfiles
+        self._watching = True
+        
+        async for changes in watchfiles.awatch(self.index.workspace_path):
+            for change_type, path in changes:
+                if change_type == watchfiles.Change.deleted:
+                    self.index.remove(path)
+                else:
+                    self.index.update(path)
+    
+    def stop(self) -> None:
+        """Stop watching."""
+        self._watching = False
+
+
+def is_watchfiles_available() -> bool:
+    """Check if watchfiles is available for file watching.
+    
+    Returns:
+        True if watchfiles is installed
+    """
+    return _check_watchfiles()

src/praisonai-agents/praisonaiagents/context/fast/search_backends.py

@@ -0,0 +1,413 @@
+"""
+Search backends for FastContext.
+
+Provides protocol-driven search backends:
+- PythonSearchBackend: Pure Python implementation (always available)
+- RipgrepBackend: Optional ripgrep subprocess (faster for large codebases)
+
+Design principles:
+- Lazy imports for optional dependencies
+- Graceful degradation if deps not installed
+- No performance impact on default behavior
+"""
+
+import logging
+from typing import Protocol, List, Dict, Any, Optional, runtime_checkable
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class SearchBackend(Protocol):
+    """Protocol for search backends.
+    
+    All backends must implement these methods to be usable
+    with FastContext.
+    """
+    
+    def grep(
+        self,
+        path: str,
+        pattern: str,
+        is_regex: bool = False,
+        case_sensitive: bool = False,
+        include_patterns: Optional[List[str]] = None,
+        exclude_patterns: Optional[List[str]] = None,
+        max_results: int = 100,
+        context_lines: int = 0
+    ) -> List[Dict[str, Any]]:
+        """Search for pattern in files.
+        
+        Args:
+            path: Directory or file to search
+            pattern: Search pattern
+            is_regex: If True, treat pattern as regex
+            case_sensitive: If True, case sensitive search
+            include_patterns: Glob patterns to include
+            exclude_patterns: Glob patterns to exclude
+            max_results: Maximum results to return
+            context_lines: Lines of context around matches
+            
+        Returns:
+            List of match dictionaries
+        """
+        ...
+    
+    def glob(
+        self,
+        path: str,
+        pattern: str,
+        max_results: int = 100
+    ) -> List[Dict[str, Any]]:
+        """Search for files matching glob pattern.
+        
+        Args:
+            path: Directory to search
+            pattern: Glob pattern
+            max_results: Maximum results
+            
+        Returns:
+            List of file info dictionaries
+        """
+        ...
+    
+    def is_available(self) -> bool:
+        """Check if this backend is available.
+        
+        Returns:
+            True if backend can be used
+        """
+        ...
+
+
+class PythonSearchBackend:
+    """Pure Python search backend using existing search_tools.
+    
+    This is always available and serves as the default fallback.
+    """
+    
+    def is_available(self) -> bool:
+        """Python backend is always available."""
+        return True
+    
+    def grep(
+        self,
+        path: str,
+        pattern: str,
+        is_regex: bool = False,
+        case_sensitive: bool = False,
+        include_patterns: Optional[List[str]] = None,
+        exclude_patterns: Optional[List[str]] = None,
+        max_results: int = 100,
+        context_lines: int = 0
+    ) -> List[Dict[str, Any]]:
+        """Search using pure Python regex."""
+        # Import existing implementation
+        from praisonaiagents.context.fast.search_tools import grep_search
+        
+        return grep_search(
+            search_path=path,
+            pattern=pattern,
+            is_regex=is_regex,
+            case_sensitive=case_sensitive,
+            include_patterns=include_patterns,
+            exclude_patterns=exclude_patterns,
+            max_results=max_results,
+            context_lines=context_lines
+        )
+    
+    def glob(
+        self,
+        path: str,
+        pattern: str,
+        max_results: int = 100
+    ) -> List[Dict[str, Any]]:
+        """Search for files using pathlib glob."""
+        from praisonaiagents.context.fast.search_tools import glob_search
+        
+        return glob_search(
+            search_path=path,
+            pattern=pattern,
+            max_results=max_results
+        )
+
+
+class RipgrepBackend:
+    """Optional ripgrep backend using subprocess.
+    
+    Faster for large codebases but requires 'rg' binary in PATH.
+    Falls back to PythonSearchBackend if not available.
+    """
+    
+    def __init__(self):
+        self._available: Optional[bool] = None
+        self._rg_path: Optional[str] = None
+        self._fallback: Optional[PythonSearchBackend] = None
+    
+    def _get_ripgrep_path(self) -> Optional[str]:
+        """Find ripgrep binary in PATH."""
+        if self._rg_path is not None:
+            return self._rg_path if self._rg_path != "" else None
+        
+        import shutil
+        self._rg_path = shutil.which("rg") or ""
+        return self._rg_path if self._rg_path else None
+    
+    def is_available(self) -> bool:
+        """Check if ripgrep is installed."""
+        if self._available is None:
+            self._available = self._get_ripgrep_path() is not None
+        return self._available
+    
+    def _get_fallback(self) -> PythonSearchBackend:
+        """Get Python fallback backend."""
+        if self._fallback is None:
+            self._fallback = PythonSearchBackend()
+        return self._fallback
+    
+    def grep(
+        self,
+        path: str,
+        pattern: str,
+        is_regex: bool = False,
+        case_sensitive: bool = False,
+        include_patterns: Optional[List[str]] = None,
+        exclude_patterns: Optional[List[str]] = None,
+        max_results: int = 100,
+        context_lines: int = 0
+    ) -> List[Dict[str, Any]]:
+        """Search using ripgrep subprocess."""
+        if not self.is_available():
+            logger.debug("Ripgrep not available, using Python fallback")
+            return self._get_fallback().grep(
+                path, pattern, is_regex, case_sensitive,
+                include_patterns, exclude_patterns, max_results, context_lines
+            )
+        
+        import subprocess
+        import json
+        import os
+        
+        rg_path = self._get_ripgrep_path()
+        if not rg_path:
+            return self._get_fallback().grep(
+                path, pattern, is_regex, case_sensitive,
+                include_patterns, exclude_patterns, max_results, context_lines
+            )
+        
+        # Build ripgrep command
+        cmd = [rg_path, "--json"]
+        
+        if not is_regex:
+            cmd.append("--fixed-strings")
+        
+        if not case_sensitive:
+            cmd.append("--ignore-case")
+        
+        if context_lines > 0:
+            cmd.extend(["--context", str(context_lines)])
+        
+        cmd.extend(["--max-count", str(max_results)])
+        
+        # Add include patterns
+        if include_patterns:
+            for p in include_patterns:
+                cmd.extend(["--glob", p])
+        
+        # Add exclude patterns
+        if exclude_patterns:
+            for p in exclude_patterns:
+                cmd.extend(["--glob", f"!{p}"])
+        
+        cmd.append(pattern)
+        cmd.append(path)
+        
+        try:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=30
+            )
+            
+            results = []
+            for line in result.stdout.splitlines():
+                try:
+                    data = json.loads(line)
+                    if data.get("type") == "match":
+                        match_data = data.get("data", {})
+                        results.append({
+                            "path": match_data.get("path", {}).get("text", ""),
+                            "absolute_path": os.path.join(path, match_data.get("path", {}).get("text", "")),
+                            "line_number": match_data.get("line_number", 0),
+                            "content": match_data.get("lines", {}).get("text", "").rstrip(),
+                        })
+                except json.JSONDecodeError:
+                    continue
+            
+            return results[:max_results]
+            
+        except (subprocess.TimeoutExpired, subprocess.SubprocessError) as e:
+            logger.warning(f"Ripgrep failed: {e}, falling back to Python")
+            return self._get_fallback().grep(
+                path, pattern, is_regex, case_sensitive,
+                include_patterns, exclude_patterns, max_results, context_lines
+            )
+    
+    def glob(
+        self,
+        path: str,
+        pattern: str,
+        max_results: int = 100
+    ) -> List[Dict[str, Any]]:
+        """Glob is not well-supported by ripgrep, use Python."""
+        return self._get_fallback().glob(path, pattern, max_results)
+
+
+def _count_files_fast(path: str, limit: int = 1000) -> int:
+    """Quickly count files in a directory (stops at limit).
+    
+    Args:
+        path: Directory to count files in
+        limit: Stop counting after this many files
+        
+    Returns:
+        Number of files found (capped at limit)
+    """
+    import os
+    
+    count = 0
+    try:
+        for root, dirs, files in os.walk(path):
+            # Skip common non-code directories
+            dirs[:] = [d for d in dirs if d not in {
+                '.git', 'node_modules', 'venv', '.venv', '__pycache__',
+                'site-packages', '.tox', 'dist', 'build', '.eggs'
+            }]
+            count += len(files)
+            if count >= limit:
+                return limit
+    except Exception:
+        pass
+    return count
+
+
+# Threshold for switching from Python to Ripgrep
+# Based on profiling: Python is faster for <500 files due to subprocess overhead
+RIPGREP_FILE_THRESHOLD = 500
+
+
+class SmartBackend:
+    """Smart backend that auto-selects Python or Ripgrep based on codebase size.
+    
+    Uses Python for small codebases (<500 files) where subprocess overhead
+    is significant, and Ripgrep for larger codebases where its speed advantage
+    (30-40x faster) outweighs the ~10ms startup overhead.
+    """
+    
+    def __init__(self, workspace_path: Optional[str] = None):
+        self._workspace_path = workspace_path
+        self._python = PythonSearchBackend()
+        self._ripgrep: Optional[RipgrepBackend] = None
+        self._file_count: Optional[int] = None
+        self._use_ripgrep: Optional[bool] = None
+    
+    def _get_ripgrep(self) -> RipgrepBackend:
+        """Lazy-load ripgrep backend."""
+        if self._ripgrep is None:
+            self._ripgrep = RipgrepBackend()
+        return self._ripgrep
+    
+    def _should_use_ripgrep(self, path: str) -> bool:
+        """Determine if ripgrep should be used based on codebase size."""
+        # Check if ripgrep is available
+        rg = self._get_ripgrep()
+        if not rg.is_available():
+            return False
+        
+        # Count files if not already counted for this path
+        if self._file_count is None or self._workspace_path != path:
+            self._workspace_path = path
+            self._file_count = _count_files_fast(path, limit=RIPGREP_FILE_THRESHOLD + 1)
+            self._use_ripgrep = self._file_count >= RIPGREP_FILE_THRESHOLD
+            
+            if self._use_ripgrep:
+                logger.debug(f"Auto-selected ripgrep backend ({self._file_count}+ files)")
+            else:
+                logger.debug(f"Auto-selected Python backend ({self._file_count} files)")
+        
+        return self._use_ripgrep or False
+    
+    def is_available(self) -> bool:
+        """Smart backend is always available."""
+        return True
+    
+    def grep(
+        self,
+        path: str,
+        pattern: str,
+        is_regex: bool = False,
+        case_sensitive: bool = False,
+        include_patterns: Optional[List[str]] = None,
+        exclude_patterns: Optional[List[str]] = None,
+        max_results: int = 100,
+        context_lines: int = 0
+    ) -> List[Dict[str, Any]]:
+        """Search using auto-selected backend."""
+        if self._should_use_ripgrep(path):
+            return self._get_ripgrep().grep(
+                path, pattern, is_regex, case_sensitive,
+                include_patterns, exclude_patterns, max_results, context_lines
+            )
+        return self._python.grep(
+            path, pattern, is_regex, case_sensitive,
+            include_patterns, exclude_patterns, max_results, context_lines
+        )
+    
+    def glob(
+        self,
+        path: str,
+        pattern: str,
+        max_results: int = 100
+    ) -> List[Dict[str, Any]]:
+        """Glob always uses Python (ripgrep doesn't support glob well)."""
+        return self._python.glob(path, pattern, max_results)
+
+
+def get_search_backend(
+    backend_type: str = "auto",
+    workspace_path: Optional[str] = None
+) -> SearchBackend:
+    """Get a search backend by type.
+    
+    Args:
+        backend_type: One of "auto", "python", "ripgrep"
+        workspace_path: Path to workspace (used for auto detection)
+        
+    Returns:
+        SearchBackend instance
+        
+    Raises:
+        ValueError: If backend_type is invalid
+        
+    Performance notes:
+        - "python": Best for small codebases (<500 files), no subprocess overhead
+        - "ripgrep": Best for large codebases (500+ files), 30-40x faster
+        - "auto": Automatically selects based on codebase size
+    """
+    if backend_type == "python":
+        return PythonSearchBackend()
+    
+    elif backend_type == "ripgrep":
+        backend = RipgrepBackend()
+        if not backend.is_available():
+            logger.warning("Ripgrep not available, using Python fallback")
+            return PythonSearchBackend()
+        return backend
+    
+    elif backend_type == "auto":
+        # Use SmartBackend that auto-selects based on codebase size
+        return SmartBackend(workspace_path)
+    
+    else:
+        raise ValueError(f"Invalid backend type: {backend_type}. Must be 'auto', 'python', or 'ripgrep'")

src/praisonai-agents/praisonaiagents/context/models.py

@@ -245,6 +245,20 @@ class ContextConfig:
     Complete context management configuration.
     
     Merges settings from CLI flags, env vars, and config files.
+    
+    Example:
+        # Enable with defaults
+        agent = Agent(instructions="...", context=True)
+        
+        # Custom configuration
+        agent = Agent(
+            instructions="...",
+            context=ContextConfig(
+                auto_compact=True,
+                session_tracking=True,      # Track goal/plan/progress
+                aggregate_memory=True,       # Concurrent multi-memory fetch
+            )
+        )
     """
     # Auto-compaction
     auto_compact: bool = True
@@ -266,6 +280,22 @@ class ContextConfig:
     # Sliding window
     keep_recent_turns: int = 5
     
+    # Session tracking (Agno pattern)
+    session_tracking: bool = False     # Enable goal/plan/progress tracking
+    track_summary: bool = True         # Auto-extract conversation summary
+    track_goal: bool = True            # Track user's objective
+    track_plan: bool = True            # Track steps to achieve goal
+    track_progress: bool = True        # Track completed steps
+    
+    # Multi-memory aggregation (CrewAI pattern)
+    aggregate_memory: bool = False     # Enable concurrent multi-source fetch
+    aggregate_sources: List[str] = field(default_factory=lambda: [
+        "memory",      # Short-term memory
+        "knowledge",   # Long-term knowledge base
+        "rag",         # RAG retrieval
+    ])
+    aggregate_max_tokens: int = 4000   # Max tokens for aggregated context
+    
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
         result = asdict(self)

src/praisonai-agents/praisonaiagents/context/session_tracker.py

@@ -0,0 +1,287 @@
+"""
+Session Context Tracker for PraisonAI Agents.
+
+Tracks session state (summary, goal, plan, progress) across conversation turns.
+Inspired by Agno's SessionContextStore pattern.
+"""
+
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, List, Optional
+from datetime import datetime
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SessionState:
+    """
+    Session state tracking data.
+    
+    Tracks:
+    - summary: What's happened in the conversation
+    - goal: User's objective for this session
+    - plan: Steps to achieve the goal
+    - progress: Completed steps
+    """
+    session_id: str = ""
+    summary: str = ""
+    goal: str = ""
+    plan: List[str] = field(default_factory=list)
+    progress: List[str] = field(default_factory=list)
+    
+    # Metadata
+    created_at: str = ""
+    updated_at: str = ""
+    turn_count: int = 0
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionState":
+        """Create from dictionary."""
+        return cls(
+            session_id=data.get("session_id", ""),
+            summary=data.get("summary", ""),
+            goal=data.get("goal", ""),
+            plan=data.get("plan", []),
+            progress=data.get("progress", []),
+            created_at=data.get("created_at", ""),
+            updated_at=data.get("updated_at", ""),
+            turn_count=data.get("turn_count", 0),
+        )
+
+
+class SessionContextTracker:
+    """
+    Tracks session context (summary/goal/plan/progress) across turns.
+    
+    Implements the Agno SessionContextStore pattern for PraisonAI.
+    Builds upon previous context rather than starting fresh each turn.
+    
+    Example:
+        tracker = SessionContextTracker(session_id="abc123")
+        
+        # After each turn, update state
+        tracker.update_from_turn(messages, llm_client)
+        
+        # Get context string for system prompt
+        context_str = tracker.to_context_string()
+    """
+    
+    def __init__(
+        self,
+        session_id: str = "",
+        track_summary: bool = True,
+        track_goal: bool = True,
+        track_plan: bool = True,
+        track_progress: bool = True,
+    ):
+        """
+        Initialize session context tracker.
+        
+        Args:
+            session_id: Session identifier
+            track_summary: Whether to track conversation summary
+            track_goal: Whether to track user's objective
+            track_plan: Whether to track steps to goal
+            track_progress: Whether to track completed steps
+        """
+        self.session_id = session_id or self._generate_session_id()
+        self.track_summary = track_summary
+        self.track_goal = track_goal
+        self.track_plan = track_plan
+        self.track_progress = track_progress
+        
+        self._state = SessionState(
+            session_id=self.session_id,
+            created_at=datetime.utcnow().isoformat() + "Z",
+            updated_at=datetime.utcnow().isoformat() + "Z",
+        )
+    
+    def _generate_session_id(self) -> str:
+        """Generate a unique session ID."""
+        import uuid
+        return str(uuid.uuid4())[:8]
+    
+    @property
+    def state(self) -> SessionState:
+        """Get current session state."""
+        return self._state
+    
+    @property
+    def summary(self) -> str:
+        """Get current summary."""
+        return self._state.summary
+    
+    @property
+    def goal(self) -> str:
+        """Get current goal."""
+        return self._state.goal
+    
+    @property
+    def plan(self) -> List[str]:
+        """Get current plan."""
+        return self._state.plan
+    
+    @property
+    def progress(self) -> List[str]:
+        """Get completed steps."""
+        return self._state.progress
+    
+    def update_summary(self, summary: str) -> None:
+        """Update conversation summary."""
+        if self.track_summary:
+            self._state.summary = summary
+            self._mark_updated()
+    
+    def update_goal(self, goal: str) -> None:
+        """Update user's objective."""
+        if self.track_goal:
+            self._state.goal = goal
+            self._mark_updated()
+    
+    def update_plan(self, plan: List[str]) -> None:
+        """Update plan steps."""
+        if self.track_plan:
+            self._state.plan = plan
+            self._mark_updated()
+    
+    def add_progress(self, step: str) -> None:
+        """Add a completed step to progress."""
+        if self.track_progress:
+            self._state.progress.append(step)
+            self._mark_updated()
+    
+    def mark_plan_step_complete(self, step_index: int) -> None:
+        """Mark a plan step as complete and add to progress."""
+        if 0 <= step_index < len(self._state.plan):
+            step = self._state.plan[step_index]
+            self.add_progress(step)
+    
+    def _mark_updated(self) -> None:
+        """Update the updated_at timestamp."""
+        self._state.updated_at = datetime.utcnow().isoformat() + "Z"
+        self._state.turn_count += 1
+    
+    def to_context_string(self) -> str:
+        """
+        Format session context for injection into system prompt.
+        
+        Returns string formatted for the agent to use as context.
+        """
+        parts = []
+        
+        if self._state.summary:
+            parts.append(f"**Conversation Summary**: {self._state.summary}")
+        
+        if self._state.goal:
+            parts.append(f"**User's Goal**: {self._state.goal}")
+        
+        if self._state.plan:
+            plan_str = "\n".join(f"  {i+1}. {step}" for i, step in enumerate(self._state.plan))
+            parts.append(f"**Plan**:\n{plan_str}")
+        
+        if self._state.progress:
+            progress_str = "\n".join(f"  ✓ {step}" for step in self._state.progress)
+            parts.append(f"**Progress**:\n{progress_str}")
+        
+        if not parts:
+            return ""
+        
+        return "\n\n".join(parts)
+    
+    def to_system_prompt_section(self) -> str:
+        """
+        Generate a system prompt section with session context.
+        
+        Follows Agno's pattern of including guidelines for using the context.
+        """
+        context = self.to_context_string()
+        
+        if not context:
+            return ""
+        
+        return f"""<session_context>
+This is a continuation of an ongoing session. Here's where things stand:
+
+{context}
+
+<session_context_guidelines>
+Use this context to maintain continuity:
+- Reference earlier decisions and conclusions naturally
+- Don't re-ask questions that have already been answered
+- Build on established understanding rather than starting fresh
+- If the user references something from "earlier," this context has the details
+
+Current messages take precedence if there's any conflict with this summary.
+</session_context_guidelines>
+</session_context>"""
+    
+    def update_from_messages(
+        self,
+        messages: List[Dict[str, Any]],
+        extract_goal: bool = True,
+    ) -> None:
+        """
+        Update session state from conversation messages.
+        
+        This is a simple heuristic-based extraction. For more accurate
+        extraction, use update_from_turn with an LLM.
+        
+        Args:
+            messages: List of conversation messages
+            extract_goal: Whether to try extracting goal from first user message
+        """
+        if not messages:
+            return
+        
+        # Extract goal from first user message if not set
+        if extract_goal and not self._state.goal:
+            for msg in messages:
+                if msg.get("role") == "user":
+                    content = msg.get("content", "")
+                    if isinstance(content, str) and len(content) < 500:
+                        self._state.goal = content[:200]
+                    break
+        
+        # Update turn count
+        user_turns = sum(1 for m in messages if m.get("role") == "user")
+        self._state.turn_count = user_turns
+        self._mark_updated()
+    
+    def clear(self) -> None:
+        """Clear session state."""
+        self._state = SessionState(
+            session_id=self.session_id,
+            created_at=self._state.created_at,
+            updated_at=datetime.utcnow().isoformat() + "Z",
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "session_id": self.session_id,
+            "track_summary": self.track_summary,
+            "track_goal": self.track_goal,
+            "track_plan": self.track_plan,
+            "track_progress": self.track_progress,
+            "state": self._state.to_dict(),
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionContextTracker":
+        """Create from dictionary."""
+        tracker = cls(
+            session_id=data.get("session_id", ""),
+            track_summary=data.get("track_summary", True),
+            track_goal=data.get("track_goal", True),
+            track_plan=data.get("track_plan", True),
+            track_progress=data.get("track_progress", True),
+        )
+        if "state" in data:
+            tracker._state = SessionState.from_dict(data["state"])
+        return tracker

src/praisonai-agents/praisonaiagents/embedding/__init__.py

@@ -0,0 +1,65 @@
+"""
+Embedding module for PraisonAI Agents.
+
+This module provides a unified embedding API that can be used throughout
+the praisonaiagents package. It consolidates embedding functionality that
+was previously duplicated in memory.py and knowledge.py.
+
+Usage:
+    >>> from praisonaiagents import embedding
+    >>> result = embedding("Hello world")
+    >>> print(result.embeddings[0][:5])
+    
+    >>> from praisonaiagents.embedding import embedding, EmbeddingResult
+    >>> result = embedding(["Hello", "World"])
+    >>> print(len(result))
+    2
+
+The module uses lazy imports to avoid loading litellm until actually needed,
+ensuring zero performance impact on import time.
+"""
+
+__all__ = [
+    'embedding',
+    'aembedding',
+    'embed',
+    'aembed',
+    'EmbeddingResult',
+    'get_dimensions',
+    'MODEL_DIMENSIONS',
+    'DEFAULT_DIMENSION',
+]
+
+
+def __getattr__(name: str):
+    """Lazy load module components to avoid importing litellm on module load."""
+    if name == 'EmbeddingResult':
+        from .result import EmbeddingResult
+        return EmbeddingResult
+    
+    if name in ('embedding', 'embed'):
+        from .embed import embedding
+        return embedding
+    
+    if name in ('aembedding', 'aembed'):
+        from .embed import aembedding
+        return aembedding
+    
+    if name == 'get_dimensions':
+        from .dimensions import get_dimensions
+        return get_dimensions
+    
+    if name == 'MODEL_DIMENSIONS':
+        from .dimensions import MODEL_DIMENSIONS
+        return MODEL_DIMENSIONS
+    
+    if name == 'DEFAULT_DIMENSION':
+        from .dimensions import DEFAULT_DIMENSION
+        return DEFAULT_DIMENSION
+    
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Return list of public names for dir()."""
+    return list(__all__)

src/praisonai-agents/praisonaiagents/embedding/dimensions.py

@@ -0,0 +1,77 @@
+"""
+Embedding dimension utilities.
+
+This module provides utilities for determining embedding dimensions
+based on model names. Consolidates the dimension lookup logic that
+was previously duplicated in memory.py and knowledge.py.
+"""
+
+from typing import Dict
+
+# Common embedding model dimensions
+# This is the single source of truth for dimension lookup
+MODEL_DIMENSIONS: Dict[str, int] = {
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+    "text-embedding-ada-002": 1536,
+    "text-embedding-002": 1536,
+    # Cohere models
+    "embed-english-v3.0": 1024,
+    "embed-multilingual-v3.0": 1024,
+    "embed-english-light-v3.0": 384,
+    "embed-multilingual-light-v3.0": 384,
+    # Voyage models
+    "voyage-3": 1024,
+    "voyage-3-lite": 512,
+    "voyage-code-3": 1024,
+    # Mistral models
+    "mistral-embed": 1024,
+    # Jina models
+    "jina-embeddings-v3": 1024,
+    "jina-embeddings-v2-base-en": 768,
+    # HuggingFace common models
+    "all-MiniLM-L6-v2": 384,
+    "all-mpnet-base-v2": 768,
+    "bge-large-en-v1.5": 1024,
+    "bge-base-en-v1.5": 768,
+    "bge-small-en-v1.5": 384,
+}
+
+# Default dimension for unknown models
+DEFAULT_DIMENSION = 1536
+
+
+def get_dimensions(model_name: str) -> int:
+    """Get embedding dimensions based on model name.
+    
+    This function checks if the model name contains any known model
+    identifiers and returns the corresponding dimension. Falls back
+    to the default dimension (1536) for unknown models.
+    
+    Args:
+        model_name: The embedding model name (e.g., "text-embedding-3-small")
+    
+    Returns:
+        The embedding dimension for the model
+    
+    Example:
+        >>> get_dimensions("text-embedding-3-small")
+        1536
+        >>> get_dimensions("openai/text-embedding-3-large")
+        3072
+        >>> get_dimensions("unknown-model")
+        1536
+    """
+    model_lower = model_name.lower()
+    
+    # Check for exact match first
+    if model_lower in MODEL_DIMENSIONS:
+        return MODEL_DIMENSIONS[model_lower]
+    
+    # Check if model name contains known model identifiers
+    for model_key, dimensions in MODEL_DIMENSIONS.items():
+        if model_key in model_lower:
+            return dimensions
+    
+    # Default to 1536 for unknown models (OpenAI standard)
+    return DEFAULT_DIMENSION

src/praisonai-agents/praisonaiagents/embedding/embed.py

@@ -0,0 +1,174 @@
+"""
+Core embedding functions.
+
+This module provides the main embedding() and aembedding() functions
+that wrap LiteLLM's embedding API with a consistent interface.
+"""
+
+from typing import Optional, Union, List, Any, Dict
+
+from .result import EmbeddingResult
+
+
+def embedding(
+    input: Union[str, List[str]],
+    model: str = "text-embedding-3-small",
+    dimensions: Optional[int] = None,
+    encoding_format: str = "float",
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> EmbeddingResult:
+    """
+    Generate embeddings for text using LiteLLM.
+    
+    This is the primary embedding function that supports all LiteLLM
+    embedding providers (OpenAI, Azure, Cohere, Voyage, etc.).
+    
+    Args:
+        input: Text or list of texts to embed
+        model: Model name (e.g., "text-embedding-3-small", "text-embedding-3-large")
+        dimensions: Optional output dimensions (for models that support it)
+        encoding_format: "float" or "base64"
+        timeout: Request timeout in seconds
+        api_key: Optional API key override
+        api_base: Optional API base URL override
+        metadata: Optional metadata for tracing
+        **kwargs: Additional arguments passed to litellm.embedding()
+        
+    Returns:
+        EmbeddingResult with embeddings list, model, usage, and metadata
+        
+    Example:
+        >>> from praisonaiagents import embedding
+        >>> result = embedding("Hello, world!")
+        >>> print(len(result.embeddings[0]))
+        1536
+        
+        >>> result = embedding(["Hello", "World"], model="text-embedding-3-large")
+        >>> print(len(result.embeddings))
+        2
+    """
+    import litellm
+    
+    call_kwargs = {
+        'model': model,
+        'input': input,
+        'encoding_format': encoding_format,
+        'timeout': timeout,
+    }
+    
+    if dimensions:
+        call_kwargs['dimensions'] = dimensions
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = litellm.embedding(**call_kwargs)
+    
+    embeddings = []
+    if hasattr(response, 'data'):
+        for item in response.data:
+            if isinstance(item, dict):
+                embeddings.append(item.get('embedding', []))
+            else:
+                embeddings.append(getattr(item, 'embedding', []))
+    
+    usage = None
+    if hasattr(response, 'usage'):
+        usage = {
+            'prompt_tokens': getattr(response.usage, 'prompt_tokens', 0),
+            'total_tokens': getattr(response.usage, 'total_tokens', 0),
+        }
+    
+    return EmbeddingResult(
+        embeddings=embeddings,
+        model=model,
+        usage=usage,
+        metadata=metadata or {},
+    )
+
+
+async def aembedding(
+    input: Union[str, List[str]],
+    model: str = "text-embedding-3-small",
+    dimensions: Optional[int] = None,
+    encoding_format: str = "float",
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> EmbeddingResult:
+    """
+    Async: Generate embeddings for text using LiteLLM.
+    
+    This is the async version of embedding() for use in async contexts.
+    See embedding() for full documentation.
+    
+    Example:
+        >>> from praisonaiagents import aembedding
+        >>> result = await aembedding("Hello, world!")
+        >>> print(len(result.embeddings[0]))
+        1536
+    """
+    import litellm
+    
+    call_kwargs = {
+        'model': model,
+        'input': input,
+        'encoding_format': encoding_format,
+        'timeout': timeout,
+    }
+    
+    if dimensions:
+        call_kwargs['dimensions'] = dimensions
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = await litellm.aembedding(**call_kwargs)
+    
+    embeddings = []
+    if hasattr(response, 'data'):
+        for item in response.data:
+            if isinstance(item, dict):
+                embeddings.append(item.get('embedding', []))
+            else:
+                embeddings.append(getattr(item, 'embedding', []))
+    
+    usage = None
+    if hasattr(response, 'usage'):
+        usage = {
+            'prompt_tokens': getattr(response.usage, 'prompt_tokens', 0),
+            'total_tokens': getattr(response.usage, 'total_tokens', 0),
+        }
+    
+    return EmbeddingResult(
+        embeddings=embeddings,
+        model=model,
+        usage=usage,
+        metadata=metadata or {},
+    )
+
+
+# Aliases for backwards compatibility and OpenAI naming convention
+embed = embedding
+"""Alias for embedding(). Use embedding() or embed() interchangeably."""
+
+aembed = aembedding
+"""Alias for aembedding(). Use aembedding() or aembed() interchangeably."""

src/praisonai-agents/praisonaiagents/embedding/result.py

@@ -0,0 +1,42 @@
+"""
+EmbeddingResult dataclass for embedding responses.
+
+This module provides a structured result type for embedding operations,
+matching the OpenAI embedding response format.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, List, Dict, Any
+
+
+@dataclass
+class EmbeddingResult:
+    """Result from embedding generation.
+    
+    Attributes:
+        embeddings: List of embedding vectors (each is a list of floats)
+        model: The model used for embedding (optional)
+        usage: Token usage information (optional)
+        metadata: Additional metadata (optional)
+    
+    Example:
+        >>> result = EmbeddingResult(
+        ...     embeddings=[[0.1, 0.2, 0.3]],
+        ...     model="text-embedding-3-small",
+        ...     usage={"prompt_tokens": 5, "total_tokens": 5}
+        ... )
+        >>> print(len(result.embeddings[0]))
+        3
+    """
+    embeddings: List[List[float]]
+    model: Optional[str] = None
+    usage: Optional[Dict[str, int]] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def __len__(self) -> int:
+        """Return the number of embeddings."""
+        return len(self.embeddings)
+    
+    def __getitem__(self, index: int) -> List[float]:
+        """Get embedding by index."""
+        return self.embeddings[index]

src/praisonai-agents/praisonaiagents/eval/__init__.py

@@ -34,6 +34,14 @@ __all__ = [
     "PerformanceMetrics",
     "ToolCallResult",
     "CriteriaScore",
+    # DRY: Common grading base classes and protocols
+    "BaseLLMGrader",
+    "GradeResult",
+    "GraderProtocol",
+    "GradeResultProtocol",
+    "ScoredResultProtocol",
+    "AsyncGraderProtocol",
+    "parse_score_reasoning",
 ]
 
 _LAZY_IMPORTS = {
@@ -50,6 +58,14 @@ _LAZY_IMPORTS = {
     "PerformanceMetrics": ("results", "PerformanceMetrics"),
     "ToolCallResult": ("results", "ToolCallResult"),
     "CriteriaScore": ("results", "CriteriaScore"),
+    # DRY: Common grading base classes and protocols
+    "BaseLLMGrader": ("grader", "BaseLLMGrader"),
+    "GradeResult": ("grader", "GradeResult"),
+    "GraderProtocol": ("protocols", "GraderProtocol"),
+    "GradeResultProtocol": ("protocols", "GradeResultProtocol"),
+    "ScoredResultProtocol": ("protocols", "ScoredResultProtocol"),
+    "AsyncGraderProtocol": ("protocols", "AsyncGraderProtocol"),
+    "parse_score_reasoning": ("grader", "parse_score_reasoning"),
 }
 
 

src/praisonai-agents/praisonaiagents/eval/accuracy.py

@@ -6,10 +6,11 @@ Evaluates agent output accuracy by comparing against expected output using LLM-a
 
 import os
 import logging
-from typing import Any, Callable, List, Optional, Union, TYPE_CHECKING
+from typing import Callable, Optional, Union, TYPE_CHECKING
 
 from .base import BaseEvaluator
 from .results import AccuracyResult, EvaluationScore
+from .grader import parse_score_reasoning
 
 if TYPE_CHECKING:
     from ..agent.agent import Agent
@@ -123,19 +124,8 @@ REASONING: [brief explanation]"""
         
         response_text = response.choices[0].message.content or ""
         
-        score = 5.0
-        reasoning = "Unable to parse response"
-        
-        lines = response_text.strip().split('\n')
-        for line in lines:
-            if line.startswith('SCORE:'):
-                try:
-                    score = float(line.replace('SCORE:', '').strip())
-                    score = max(1.0, min(10.0, score))
-                except ValueError:
-                    pass
-            elif line.startswith('REASONING:'):
-                reasoning = line.replace('REASONING:', '').strip()
+        # DRY: Use common parsing function from grader module
+        score, reasoning = parse_score_reasoning(response_text)
         
         return EvaluationScore(
             score=score,

src/praisonai-agents/praisonaiagents/eval/criteria.py

@@ -6,10 +6,11 @@ Evaluates agent outputs against custom criteria using LLM-as-judge.
 
 import os
 import logging
-from typing import Callable, List, Literal, Optional, Union, TYPE_CHECKING
+from typing import Callable, Literal, Optional, Union, TYPE_CHECKING
 
 from .base import BaseEvaluator
 from .results import CriteriaResult, CriteriaScore
+from .grader import parse_score_reasoning
 
 if TYPE_CHECKING:
     from ..agent.agent import Agent
@@ -146,28 +147,24 @@ REASONING: [brief explanation]"""
         
         response_text = response.choices[0].message.content or ""
         
-        score = 5.0
-        passed = False
-        reasoning = "Unable to parse response"
-        
-        lines = response_text.strip().split('\n')
-        for line in lines:
-            if self.scoring_type == "binary":
+        # DRY: Use common parsing for numeric mode, custom for binary
+        if self.scoring_type == "binary":
+            score = 5.0
+            passed = False
+            reasoning = "Unable to parse response"
+            
+            lines = response_text.strip().split('\n')
+            for line in lines:
                 if line.startswith('RESULT:'):
                     result_str = line.replace('RESULT:', '').strip().upper()
                     passed = result_str == 'PASS'
                     score = 10.0 if passed else 0.0
-            else:
-                if line.startswith('SCORE:'):
-                    try:
-                        score = float(line.replace('SCORE:', '').strip())
-                        score = max(1.0, min(10.0, score))
-                        passed = score >= self.threshold
-                    except ValueError:
-                        pass
-            
-            if line.startswith('REASONING:'):
-                reasoning = line.replace('REASONING:', '').strip()
+                elif line.startswith('REASONING:'):
+                    reasoning = line.replace('REASONING:', '').strip()
+        else:
+            # DRY: Use common parsing function from grader module
+            score, reasoning = parse_score_reasoning(response_text)
+            passed = score >= self.threshold
         
         return CriteriaScore(
             score=score,

src/praisonai-agents/praisonaiagents/eval/grader.py

@@ -0,0 +1,378 @@
+"""
+Base LLM Grader for PraisonAI Agents.
+
+Provides a common base class for LLM-as-judge grading to eliminate duplication
+across AccuracyEvaluator, CriteriaEvaluator, and TrainingGrader.
+
+DRY: This module extracts the common prompt building and response parsing logic
+that was duplicated across multiple evaluators.
+"""
+
+import os
+import logging
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GradeResult:
+    """
+    Result from grading an agent output.
+    
+    This is the common result type used by all graders.
+    
+    Attributes:
+        score: Quality score (1-10)
+        reasoning: Explanation for the score
+        suggestions: List of improvement suggestions (optional)
+        input_text: The input that was given
+        output: The output that was graded
+        expected_output: Optional expected output
+        timestamp: When grading occurred
+    """
+    score: float
+    reasoning: str
+    suggestions: List[str] = field(default_factory=list)
+    input_text: str = ""
+    output: str = ""
+    expected_output: Optional[str] = None
+    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat())
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GradeResult":
+        """Create from dictionary."""
+        return cls(
+            score=data.get("score", 5.0),
+            reasoning=data.get("reasoning", ""),
+            suggestions=data.get("suggestions", []),
+            input_text=data.get("input_text", ""),
+            output=data.get("output", ""),
+            expected_output=data.get("expected_output"),
+            timestamp=data.get("timestamp", datetime.utcnow().isoformat()),
+        )
+
+
+class BaseLLMGrader:
+    """
+    Base class for LLM-as-judge grading.
+    
+    Provides common functionality for:
+    - Building grading prompts
+    - Parsing LLM responses (SCORE/REASONING/SUGGESTIONS format)
+    - Lazy litellm import
+    - Sync and async grading
+    
+    Subclasses can override:
+    - _build_prompt() for custom prompt templates
+    - _parse_response() for custom response parsing
+    
+    Usage:
+        grader = BaseLLMGrader()
+        result = grader.grade(
+            input_text="What is Python?",
+            output="Python is a programming language",
+            expected_output="Python is a high-level programming language"
+        )
+        print(f"Score: {result.score}/10")
+    """
+    
+    # Default prompt template - can be overridden by subclasses
+    PROMPT_TEMPLATE = """You are an expert evaluator for AI agent outputs. Your task is to grade the quality of an agent's response.
+
+INPUT (what the agent was asked):
+{input_text}
+
+AGENT OUTPUT (what the agent responded):
+{output}
+{expected_section}
+GRADING CRITERIA:
+- Accuracy: Is the response factually correct?
+- Completeness: Does it fully address the input?
+- Clarity: Is it well-written and easy to understand?
+- Relevance: Does it stay on topic?
+
+SCORING GUIDELINES:
+- 10: Perfect - Excellent in all criteria
+- 8-9: Very Good - Minor improvements possible
+- 6-7: Good - Some issues but mostly correct
+- 4-5: Fair - Significant issues
+- 2-3: Poor - Major problems
+- 1: Very Poor - Completely wrong or irrelevant
+
+Respond in this EXACT format:
+SCORE: [number 1-10]
+REASONING: [brief explanation of the score]
+SUGGESTIONS:
+- [first suggestion for improvement]
+- [second suggestion if applicable]
+- [third suggestion if applicable]
+
+If no suggestions are needed, write "SUGGESTIONS: None"
+"""
+    
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        temperature: float = 0.1,
+        max_tokens: int = 500,
+    ):
+        """
+        Initialize the grader.
+        
+        Args:
+            model: LLM model to use for grading (default: gpt-4o-mini)
+            temperature: Temperature for LLM calls (default: 0.1 for consistency)
+            max_tokens: Maximum tokens for LLM response
+        """
+        self.model = model or os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini")
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+    
+    def _build_prompt(
+        self,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str] = None,
+    ) -> str:
+        """
+        Build the grading prompt.
+        
+        Override this method in subclasses for custom prompt templates.
+        
+        Args:
+            input_text: The input given to the agent
+            output: The agent's output to grade
+            expected_output: Optional expected output for comparison
+            
+        Returns:
+            The prompt string for the LLM
+        """
+        expected_section = ""
+        if expected_output:
+            expected_section = f"""
+EXPECTED OUTPUT (ideal response):
+{expected_output}
+"""
+        
+        return self.PROMPT_TEMPLATE.format(
+            input_text=input_text,
+            output=output,
+            expected_section=expected_section,
+        )
+    
+    def _parse_response(
+        self,
+        response_text: str,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str],
+    ) -> GradeResult:
+        """
+        Parse the LLM response into a GradeResult.
+        
+        Override this method in subclasses for custom parsing logic.
+        
+        Args:
+            response_text: Raw LLM response
+            input_text: Original input
+            output: Original output
+            expected_output: Original expected output
+            
+        Returns:
+            Parsed GradeResult
+        """
+        score = 5.0  # Default
+        reasoning = "Unable to parse response"
+        suggestions: List[str] = []
+        
+        lines = response_text.strip().split('\n')
+        in_suggestions = False
+        
+        for line in lines:
+            line = line.strip()
+            
+            if line.startswith('SCORE:'):
+                try:
+                    score_str = line.replace('SCORE:', '').strip()
+                    score = float(score_str)
+                    # Clamp to valid range
+                    score = max(1.0, min(10.0, score))
+                except ValueError:
+                    pass
+            
+            elif line.startswith('REASONING:'):
+                reasoning = line.replace('REASONING:', '').strip()
+            
+            elif line.startswith('SUGGESTIONS:'):
+                in_suggestions = True
+                rest = line.replace('SUGGESTIONS:', '').strip()
+                if rest.lower() != 'none' and rest:
+                    suggestions.append(rest)
+            
+            elif in_suggestions and line.startswith('-'):
+                suggestion = line.lstrip('- ').strip()
+                if suggestion and suggestion.lower() != 'none':
+                    suggestions.append(suggestion)
+        
+        return GradeResult(
+            score=score,
+            reasoning=reasoning,
+            suggestions=suggestions,
+            input_text=input_text,
+            output=output,
+            expected_output=expected_output,
+        )
+    
+    def _get_litellm(self):
+        """
+        Lazy import litellm.
+        
+        Returns:
+            The litellm module
+            
+        Raises:
+            ImportError: If litellm is not installed
+        """
+        try:
+            import litellm
+            return litellm
+        except ImportError:
+            raise ImportError(
+                "litellm is required for LLM grading. "
+                "Install with: pip install litellm"
+            )
+    
+    def grade(
+        self,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str] = None,
+    ) -> GradeResult:
+        """
+        Grade an agent output.
+        
+        Args:
+            input_text: The input given to the agent
+            output: The agent's output to grade
+            expected_output: Optional expected output for comparison
+            
+        Returns:
+            GradeResult with score, reasoning, and suggestions
+        """
+        litellm = self._get_litellm()
+        
+        prompt = self._build_prompt(input_text, output, expected_output)
+        
+        try:
+            response = litellm.completion(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+            )
+            
+            response_text = response.choices[0].message.content or ""
+            return self._parse_response(response_text, input_text, output, expected_output)
+            
+        except Exception as e:
+            logger.warning(f"Grading failed: {e}")
+            return GradeResult(
+                score=5.0,
+                reasoning=f"Grading error: {str(e)}",
+                suggestions=[],
+                input_text=input_text,
+                output=output,
+                expected_output=expected_output,
+            )
+    
+    async def grade_async(
+        self,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str] = None,
+    ) -> GradeResult:
+        """
+        Grade an agent output asynchronously.
+        
+        Args:
+            input_text: The input given to the agent
+            output: The agent's output to grade
+            expected_output: Optional expected output for comparison
+            
+        Returns:
+            GradeResult with score, reasoning, and suggestions
+        """
+        litellm = self._get_litellm()
+        
+        prompt = self._build_prompt(input_text, output, expected_output)
+        
+        try:
+            response = await litellm.acompletion(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+            )
+            
+            response_text = response.choices[0].message.content or ""
+            return self._parse_response(response_text, input_text, output, expected_output)
+            
+        except Exception as e:
+            logger.warning(f"Async grading failed: {e}")
+            return GradeResult(
+                score=5.0,
+                reasoning=f"Grading error: {str(e)}",
+                suggestions=[],
+                input_text=input_text,
+                output=output,
+                expected_output=expected_output,
+            )
+
+
+def parse_score_reasoning(response_text: str) -> tuple:
+    """
+    Parse SCORE and REASONING from LLM response.
+    
+    DRY: Static utility function for parsing LLM-as-judge responses.
+    Used by AccuracyEvaluator, CriteriaEvaluator, and BaseLLMGrader.
+    
+    Args:
+        response_text: Raw LLM response text
+        
+    Returns:
+        Tuple of (score: float, reasoning: str)
+    """
+    score = 5.0  # Default
+    reasoning = "Unable to parse response"
+    
+    lines = response_text.strip().split('\n')
+    for line in lines:
+        line = line.strip()
+        
+        if line.startswith('SCORE:'):
+            try:
+                score_str = line.replace('SCORE:', '').strip()
+                score = float(score_str)
+                # Clamp to valid range
+                score = max(1.0, min(10.0, score))
+            except ValueError:
+                pass
+        
+        elif line.startswith('REASONING:'):
+            reasoning = line.replace('REASONING:', '').strip()
+    
+    return score, reasoning
+
+
+__all__ = [
+    'GradeResult',
+    'BaseLLMGrader',
+    'parse_score_reasoning',
+]

src/praisonai-agents/praisonaiagents/eval/protocols.py

@@ -0,0 +1,128 @@
+"""
+Protocols for PraisonAI Agents Evaluation Framework.
+
+Provides Protocol interfaces for grading and evaluation to enable:
+- DRY: Common grading logic shared across eval and train modules
+- Mocking: Easy testing without real LLM calls
+- Extensibility: Custom grader implementations
+
+These protocols are lightweight and have zero performance impact.
+"""
+from typing import Protocol, runtime_checkable, Optional, List, Dict, Any
+
+
+@runtime_checkable
+class GradeResultProtocol(Protocol):
+    """
+    Protocol for grading results.
+    
+    Defines the minimal interface for any grading result.
+    """
+    score: float
+    reasoning: str
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary."""
+        ...
+
+
+@runtime_checkable
+class GraderProtocol(Protocol):
+    """
+    Protocol for LLM-as-judge grading.
+    
+    Defines the interface for grading agent outputs using an LLM.
+    Implementations include AccuracyEvaluator, CriteriaEvaluator, and TrainingGrader.
+    
+    Example:
+        ```python
+        class MockGrader:
+            model = "mock"
+            temperature = 0.1
+            
+            def grade(self, input_text, output, expected_output=None):
+                return MockGradeResult(score=8.0, reasoning="Mock")
+        
+        # Use in tests
+        grader: GraderProtocol = MockGrader()
+        result = grader.grade("test", "output")
+        ```
+    """
+    model: str
+    temperature: float
+    
+    def grade(
+        self,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str] = None,
+    ) -> GradeResultProtocol:
+        """
+        Grade an agent output.
+        
+        Args:
+            input_text: The input given to the agent
+            output: The agent's output to grade
+            expected_output: Optional expected output for comparison
+            
+        Returns:
+            A GradeResultProtocol with score and reasoning
+        """
+        ...
+
+
+@runtime_checkable
+class ScoredResultProtocol(Protocol):
+    """
+    Protocol for scored evaluation results.
+    
+    Defines the interface for results that have scores and statistics.
+    """
+    @property
+    def scores(self) -> List[float]:
+        """Get list of all scores."""
+        ...
+    
+    @property
+    def avg_score(self) -> float:
+        """Calculate average score."""
+        ...
+    
+    @property
+    def passed(self) -> bool:
+        """Check if evaluation passed."""
+        ...
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary."""
+        ...
+    
+    def print_summary(self) -> None:
+        """Print a summary of the results."""
+        ...
+
+
+@runtime_checkable
+class AsyncGraderProtocol(Protocol):
+    """
+    Async Protocol for LLM-as-judge grading.
+    """
+    model: str
+    temperature: float
+    
+    async def grade_async(
+        self,
+        input_text: str,
+        output: str,
+        expected_output: Optional[str] = None,
+    ) -> GradeResultProtocol:
+        """Grade an agent output asynchronously."""
+        ...
+
+
+__all__ = [
+    'GradeResultProtocol',
+    'GraderProtocol',
+    'ScoredResultProtocol',
+    'AsyncGraderProtocol',
+]

src/praisonai-agents/praisonaiagents/knowledge/knowledge.py

@@ -110,22 +110,8 @@ class MongoDBMemory:
     
     def _get_embedding_dimensions(self, model_name: str) -> int:
         """Get embedding dimensions based on model name."""
-        # Common embedding model dimensions
-        model_dimensions = {
-            "text-embedding-3-small": 1536,
-            "text-embedding-3-large": 3072,
-            "text-embedding-ada-002": 1536,
-            "text-embedding-002": 1536,
-            # Add more models as needed
-        }
-        
-        # Check if model name contains known model identifiers
-        for model_key, dimensions in model_dimensions.items():
-            if model_key in model_name.lower():
-                return dimensions
-        
-        # Default to 1536 for unknown models (OpenAI standard)
-        return 1536
+        from praisonaiagents.embedding import get_dimensions
+        return get_dimensions(model_name)
     
     def _create_indexes(self):
         """Create necessary indexes for MongoDB."""
@@ -173,16 +159,11 @@ class MongoDBMemory:
             logging.warning(f"Could not create vector search index: {e}")
     
     def _get_embedding(self, text):
-        """Get embedding for text using litellm for unified provider support."""
+        """Get embedding for text using the unified embedding module."""
         try:
-            import litellm
-            
-            # Use litellm.embedding() for unified multi-provider support
-            response = litellm.embedding(
-                model=self.embedding_model_name,
-                input=text
-            )
-            return response.data[0]["embedding"]
+            from praisonaiagents.embedding import embedding
+            result = embedding(text, model=self.embedding_model_name)
+            return result.embeddings[0] if result.embeddings else None
         except Exception as e:
             logging.error(f"Error getting embedding: {e}")
             return None

src/praisonai-agents/praisonaiagents/llm/llm.py

@@ -1210,8 +1210,13 @@ Now provide your final answer using this result."""
             # Only append JSON schema for non-Gemini models or Gemini models without structured output support
             if (output_json or output_pydantic) and not is_gemini_with_structured_output:
                 schema_model = output_json or output_pydantic
-                if schema_model and hasattr(schema_model, 'model_json_schema'):
-                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(schema_model.model_json_schema())}"
+                if schema_model:
+                    # Handle Pydantic model
+                    if hasattr(schema_model, 'model_json_schema'):
+                        system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(schema_model.model_json_schema())}"
+                    # Handle inline dict schema (Option A from YAML)
+                    elif isinstance(schema_model, dict):
+                        system_prompt += f"\nReturn ONLY a JSON object that matches this schema: {json.dumps(schema_model)}"
             
             # Skip system messages for legacy o1 models as they don't support them
             if not self._needs_system_message_skip():
@@ -3771,20 +3776,46 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
             params.pop(param, None)
         
         if output_json or output_pydantic:
+            from .model_capabilities import supports_structured_outputs
+            schema_model = output_json or output_pydantic
             
-            # Check if this is a Gemini model that supports native structured outputs
-            if self._is_gemini_model():
-                from .model_capabilities import supports_structured_outputs
-                schema_model = output_json or output_pydantic
-                
-                if schema_model and hasattr(schema_model, 'model_json_schema') and supports_structured_outputs(self.model):
-                    schema = schema_model.model_json_schema()
-                    
-                    # Gemini uses response_mime_type and response_schema
-                    params['response_mime_type'] = 'application/json'
-                    params['response_schema'] = schema
-                    
-                    logging.debug(f"Using Gemini native structured output with schema: {json.dumps(schema, indent=2)}")
+            # Check if model supports structured outputs
+            if supports_structured_outputs(self.model):
+                # Check if this is a Gemini model (uses different params)
+                if self._is_gemini_model():
+                    if schema_model and hasattr(schema_model, 'model_json_schema'):
+                        schema = schema_model.model_json_schema()
+                        # Gemini uses response_mime_type and response_schema
+                        params['response_mime_type'] = 'application/json'
+                        params['response_schema'] = schema
+                        logging.debug(f"Using Gemini native structured output with schema: {json.dumps(schema, indent=2)}")
+                else:
+                    # OpenAI and other models use response_format
+                    # Handle Pydantic model
+                    if schema_model and hasattr(schema_model, 'model_json_schema'):
+                        schema = schema_model.model_json_schema()
+                        # Use json_schema format for strict validation
+                        params['response_format'] = {
+                            "type": "json_schema",
+                            "json_schema": {
+                                "name": getattr(schema_model, '__name__', 'response'),
+                                "schema": schema,
+                                "strict": True
+                            }
+                        }
+                        logging.debug(f"Using OpenAI structured output with schema: {json.dumps(schema, indent=2)}")
+                    # Handle dict schema (inline JSON schema from YAML)
+                    elif isinstance(schema_model, dict):
+                        # Inline JSON schema - wrap in response_format
+                        params['response_format'] = {
+                            "type": "json_schema",
+                            "json_schema": {
+                                "name": "structured_output",
+                                "schema": schema_model,
+                                "strict": True
+                            }
+                        }
+                        logging.debug(f"Using inline JSON schema: {json.dumps(schema_model, indent=2)}")
         
         # Add tool_choice="auto" when tools are provided (unless already specified)
         # This encourages the model to use tools when appropriate

src/praisonai-agents/praisonaiagents/llm/openai_client.py

@@ -342,9 +342,19 @@ class OpenAIClient:
         if system_prompt:
             # Append JSON schema if needed
             if output_json:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_json.model_json_schema())}"
+                # Handle Pydantic model
+                if hasattr(output_json, 'model_json_schema'):
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_json.model_json_schema())}"
+                # Handle inline dict schema (Option A from YAML)
+                elif isinstance(output_json, dict):
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this schema: {json.dumps(output_json)}"
             elif output_pydantic:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_pydantic.model_json_schema())}"
+                # Handle Pydantic model
+                if hasattr(output_pydantic, 'model_json_schema'):
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_pydantic.model_json_schema())}"
+                # Handle inline dict schema
+                elif isinstance(output_pydantic, dict):
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this schema: {json.dumps(output_pydantic)}"
             
             messages.append({"role": "system", "content": system_prompt})
         

src/praisonai-agents/praisonaiagents/mcp/__init__.py

@@ -54,6 +54,35 @@ def __getattr__(name):
     elif name == "get_thread_local_event_loop":
         from .mcp_schema_utils import get_thread_local_event_loop
         return get_thread_local_event_loop
+    # Auth storage (lazy loaded)
+    elif name == "MCPAuthStorage":
+        from .mcp_auth_storage import MCPAuthStorage
+        return MCPAuthStorage
+    elif name == "get_default_auth_filepath":
+        from .mcp_auth_storage import get_default_auth_filepath
+        return get_default_auth_filepath
+    # OAuth callback utilities (lazy loaded)
+    elif name == "OAuthCallbackHandler":
+        from .mcp_oauth_callback import OAuthCallbackHandler
+        return OAuthCallbackHandler
+    elif name == "generate_state":
+        from .mcp_oauth_callback import generate_state
+        return generate_state
+    elif name == "generate_code_verifier":
+        from .mcp_oauth_callback import generate_code_verifier
+        return generate_code_verifier
+    elif name == "generate_code_challenge":
+        from .mcp_oauth_callback import generate_code_challenge
+        return generate_code_challenge
+    elif name == "get_redirect_url":
+        from .mcp_oauth_callback import get_redirect_url
+        return get_redirect_url
+    elif name == "OAUTH_CALLBACK_PORT":
+        from .mcp_oauth_callback import OAUTH_CALLBACK_PORT
+        return OAUTH_CALLBACK_PORT
+    elif name == "OAUTH_CALLBACK_PATH":
+        from .mcp_oauth_callback import OAUTH_CALLBACK_PATH
+        return OAUTH_CALLBACK_PATH
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 __all__ = [
@@ -76,4 +105,15 @@ __all__ = [
     "SecurityConfig",
     "BaseTransport",
     "TransportConfig",
+    # Auth storage (lazy loaded)
+    "MCPAuthStorage",
+    "get_default_auth_filepath",
+    # OAuth callback utilities (lazy loaded)
+    "OAuthCallbackHandler",
+    "generate_state",
+    "generate_code_verifier",
+    "generate_code_challenge",
+    "get_redirect_url",
+    "OAUTH_CALLBACK_PORT",
+    "OAUTH_CALLBACK_PATH",
 ]

src/praisonai-agents/praisonaiagents/mcp/mcp_auth_storage.py

@@ -0,0 +1,294 @@
+"""
+MCP Authentication Storage.
+
+This module provides persistent storage for OAuth tokens and client information
+for MCP servers. It handles:
+- Access/refresh token storage
+- Client registration info (for dynamic registration)
+- PKCE code verifiers
+- OAuth state for CSRF protection
+
+Storage is file-based with restricted permissions (0600) for security.
+"""
+
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+def get_default_auth_filepath() -> str:
+    """
+    Get the default filepath for MCP auth storage.
+    
+    Returns:
+        Path to mcp-auth.json in user's praison config directory
+    """
+    # Use ~/.praison/ directory for config
+    home = Path.home()
+    praison_dir = home / ".praison"
+    praison_dir.mkdir(parents=True, exist_ok=True)
+    return str(praison_dir / "mcp-auth.json")
+
+
+class MCPAuthStorage:
+    """
+    Persistent storage for MCP OAuth credentials.
+    
+    This class manages storage of OAuth tokens, client information,
+    and temporary OAuth flow data (code verifiers, state) for MCP servers.
+    
+    Example:
+        ```python
+        from praisonaiagents.mcp.mcp_auth_storage import MCPAuthStorage
+        
+        storage = MCPAuthStorage()
+        
+        # Store tokens after OAuth flow
+        storage.set_tokens("github", {
+            "access_token": "gho_xxx",
+            "refresh_token": "ghr_xxx",
+            "expires_at": 1234567890
+        }, server_url="https://api.github.com/mcp")
+        
+        # Check if authenticated
+        entry = storage.get("github")
+        if entry and entry.get("tokens"):
+            print("Authenticated!")
+        ```
+    """
+    
+    def __init__(self, filepath: Optional[str] = None):
+        """
+        Initialize auth storage.
+        
+        Args:
+            filepath: Path to auth storage file. Defaults to ~/.praison/mcp-auth.json
+        """
+        self.filepath = filepath or get_default_auth_filepath()
+        self._ensure_directory()
+    
+    def _ensure_directory(self) -> None:
+        """Ensure the parent directory exists."""
+        parent = Path(self.filepath).parent
+        parent.mkdir(parents=True, exist_ok=True)
+    
+    def _read(self) -> Dict[str, Any]:
+        """Read all entries from storage file."""
+        try:
+            with open(self.filepath, 'r') as f:
+                return json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return {}
+    
+    def _write(self, data: Dict[str, Any]) -> None:
+        """Write all entries to storage file with secure permissions."""
+        # Write to file
+        with open(self.filepath, 'w') as f:
+            json.dump(data, f, indent=2)
+        
+        # Set file permissions to 0600 (owner read/write only)
+        if os.name != 'nt':  # Unix-like systems
+            os.chmod(self.filepath, 0o600)
+    
+    def get(self, mcp_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Get auth entry for an MCP server.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            
+        Returns:
+            Auth entry dict or None if not found
+        """
+        data = self._read()
+        return data.get(mcp_name)
+    
+    def get_for_url(self, mcp_name: str, server_url: str) -> Optional[Dict[str, Any]]:
+        """
+        Get auth entry and validate it's for the correct URL.
+        
+        Returns None if URL has changed (credentials are invalid).
+        
+        Args:
+            mcp_name: Name of the MCP server
+            server_url: Expected server URL
+            
+        Returns:
+            Auth entry if URL matches, None otherwise
+        """
+        entry = self.get(mcp_name)
+        if not entry:
+            return None
+        
+        # If no server_url stored, consider it invalid (old version)
+        if not entry.get("server_url"):
+            return None
+        
+        # If URL changed, credentials are invalid
+        if entry.get("server_url") != server_url:
+            return None
+        
+        return entry
+    
+    def all(self) -> Dict[str, Dict[str, Any]]:
+        """
+        Get all stored auth entries.
+        
+        Returns:
+            Dict mapping MCP server names to their auth entries
+        """
+        return self._read()
+    
+    def _set(self, mcp_name: str, entry: Dict[str, Any], server_url: Optional[str] = None) -> None:
+        """
+        Set an auth entry for an MCP server.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            entry: Auth entry dict
+            server_url: Optional server URL to track
+        """
+        data = self._read()
+        
+        # Update server_url if provided
+        if server_url:
+            entry["server_url"] = server_url
+        
+        data[mcp_name] = entry
+        self._write(data)
+    
+    def set_tokens(
+        self,
+        mcp_name: str,
+        tokens: Dict[str, Any],
+        server_url: Optional[str] = None
+    ) -> None:
+        """
+        Set OAuth tokens for an MCP server.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            tokens: Token dict with access_token, refresh_token, expires_at, scope
+            server_url: Server URL these tokens are for
+        """
+        entry = self.get(mcp_name) or {}
+        entry["tokens"] = tokens
+        self._set(mcp_name, entry, server_url)
+    
+    def set_client_info(
+        self,
+        mcp_name: str,
+        client_info: Dict[str, Any],
+        server_url: Optional[str] = None
+    ) -> None:
+        """
+        Set client info from dynamic registration.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            client_info: Client info dict with client_id, client_secret, etc.
+            server_url: Server URL this client is registered with
+        """
+        entry = self.get(mcp_name) or {}
+        entry["client_info"] = client_info
+        self._set(mcp_name, entry, server_url)
+    
+    def remove(self, mcp_name: str) -> None:
+        """
+        Remove auth entry for an MCP server.
+        
+        Args:
+            mcp_name: Name of the MCP server
+        """
+        data = self._read()
+        if mcp_name in data:
+            del data[mcp_name]
+            self._write(data)
+    
+    def set_code_verifier(self, mcp_name: str, code_verifier: str) -> None:
+        """
+        Set PKCE code verifier for OAuth flow.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            code_verifier: PKCE code verifier string
+        """
+        entry = self.get(mcp_name) or {}
+        entry["code_verifier"] = code_verifier
+        self._set(mcp_name, entry)
+    
+    def clear_code_verifier(self, mcp_name: str) -> None:
+        """
+        Clear PKCE code verifier after use.
+        
+        Args:
+            mcp_name: Name of the MCP server
+        """
+        entry = self.get(mcp_name)
+        if entry and "code_verifier" in entry:
+            del entry["code_verifier"]
+            self._set(mcp_name, entry)
+    
+    def set_oauth_state(self, mcp_name: str, state: str) -> None:
+        """
+        Set OAuth state for CSRF protection.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            state: Random state string
+        """
+        entry = self.get(mcp_name) or {}
+        entry["oauth_state"] = state
+        self._set(mcp_name, entry)
+    
+    def get_oauth_state(self, mcp_name: str) -> Optional[str]:
+        """
+        Get OAuth state for an MCP server.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            
+        Returns:
+            OAuth state string or None
+        """
+        entry = self.get(mcp_name)
+        if entry:
+            return entry.get("oauth_state")
+        return None
+    
+    def clear_oauth_state(self, mcp_name: str) -> None:
+        """
+        Clear OAuth state after use.
+        
+        Args:
+            mcp_name: Name of the MCP server
+        """
+        entry = self.get(mcp_name)
+        if entry and "oauth_state" in entry:
+            del entry["oauth_state"]
+            self._set(mcp_name, entry)
+    
+    def is_token_expired(self, mcp_name: str) -> Optional[bool]:
+        """
+        Check if stored tokens are expired.
+        
+        Args:
+            mcp_name: Name of the MCP server
+            
+        Returns:
+            None if no tokens exist
+            False if no expiry or not expired
+            True if expired
+        """
+        entry = self.get(mcp_name)
+        if not entry or not entry.get("tokens"):
+            return None
+        
+        tokens = entry["tokens"]
+        expires_at = tokens.get("expires_at")
+        if not expires_at:
+            return False
+        
+        return expires_at < time.time()

src/praisonai-agents/praisonaiagents/mcp/mcp_oauth_callback.py

@@ -0,0 +1,256 @@
+"""
+MCP OAuth Callback Handler.
+
+This module provides utilities for handling OAuth 2.1 callback flow for MCP servers.
+It includes:
+- Local callback server for receiving authorization codes
+- PKCE code verifier/challenge generation
+- OAuth state generation for CSRF protection
+
+The callback server listens on localhost for the OAuth redirect after
+browser-based authorization.
+"""
+
+import asyncio
+import base64
+import hashlib
+import secrets
+import threading
+from typing import Dict, Optional
+
+
+# Default callback configuration (same as OpenCode for consistency)
+OAUTH_CALLBACK_PORT = 19876
+OAUTH_CALLBACK_PATH = "/mcp/oauth/callback"
+
+
+def get_redirect_url(port: int = OAUTH_CALLBACK_PORT) -> str:
+    """
+    Get the OAuth redirect URL.
+    
+    Args:
+        port: Port for the callback server (default: 19876)
+        
+    Returns:
+        Full redirect URL for OAuth providers to whitelist
+    """
+    return f"http://127.0.0.1:{port}{OAUTH_CALLBACK_PATH}"
+
+
+def generate_state() -> str:
+    """
+    Generate a random OAuth state parameter for CSRF protection.
+    
+    Returns:
+        URL-safe random string
+    """
+    return secrets.token_urlsafe(32)
+
+
+def generate_code_verifier() -> str:
+    """
+    Generate a PKCE code verifier.
+    
+    Per RFC 7636, the code verifier is a high-entropy cryptographic
+    random string between 43 and 128 characters.
+    
+    Returns:
+        URL-safe random string (43-128 chars)
+    """
+    # Generate 32 bytes = 43 base64url characters (without padding)
+    return secrets.token_urlsafe(32)
+
+
+def generate_code_challenge(code_verifier: str) -> str:
+    """
+    Generate a PKCE code challenge from a code verifier.
+    
+    Uses S256 method: BASE64URL(SHA256(code_verifier))
+    
+    Args:
+        code_verifier: The code verifier string
+        
+    Returns:
+        Base64url-encoded SHA256 hash (no padding)
+    """
+    # SHA256 hash of the verifier
+    digest = hashlib.sha256(code_verifier.encode('ascii')).digest()
+    
+    # Base64url encode without padding
+    challenge = base64.urlsafe_b64encode(digest).decode('ascii')
+    
+    # Remove padding (= characters)
+    return challenge.rstrip('=')
+
+
+class OAuthCallbackHandler:
+    """
+    Handler for OAuth callback flow.
+    
+    This class manages the state of OAuth callbacks, allowing
+    the application to wait for authorization codes after
+    redirecting the user to the authorization URL.
+    
+    Example:
+        ```python
+        from praisonaiagents.mcp.mcp_oauth_callback import (
+            OAuthCallbackHandler, generate_state, get_redirect_url
+        )
+        
+        handler = OAuthCallbackHandler()
+        state = generate_state()
+        
+        # Build authorization URL with state and redirect_uri
+        auth_url = f"{auth_endpoint}?state={state}&redirect_uri={get_redirect_url()}"
+        
+        # Open browser for user authorization
+        webbrowser.open(auth_url)
+        
+        # Wait for callback (blocks until callback received or timeout)
+        code = handler.wait_for_callback(state, timeout=300)
+        
+        # Exchange code for tokens
+        tokens = exchange_code_for_tokens(code)
+        ```
+    """
+    
+    def __init__(self):
+        """Initialize the callback handler."""
+        self._callbacks: Dict[str, str] = {}  # state -> code
+        self._events: Dict[str, threading.Event] = {}  # state -> event
+        self._lock = threading.Lock()
+    
+    def receive_callback(self, state: str, code: str) -> None:
+        """
+        Receive an OAuth callback with authorization code.
+        
+        Called by the callback server when it receives a redirect.
+        
+        Args:
+            state: OAuth state parameter
+            code: Authorization code from the OAuth provider
+        """
+        with self._lock:
+            self._callbacks[state] = code
+            if state in self._events:
+                self._events[state].set()
+    
+    def get_code(self, state: str) -> Optional[str]:
+        """
+        Get the authorization code for a state.
+        
+        Args:
+            state: OAuth state parameter
+            
+        Returns:
+            Authorization code or None if not received
+        """
+        with self._lock:
+            return self._callbacks.get(state)
+    
+    def wait_for_callback(self, state: str, timeout: float = 300.0) -> str:
+        """
+        Wait for an OAuth callback with the given state.
+        
+        Blocks until the callback is received or timeout.
+        
+        Args:
+            state: OAuth state parameter to wait for
+            timeout: Maximum time to wait in seconds
+            
+        Returns:
+            Authorization code
+            
+        Raises:
+            TimeoutError: If callback not received within timeout
+            ValueError: If state doesn't match
+        """
+        # Check if already received
+        code = self.get_code(state)
+        if code:
+            return code
+        
+        # Create event for this state
+        with self._lock:
+            if state not in self._events:
+                self._events[state] = threading.Event()
+            event = self._events[state]
+        
+        # Wait for callback
+        if not event.wait(timeout=timeout):
+            raise TimeoutError(f"OAuth callback not received within {timeout} seconds")
+        
+        # Get the code
+        code = self.get_code(state)
+        if not code:
+            # Check if a different state was received
+            with self._lock:
+                if self._callbacks:
+                    received_states = list(self._callbacks.keys())
+                    raise ValueError(
+                        f"OAuth state mismatch. Expected '{state}', "
+                        f"received: {received_states}"
+                    )
+            raise ValueError(f"OAuth callback received but no code for state '{state}'")
+        
+        return code
+    
+    async def async_wait_for_callback(self, state: str, timeout: float = 300.0) -> str:
+        """
+        Async version of wait_for_callback.
+        
+        Args:
+            state: OAuth state parameter to wait for
+            timeout: Maximum time to wait in seconds
+            
+        Returns:
+            Authorization code
+            
+        Raises:
+            asyncio.TimeoutError: If callback not received within timeout
+            ValueError: If state doesn't match
+        """
+        # Check if already received
+        code = self.get_code(state)
+        if code:
+            return code
+        
+        # Create event for this state
+        with self._lock:
+            if state not in self._events:
+                self._events[state] = threading.Event()
+            event = self._events[state]
+        
+        # Poll for callback with async sleep
+        start_time = asyncio.get_event_loop().time()
+        while True:
+            if event.is_set():
+                break
+            
+            elapsed = asyncio.get_event_loop().time() - start_time
+            if elapsed >= timeout:
+                raise asyncio.TimeoutError(
+                    f"OAuth callback not received within {timeout} seconds"
+                )
+            
+            await asyncio.sleep(0.05)  # 50ms poll interval
+        
+        # Get the code
+        code = self.get_code(state)
+        if not code:
+            raise ValueError(f"OAuth callback received but no code for state '{state}'")
+        
+        return code
+    
+    def clear_state(self, state: str) -> None:
+        """
+        Clear callback data for a state.
+        
+        Should be called after successfully processing a callback.
+        
+        Args:
+            state: OAuth state parameter to clear
+        """
+        with self._lock:
+            self._callbacks.pop(state, None)
+            self._events.pop(state, None)

src/praisonai-agents/praisonaiagents/memory/learn/stores.py

@@ -11,13 +11,17 @@ Each store handles a specific type of learning data:
 - ImprovementStore: Self-improvement proposals
 """
 
-import json
-import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+# DRY: Import base storage for thread safety
+from ...storage.base import BaseJSONStore
+
+if TYPE_CHECKING:
+    from ...storage.protocols import StorageBackendProtocol
 
 
 @dataclass
@@ -50,18 +54,31 @@ class LearnEntry:
 
 
 class BaseStore(ABC):
-    """Abstract base class for learning stores."""
+    """
+    Abstract base class for learning stores.
+    
+    DRY: Uses BaseJSONStore internally for thread-safe, file-locked storage.
+    Supports pluggable backends (file, sqlite, etc.) via the backend parameter.
+    """
     
     def __init__(
         self,
         store_path: Optional[str] = None,
         user_id: Optional[str] = None,
         scope: str = "private",
+        backend: Optional["StorageBackendProtocol"] = None,
     ):
         self.user_id = user_id or "default"
         self.scope = scope
+        self._backend = backend
         self.store_path = store_path or self._default_path()
         self._entries: Dict[str, LearnEntry] = {}
+        
+        # DRY: Use BaseJSONStore for thread-safe storage
+        self._store = BaseJSONStore(
+            storage_path=self.store_path,
+            backend=backend,
+        )
         self._load()
     
     @property
@@ -77,26 +94,15 @@ class BaseStore(ABC):
         return str(base / f"{self.store_name}.json")
     
     def _load(self) -> None:
-        """Load entries from storage."""
-        if os.path.exists(self.store_path):
-            try:
-                with open(self.store_path, "r") as f:
-                    data = json.load(f)
-                    self._entries = {
-                        k: LearnEntry.from_dict(v) for k, v in data.items()
-                    }
-            except (json.JSONDecodeError, IOError):
-                self._entries = {}
+        """Load entries from storage using BaseJSONStore."""
+        data = self._store.load()
+        self._entries = {
+            k: LearnEntry.from_dict(v) for k, v in data.items()
+        }
     
     def _save(self) -> None:
-        """Save entries to storage."""
-        os.makedirs(os.path.dirname(self.store_path), exist_ok=True)
-        with open(self.store_path, "w") as f:
-            json.dump(
-                {k: v.to_dict() for k, v in self._entries.items()},
-                f,
-                indent=2,
-            )
+        """Save entries to storage using BaseJSONStore."""
+        self._store.save({k: v.to_dict() for k, v in self._entries.items()})
     
     def add(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> LearnEntry:
         """Add a new entry."""

src/praisonai-agents/praisonaiagents/memory/memory.py

@@ -463,52 +463,19 @@ class Memory:
             self._log_verbose(f"Error creating vector search indexes: {e}", logging.WARNING)
 
     def _get_embedding(self, text: str) -> List[float]:
-        """Get embedding for text using available embedding services."""
+        """Get embedding for text using the unified embedding module."""
         try:
-            if _check_litellm():
-                # Use LiteLLM for consistency with the rest of the codebase
-                import litellm
-                
-                response = litellm.embedding(
-                    model=self.embedding_model,
-                    input=text
-                )
-                return response.data[0]["embedding"]
-            elif _check_openai():
-                # Fallback to OpenAI client
-                from openai import OpenAI
-                client = OpenAI()
-                
-                response = client.embeddings.create(
-                    input=text,
-                    model=self.embedding_model
-                )
-                return response.data[0].embedding
-            else:
-                self._log_verbose("Neither litellm nor openai available for embeddings", logging.WARNING)
-                return None
+            from praisonaiagents.embedding import embedding
+            result = embedding(text, model=self.embedding_model)
+            return result.embeddings[0] if result.embeddings else None
         except Exception as e:
             self._log_verbose(f"Error getting embedding: {e}", logging.ERROR)
             return None
 
     def _get_embedding_dimensions(self, model_name: str) -> int:
         """Get embedding dimensions based on model name."""
-        # Common embedding model dimensions
-        model_dimensions = {
-            "text-embedding-3-small": 1536,
-            "text-embedding-3-large": 3072,
-            "text-embedding-ada-002": 1536,
-            "text-embedding-002": 1536,
-            # Add more models as needed
-        }
-        
-        # Check if model name contains known model identifiers
-        for model_key, dimensions in model_dimensions.items():
-            if model_key in model_name.lower():
-                return dimensions
-        
-        # Default to 1536 for unknown models (OpenAI standard)
-        return 1536
+        from praisonaiagents.embedding import get_dimensions
+        return get_dimensions(model_name)
 
     # -------------------------------------------------------------------------
     #                      Basic Quality Score Computation
@@ -690,27 +657,12 @@ class Memory:
             
         elif self.use_rag and hasattr(self, "chroma_col"):
             try:
-                if _check_litellm():
-                    # Use LiteLLM for consistency with the rest of the codebase
-                    import litellm
-                    
-                    response = litellm.embedding(
-                        model=self.embedding_model,
-                        input=query
-                    )
-                    query_embedding = response.data[0]["embedding"]
-                elif _check_openai():
-                    # Fallback to OpenAI client
-                    from openai import OpenAI
-                    client = OpenAI()
-                    
-                    response = client.embeddings.create(
-                        input=query,
-                        model=self.embedding_model
-                    )
-                    query_embedding = response.data[0].embedding
-                else:
-                    self._log_verbose("Neither litellm nor openai available for embeddings", logging.WARNING)
+                from praisonaiagents.embedding import embedding
+                result = embedding(query, model=self.embedding_model)
+                query_embedding = result.embeddings[0] if result.embeddings else None
+                
+                if query_embedding is None:
+                    self._log_verbose("Failed to get embedding for query", logging.WARNING)
                     return []
                 
                 resp = self.chroma_col.query(
@@ -853,39 +805,20 @@ class Memory:
         # Store in vector database if enabled
         if self.use_rag and hasattr(self, "chroma_col"):
             try:
-                if _check_litellm():
-                    # Use LiteLLM for consistency with the rest of the codebase
-                    import litellm
-                    
-                    logger.info("Getting embeddings from LiteLLM...")
-                    logger.trace(f"Embedding input text: {text}")
-                    
-                    response = litellm.embedding(
-                        model=self.embedding_model,
-                        input=text
-                    )
-                    embedding = response.data[0]["embedding"]
-                    logger.info("Successfully got embeddings from LiteLLM")
-                    logger.trace(f"Received embedding of length: {len(embedding)}")
-                    
-                elif _check_openai():
-                    # Fallback to OpenAI client
-                    from openai import OpenAI
-                    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-                    
-                    logger.info("Getting embeddings from OpenAI...")
-                    logger.trace(f"Embedding input text: {text}")
-                    
-                    response = client.embeddings.create(
-                        input=text,
-                        model=self.embedding_model
-                    )
-                    embedding = response.data[0].embedding
-                    logger.info("Successfully got embeddings from OpenAI")
-                    logger.trace(f"Received embedding of length: {len(embedding)}")
-                else:
-                    logger.warning("Neither litellm nor openai available for embeddings")
+                from praisonaiagents.embedding import embedding as get_embedding
+                
+                logger.info("Getting embeddings...")
+                logger.trace(f"Embedding input text: {text}")
+                
+                result = get_embedding(text, model=self.embedding_model)
+                embedding = result.embeddings[0] if result.embeddings else None
+                
+                if embedding is None:
+                    logger.warning("Failed to get embedding")
                     return
+                    
+                logger.info("Successfully got embeddings")
+                logger.trace(f"Received embedding of length: {len(embedding)}")
                 
                 # Sanitize metadata for ChromaDB
                 sanitized_metadata = self._sanitize_metadata(metadata)
@@ -1006,27 +939,12 @@ class Memory:
 
         elif self.use_rag and hasattr(self, "chroma_col"):
             try:
-                if _check_litellm():
-                    # Use LiteLLM for consistency with the rest of the codebase
-                    import litellm
-                    
-                    response = litellm.embedding(
-                        model=self.embedding_model,
-                        input=query
-                    )
-                    query_embedding = response.data[0]["embedding"]
-                elif _check_openai():
-                    # Fallback to OpenAI client
-                    from openai import OpenAI
-                    client = OpenAI()
-                    
-                    response = client.embeddings.create(
-                        input=query,
-                        model=self.embedding_model
-                    )
-                    query_embedding = response.data[0].embedding
-                else:
-                    self._log_verbose("Neither litellm nor openai available for embeddings", logging.WARNING)
+                from praisonaiagents.embedding import embedding as get_embedding
+                result = get_embedding(query, model=self.embedding_model)
+                query_embedding = result.embeddings[0] if result.embeddings else None
+                
+                if query_embedding is None:
+                    self._log_verbose("Failed to get embedding for query", logging.WARNING)
                     return []
                 
                 # Search ChromaDB with embedding

src/praisonai-agents/praisonaiagents/memory/workflows.py

@@ -1152,12 +1152,18 @@ class WorkflowManager:
         text: str,
         variables: Dict[str, Any]
     ) -> str:
-        """Substitute {{variable}} placeholders in text."""
-        def replace(match):
-            var_name = match.group(1).strip()
-            return str(variables.get(var_name, match.group(0)))
+        """
+        Substitute {{variable}} placeholders in text.
         
-        return re.sub(r'\{\{(\w+)\}\}', replace, text)
+        Delegates to shared substitute_variables() utility for DRY compliance.
+        
+        Resolution order:
+        1. Dynamic variable providers ({{now}}, {{today}}, {{uuid}}, etc.)
+        2. Static variables from workflow variables dict
+        3. Keep original placeholder if not found
+        """
+        from praisonaiagents.utils.variables import substitute_variables
+        return substitute_variables(text, variables)
     
     def _build_step_context(
         self,

src/praisonai-agents/praisonaiagents/storage/__init__.py

@@ -0,0 +1,66 @@
+"""
+PraisonAI Agents Storage Framework.
+
+Provides common storage abstractions for JSON-based persistence with:
+- Thread-safe operations
+- File locking for concurrent access
+- Common session info dataclass
+
+Zero performance impact when not in use through lazy loading.
+
+Example:
+    >>> from praisonaiagents.storage import BaseJSONStore, BaseSessionInfo
+    >>> class MyStore(BaseJSONStore):
+    ...     pass
+"""
+
+__all__ = [
+    "BaseJSONStore",
+    "AsyncBaseJSONStore",
+    "BaseSessionInfo",
+    "JSONStoreProtocol",
+    "SessionInfoProtocol",
+    "StorageBackendProtocol",
+    "AsyncStorageBackendProtocol",
+    "FileLock",
+    "list_json_sessions",
+    "cleanup_old_sessions",
+    # Backend implementations
+    "FileBackend",
+    "SQLiteBackend",
+    "RedisBackend",
+    "get_backend",
+]
+
+_LAZY_IMPORTS = {
+    "BaseJSONStore": ("base", "BaseJSONStore"),
+    "BaseSessionInfo": ("models", "BaseSessionInfo"),
+    "JSONStoreProtocol": ("protocols", "JSONStoreProtocol"),
+    "SessionInfoProtocol": ("protocols", "SessionInfoProtocol"),
+    "StorageBackendProtocol": ("protocols", "StorageBackendProtocol"),
+    "AsyncStorageBackendProtocol": ("protocols", "AsyncStorageBackendProtocol"),
+    "FileLock": ("base", "FileLock"),
+    "AsyncBaseJSONStore": ("base", "AsyncBaseJSONStore"),
+    "list_json_sessions": ("base", "list_json_sessions"),
+    "cleanup_old_sessions": ("base", "cleanup_old_sessions"),
+    # Backend implementations
+    "FileBackend": ("backends", "FileBackend"),
+    "SQLiteBackend": ("backends", "SQLiteBackend"),
+    "RedisBackend": ("backends", "RedisBackend"),
+    "get_backend": ("backends", "get_backend"),
+}
+
+
+def __getattr__(name: str):
+    """Lazy import mechanism for zero-cost imports when not used."""
+    if name in _LAZY_IMPORTS:
+        module_name, attr_name = _LAZY_IMPORTS[name]
+        import importlib
+        module = importlib.import_module(f".{module_name}", __name__)
+        return getattr(module, attr_name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Return list of available attributes for tab completion."""
+    return __all__

src/praisonai-agents/praisonaiagents/storage/backends.py

@@ -0,0 +1,495 @@
+"""
+Storage Backend Implementations for PraisonAI Agents.
+
+Provides concrete implementations of StorageBackendProtocol:
+- FileBackend: JSON file-based storage (default)
+- SQLiteBackend: SQLite database storage (optional, zero external deps)
+
+These backends enable switching between file-based and database-based storage
+without changing application code.
+
+DRY: All backends implement StorageBackendProtocol for consistent interface.
+"""
+
+import json
+import os
+import threading
+import tempfile
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class FileBackend:
+    """
+    JSON file-based storage backend.
+    
+    Default backend that stores each key as a separate JSON file.
+    Thread-safe with file locking support.
+    
+    Example:
+        ```python
+        backend = FileBackend(storage_dir="~/.praison/data")
+        backend.save("session_123", {"messages": []})
+        data = backend.load("session_123")
+        ```
+    """
+    
+    def __init__(
+        self,
+        storage_dir: str = "~/.praison/storage",
+        suffix: str = ".json",
+        pretty: bool = True,
+    ):
+        """
+        Initialize the file backend.
+        
+        Args:
+            storage_dir: Directory to store files
+            suffix: File suffix (default: .json)
+            pretty: Use pretty-printed JSON
+        """
+        self.storage_dir = Path(os.path.expanduser(storage_dir))
+        self.suffix = suffix
+        self.pretty = pretty
+        self._lock = threading.Lock()
+        
+        # Create directory if needed
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _key_to_path(self, key: str) -> Path:
+        """Convert key to file path."""
+        # Sanitize key for filesystem
+        safe_key = "".join(c if c.isalnum() or c in "-_" else "_" for c in key)
+        return self.storage_dir / f"{safe_key}{self.suffix}"
+    
+    def save(self, key: str, data: Dict[str, Any]) -> None:
+        """Save data with the given key."""
+        with self._lock:
+            file_path = self._key_to_path(key)
+            
+            # Atomic write via temp file
+            try:
+                with tempfile.NamedTemporaryFile(
+                    mode="w",
+                    encoding="utf-8",
+                    dir=str(self.storage_dir),
+                    delete=False,
+                    suffix=".tmp"
+                ) as f:
+                    if self.pretty:
+                        json.dump(data, f, indent=2, default=str, ensure_ascii=False)
+                    else:
+                        json.dump(data, f, default=str, ensure_ascii=False)
+                    temp_path = f.name
+                
+                os.replace(temp_path, file_path)
+            except Exception as e:
+                logger.error(f"Failed to save {key}: {e}")
+                try:
+                    if 'temp_path' in locals():
+                        os.remove(temp_path)
+                except Exception:
+                    pass
+                raise
+    
+    def load(self, key: str) -> Optional[Dict[str, Any]]:
+        """Load data by key."""
+        file_path = self._key_to_path(key)
+        
+        if not file_path.exists():
+            return None
+        
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError) as e:
+            logger.warning(f"Failed to load {key}: {e}")
+            return None
+    
+    def delete(self, key: str) -> bool:
+        """Delete data by key."""
+        with self._lock:
+            file_path = self._key_to_path(key)
+            
+            if file_path.exists():
+                try:
+                    file_path.unlink()
+                    return True
+                except Exception as e:
+                    logger.error(f"Failed to delete {key}: {e}")
+                    return False
+            return False
+    
+    def list_keys(self, prefix: str = "") -> List[str]:
+        """List all keys, optionally filtered by prefix."""
+        keys = []
+        
+        for file_path in self.storage_dir.iterdir():
+            if file_path.is_file() and file_path.suffix == self.suffix:
+                key = file_path.stem
+                if not prefix or key.startswith(prefix):
+                    keys.append(key)
+        
+        return sorted(keys)
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        return self._key_to_path(key).exists()
+    
+    def clear(self) -> int:
+        """Clear all data. Returns number of items deleted."""
+        count = 0
+        with self._lock:
+            for file_path in self.storage_dir.iterdir():
+                if file_path.is_file() and file_path.suffix == self.suffix:
+                    try:
+                        file_path.unlink()
+                        count += 1
+                    except Exception:
+                        pass
+        return count
+
+
+class SQLiteBackend:
+    """
+    SQLite-based storage backend.
+    
+    Uses Python's built-in sqlite3 module (zero external dependencies).
+    Better for concurrent access and large datasets.
+    
+    Example:
+        ```python
+        backend = SQLiteBackend(db_path="~/.praison/data.db")
+        backend.save("session_123", {"messages": []})
+        data = backend.load("session_123")
+        ```
+    """
+    
+    def __init__(
+        self,
+        db_path: str = "~/.praison/storage.db",
+        table_name: str = "praison_storage",
+        auto_create: bool = True,
+    ):
+        """
+        Initialize the SQLite backend.
+        
+        Args:
+            db_path: Path to SQLite database file
+            table_name: Name of the storage table
+            auto_create: Create table if it doesn't exist
+        """
+        self.db_path = os.path.expanduser(db_path)
+        self.table_name = table_name
+        self._local = threading.local()
+        
+        # Ensure directory exists
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+        
+        if auto_create:
+            self._create_table()
+    
+    def _get_conn(self):
+        """Get thread-local connection."""
+        # Lazy import sqlite3
+        import sqlite3
+        
+        if not hasattr(self._local, "conn") or self._local.conn is None:
+            self._local.conn = sqlite3.connect(
+                self.db_path,
+                check_same_thread=False
+            )
+            self._local.conn.row_factory = sqlite3.Row
+        return self._local.conn
+    
+    def _create_table(self) -> None:
+        """Create storage table if it doesn't exist."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        cur.execute(f"""
+            CREATE TABLE IF NOT EXISTS {self.table_name} (
+                key TEXT PRIMARY KEY,
+                data TEXT NOT NULL,
+                created_at REAL DEFAULT (strftime('%s', 'now')),
+                updated_at REAL DEFAULT (strftime('%s', 'now'))
+            )
+        """)
+        
+        # Index for prefix queries
+        cur.execute(f"""
+            CREATE INDEX IF NOT EXISTS idx_{self.table_name}_key 
+            ON {self.table_name}(key)
+        """)
+        
+        conn.commit()
+    
+    def save(self, key: str, data: Dict[str, Any]) -> None:
+        """Save data with the given key."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        json_data = json.dumps(data, default=str, ensure_ascii=False)
+        
+        cur.execute(f"""
+            INSERT INTO {self.table_name} (key, data, updated_at)
+            VALUES (?, ?, strftime('%s', 'now'))
+            ON CONFLICT(key) DO UPDATE SET
+                data = excluded.data,
+                updated_at = strftime('%s', 'now')
+        """, (key, json_data))
+        
+        conn.commit()
+    
+    def load(self, key: str) -> Optional[Dict[str, Any]]:
+        """Load data by key."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        cur.execute(f"""
+            SELECT data FROM {self.table_name} WHERE key = ?
+        """, (key,))
+        
+        row = cur.fetchone()
+        if row:
+            try:
+                return json.loads(row["data"])
+            except json.JSONDecodeError:
+                return None
+        return None
+    
+    def delete(self, key: str) -> bool:
+        """Delete data by key."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        cur.execute(f"""
+            DELETE FROM {self.table_name} WHERE key = ?
+        """, (key,))
+        
+        deleted = cur.rowcount > 0
+        conn.commit()
+        return deleted
+    
+    def list_keys(self, prefix: str = "") -> List[str]:
+        """List all keys, optionally filtered by prefix."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        if prefix:
+            cur.execute(f"""
+                SELECT key FROM {self.table_name}
+                WHERE key LIKE ?
+                ORDER BY key
+            """, (f"{prefix}%",))
+        else:
+            cur.execute(f"""
+                SELECT key FROM {self.table_name}
+                ORDER BY key
+            """)
+        
+        return [row["key"] for row in cur.fetchall()]
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        cur.execute(f"""
+            SELECT 1 FROM {self.table_name} WHERE key = ? LIMIT 1
+        """, (key,))
+        
+        return cur.fetchone() is not None
+    
+    def clear(self) -> int:
+        """Clear all data. Returns number of items deleted."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        
+        cur.execute(f"SELECT COUNT(*) as count FROM {self.table_name}")
+        count = cur.fetchone()["count"]
+        
+        cur.execute(f"DELETE FROM {self.table_name}")
+        conn.commit()
+        
+        return count
+    
+    def close(self) -> None:
+        """Close the database connection."""
+        if hasattr(self._local, "conn") and self._local.conn:
+            self._local.conn.close()
+            self._local.conn = None
+
+
+class RedisBackend:
+    """
+    Redis-based storage backend.
+    
+    Uses Redis for high-speed caching and ephemeral data storage.
+    Requires the `redis` package (optional dependency).
+    
+    Example:
+        ```python
+        backend = RedisBackend(url="redis://localhost:6379")
+        backend.save("session_123", {"messages": []})
+        data = backend.load("session_123")
+        ```
+    """
+    
+    def __init__(
+        self,
+        url: str = "redis://localhost:6379",
+        prefix: str = "praison:",
+        ttl: Optional[int] = None,
+        db: int = 0,
+    ):
+        """
+        Initialize the Redis backend.
+        
+        Args:
+            url: Redis connection URL
+            prefix: Key prefix for all stored data
+            ttl: Optional TTL in seconds for all keys
+            db: Redis database number
+        """
+        self.url = url
+        self.prefix = prefix
+        self.ttl = ttl
+        self.db = db
+        self._client = None
+    
+    def _get_client(self):
+        """Lazy initialize Redis client."""
+        if self._client is None:
+            try:
+                import redis
+            except ImportError:
+                raise ImportError(
+                    "Redis backend requires the 'redis' package. "
+                    "Install with: pip install redis"
+                )
+            self._client = redis.from_url(self.url, db=self.db)
+        return self._client
+    
+    def _make_key(self, key: str) -> str:
+        """Create prefixed key."""
+        return f"{self.prefix}{key}"
+    
+    def save(self, key: str, data: Dict[str, Any]) -> None:
+        """Save data with the given key."""
+        client = self._get_client()
+        full_key = self._make_key(key)
+        json_data = json.dumps(data, default=str, ensure_ascii=False)
+        
+        if self.ttl:
+            client.setex(full_key, self.ttl, json_data)
+        else:
+            client.set(full_key, json_data)
+    
+    def load(self, key: str) -> Optional[Dict[str, Any]]:
+        """Load data by key."""
+        client = self._get_client()
+        full_key = self._make_key(key)
+        
+        value = client.get(full_key)
+        if value:
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                return None
+        return None
+    
+    def delete(self, key: str) -> bool:
+        """Delete data by key."""
+        client = self._get_client()
+        full_key = self._make_key(key)
+        return client.delete(full_key) > 0
+    
+    def list_keys(self, prefix: str = "") -> List[str]:
+        """List all keys, optionally filtered by prefix."""
+        client = self._get_client()
+        pattern = self._make_key(f"{prefix}*")
+        
+        keys = []
+        for key in client.keys(pattern):
+            # Remove the prefix to return clean keys
+            key_str = key.decode() if isinstance(key, bytes) else key
+            clean_key = key_str[len(self.prefix):]
+            keys.append(clean_key)
+        
+        return sorted(keys)
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        client = self._get_client()
+        full_key = self._make_key(key)
+        return client.exists(full_key) > 0
+    
+    def clear(self) -> int:
+        """Clear all data with our prefix. Returns number of items deleted."""
+        client = self._get_client()
+        pattern = self._make_key("*")
+        keys = list(client.keys(pattern))
+        
+        if keys:
+            return client.delete(*keys)
+        return 0
+    
+    def set_ttl(self, key: str, ttl: int) -> bool:
+        """Set TTL on a specific key."""
+        client = self._get_client()
+        full_key = self._make_key(key)
+        return client.expire(full_key, ttl)
+    
+    def close(self) -> None:
+        """Close the Redis connection."""
+        if self._client:
+            self._client.close()
+            self._client = None
+
+
+def get_backend(
+    backend_type: str = "file",
+    **kwargs
+) -> Any:
+    """
+    Factory function to get a storage backend.
+    
+    Args:
+        backend_type: Type of backend ("file", "sqlite", or "redis")
+        **kwargs: Backend-specific arguments
+        
+    Returns:
+        Storage backend instance
+        
+    Example:
+        ```python
+        # File backend (default)
+        backend = get_backend("file", storage_dir="~/.praison/data")
+        
+        # SQLite backend
+        backend = get_backend("sqlite", db_path="~/.praison/data.db")
+        
+        # Redis backend
+        backend = get_backend("redis", url="redis://localhost:6379")
+        ```
+    """
+    if backend_type == "file":
+        return FileBackend(**kwargs)
+    elif backend_type == "sqlite":
+        return SQLiteBackend(**kwargs)
+    elif backend_type == "redis":
+        return RedisBackend(**kwargs)
+    else:
+        raise ValueError(f"Unknown backend type: {backend_type}")
+
+
+__all__ = [
+    'FileBackend',
+    'SQLiteBackend',
+    'RedisBackend',
+    'get_backend',
+]

src/praisonai-agents/praisonaiagents/storage/base.py

@@ -0,0 +1,537 @@
+"""
+Base JSON Store for PraisonAI Agents.
+
+Provides a common base class for JSON-based storage to eliminate duplication
+across BaseStore, TrainingStorage, and similar classes.
+
+DRY: This module extracts the common JSON load/save pattern with:
+- Thread-safe operations
+- File locking for concurrent access
+- Atomic writes
+"""
+
+import json
+import os
+import threading
+import tempfile
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
+from datetime import datetime
+
+from .models import BaseSessionInfo
+
+if TYPE_CHECKING:
+    from .protocols import StorageBackendProtocol
+
+logger = logging.getLogger(__name__)
+
+
+class FileLock:
+    """
+    Simple file-based lock for cross-process synchronization.
+    
+    Uses a .lock file alongside the target file.
+    """
+    
+    def __init__(self, path: Path, timeout: float = 10.0):
+        """
+        Initialize the file lock.
+        
+        Args:
+            path: Path to the file to lock
+            timeout: Maximum time to wait for lock (seconds)
+        """
+        self.lock_path = Path(str(path) + ".lock")
+        self.timeout = timeout
+        self._fd = None
+    
+    def __enter__(self):
+        """Acquire the lock."""
+        import time
+        start = time.time()
+        
+        while True:
+            try:
+                # Try to create lock file exclusively
+                self._fd = os.open(
+                    str(self.lock_path),
+                    os.O_CREAT | os.O_EXCL | os.O_WRONLY
+                )
+                break
+            except FileExistsError:
+                if time.time() - start > self.timeout:
+                    # Force remove stale lock
+                    try:
+                        self.lock_path.unlink()
+                    except Exception:
+                        pass
+                    continue
+                time.sleep(0.01)
+        
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Release the lock."""
+        if self._fd is not None:
+            try:
+                os.close(self._fd)
+            except Exception:
+                pass
+        try:
+            self.lock_path.unlink()
+        except Exception:
+            pass
+        return False
+
+
+class BaseJSONStore:
+    """
+    Base class for JSON-based storage.
+    
+    Provides thread-safe, file-locked JSON persistence with:
+    - Atomic writes (write to temp, then rename)
+    - Thread safety (threading.Lock)
+    - File locking (FileLock for cross-process)
+    - Common load/save pattern
+    - Optional pluggable backend (file, sqlite, etc.)
+    
+    Subclasses should:
+    - Set storage_path in __init__
+    - Override _default_data() for initial data structure
+    
+    Usage:
+        class MyStore(BaseJSONStore):
+            def __init__(self, path: Path):
+                super().__init__(path)
+            
+            def _default_data(self) -> Dict[str, Any]:
+                return {"items": []}
+            
+            def add_item(self, item: str):
+                with self._lock:
+                    self._data["items"].append(item)
+                    self._save()
+        
+        # With SQLite backend:
+        from praisonaiagents.storage import SQLiteBackend
+        store = MyStore(path, backend=SQLiteBackend("data.db"))
+    """
+    
+    def __init__(
+        self,
+        storage_path: Union[Path, str],
+        create_if_missing: bool = True,
+        use_file_lock: bool = True,
+        lock_timeout: float = 10.0,
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
+        """
+        Initialize the store.
+        
+        Args:
+            storage_path: Path to the JSON file (or key for backend)
+            create_if_missing: Create parent directories if needed
+            use_file_lock: Use file locking for cross-process safety
+            lock_timeout: Timeout for file lock acquisition
+            backend: Optional storage backend (file, sqlite, etc.)
+                     If provided, storage_path is used as the key.
+        """
+        self.storage_path = Path(storage_path)
+        self.use_file_lock = use_file_lock
+        self.lock_timeout = lock_timeout
+        self._lock = threading.Lock()
+        self._data: Dict[str, Any] = {}
+        self._backend = backend
+        
+        if backend is None:
+            # File-based storage (default)
+            if create_if_missing:
+                self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            if self.storage_path.exists():
+                self._load()
+            else:
+                self._data = self._default_data()
+        else:
+            # Backend-based storage
+            self._storage_key = self.storage_path.stem
+            loaded = backend.load(self._storage_key)
+            if loaded is not None:
+                self._data = loaded
+            else:
+                self._data = self._default_data()
+    
+    def _default_data(self) -> Dict[str, Any]:
+        """
+        Return default data structure for new stores.
+        
+        Override in subclasses for custom initial data.
+        """
+        return {}
+    
+    def _load(self) -> None:
+        """Load data from storage file or backend."""
+        if self._backend is not None:
+            loaded = self._backend.load(self._storage_key)
+            if loaded is not None:
+                self._data = loaded
+            else:
+                self._data = self._default_data()
+            return
+        
+        try:
+            if self.use_file_lock:
+                with FileLock(self.storage_path, self.lock_timeout):
+                    with open(self.storage_path, "r", encoding="utf-8") as f:
+                        self._data = json.load(f)
+            else:
+                with open(self.storage_path, "r", encoding="utf-8") as f:
+                    self._data = json.load(f)
+        except (json.JSONDecodeError, IOError) as e:
+            logger.warning(f"Failed to load {self.storage_path}: {e}")
+            self._data = self._default_data()
+    
+    def _save(self) -> None:
+        """Save data to storage file or backend with atomic write."""
+        if self._backend is not None:
+            try:
+                self._backend.save(self._storage_key, self._data)
+            except Exception as e:
+                logger.error(f"Failed to save to backend: {e}")
+                raise
+            return
+        
+        try:
+            dir_path = self.storage_path.parent
+            dir_path.mkdir(parents=True, exist_ok=True)
+            
+            # Write to temp file first
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=str(dir_path),
+                delete=False,
+                suffix=".tmp"
+            ) as f:
+                json.dump(self._data, f, indent=2, default=str, ensure_ascii=False)
+                temp_path = f.name
+            
+            # Atomic rename
+            if self.use_file_lock:
+                with FileLock(self.storage_path, self.lock_timeout):
+                    os.replace(temp_path, self.storage_path)
+            else:
+                os.replace(temp_path, self.storage_path)
+                
+        except Exception as e:
+            logger.error(f"Failed to save {self.storage_path}: {e}")
+            # Clean up temp file if it exists
+            try:
+                if 'temp_path' in locals():
+                    os.remove(temp_path)
+            except Exception:
+                pass
+            raise
+    
+    def load(self) -> Dict[str, Any]:
+        """
+        Load and return data from storage.
+        
+        Returns:
+            The stored data dictionary
+        """
+        with self._lock:
+            self._load()
+            return self._data.copy()
+    
+    def save(self, data: Dict[str, Any]) -> None:
+        """
+        Save data to storage.
+        
+        Args:
+            data: Data dictionary to save
+        """
+        with self._lock:
+            self._data = data
+            self._save()
+    
+    def exists(self) -> bool:
+        """Check if storage file or backend key exists."""
+        if self._backend is not None:
+            return self._backend.exists(self._storage_key)
+        return self.storage_path.exists()
+    
+    def clear(self) -> None:
+        """Clear all stored data."""
+        with self._lock:
+            self._data = self._default_data()
+            self._save()
+    
+    def delete(self) -> bool:
+        """
+        Delete the storage file or backend key.
+        
+        Returns:
+            True if deleted, False if file didn't exist
+        """
+        with self._lock:
+            if self._backend is not None:
+                result = self._backend.delete(self._storage_key)
+                if result:
+                    self._data = self._default_data()
+                return result
+            
+            try:
+                if self.storage_path.exists():
+                    self.storage_path.unlink()
+                    self._data = self._default_data()
+                    return True
+                return False
+            except Exception as e:
+                logger.error(f"Failed to delete {self.storage_path}: {e}")
+                return False
+
+
+def list_json_sessions(
+    storage_dir: Path,
+    suffix: str = ".json",
+    limit: int = 50,
+) -> List[BaseSessionInfo]:
+    """
+    List all sessions in a directory.
+    
+    Common utility for listing training sessions, traces, etc.
+    
+    Args:
+        storage_dir: Directory to search
+        suffix: File suffix to filter by
+        limit: Maximum number of sessions to return
+        
+    Returns:
+        List of BaseSessionInfo, sorted by modification time (newest first)
+    """
+    if not storage_dir.exists():
+        return []
+    
+    sessions = []
+    for file_path in storage_dir.iterdir():
+        if file_path.is_file() and file_path.suffix == suffix:
+            try:
+                info = BaseSessionInfo.from_path(file_path)
+                
+                # Try to get item count from file
+                if suffix == ".json":
+                    try:
+                        with open(file_path, "r") as f:
+                            data = json.load(f)
+                            # Common patterns for item count
+                            if "iterations" in data:
+                                info.item_count = len(data["iterations"])
+                            elif "messages" in data:
+                                info.item_count = len(data["messages"])
+                            elif "items" in data:
+                                info.item_count = len(data["items"])
+                    except Exception:
+                        pass
+                elif suffix == ".jsonl":
+                    try:
+                        with open(file_path, "r") as f:
+                            info.item_count = sum(1 for _ in f)
+                    except Exception:
+                        pass
+                
+                sessions.append(info)
+            except Exception as e:
+                logger.warning(f"Failed to stat {file_path}: {e}")
+    
+    # Sort by modification time (newest first)
+    sessions.sort(key=lambda s: s.modified_at, reverse=True)
+    
+    return sessions[:limit]
+
+
+def cleanup_old_sessions(
+    storage_dir: Path,
+    suffix: str = ".json",
+    max_age_days: int = 30,
+    max_size_mb: int = 100,
+) -> int:
+    """
+    Clean up old session files.
+    
+    Common utility for cleaning up training sessions, traces, etc.
+    
+    Args:
+        storage_dir: Directory to clean
+        suffix: File suffix to filter by
+        max_age_days: Delete sessions older than this
+        max_size_mb: Delete oldest sessions if total size exceeds this
+        
+    Returns:
+        Number of files deleted
+    """
+    if not storage_dir.exists():
+        return 0
+    
+    deleted = 0
+    now = datetime.now()
+    
+    # Get all sessions sorted by age (oldest first)
+    sessions = list_json_sessions(storage_dir, suffix, limit=10000)
+    sessions.sort(key=lambda s: s.modified_at)
+    
+    # Delete old sessions
+    for session in sessions:
+        age_days = (now - session.modified_at).days
+        if age_days > max_age_days:
+            try:
+                session.path.unlink()
+                deleted += 1
+            except Exception:
+                pass
+    
+    # Check total size and delete oldest if needed
+    remaining = [s for s in sessions if s.path.exists()]
+    total_size_mb = sum(s.size_bytes for s in remaining) / (1024 * 1024)
+    
+    if total_size_mb > max_size_mb:
+        for session in remaining:
+            if total_size_mb <= max_size_mb:
+                break
+            try:
+                session.path.unlink()
+                total_size_mb -= session.size_bytes / (1024 * 1024)
+                deleted += 1
+            except Exception:
+                pass
+    
+    return deleted
+
+
+class AsyncBaseJSONStore:
+    """
+    Async version of BaseJSONStore.
+    
+    Provides async-safe JSON persistence with optional pluggable backends.
+    Uses asyncio locks for thread safety in async contexts.
+    
+    Usage:
+        store = AsyncBaseJSONStore("data.json")
+        await store.load_async()
+        await store.save_async({"key": "value"})
+    """
+    
+    def __init__(
+        self,
+        storage_path: Union[Path, str],
+        create_if_missing: bool = True,
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
+        """
+        Initialize the async store.
+        
+        Args:
+            storage_path: Path to the JSON file (or key for backend)
+            create_if_missing: Create parent directories if needed
+            backend: Optional async storage backend
+        """
+        self.storage_path = Path(storage_path)
+        self._backend = backend
+        self._data: Dict[str, Any] = {}
+        self._async_lock: Optional[Any] = None  # Lazy init asyncio.Lock
+        
+        if backend is None and create_if_missing:
+            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        self._storage_key = self.storage_path.stem
+    
+    def _get_async_lock(self):
+        """Lazy initialize asyncio lock."""
+        if self._async_lock is None:
+            import asyncio
+            self._async_lock = asyncio.Lock()
+        return self._async_lock
+    
+    def _default_data(self) -> Dict[str, Any]:
+        """Return default data structure."""
+        return {}
+    
+    async def load_async(self) -> Dict[str, Any]:
+        """Load data asynchronously."""
+        async with self._get_async_lock():
+            if self._backend is not None:
+                # Check if backend has async method
+                if hasattr(self._backend, 'load') and callable(self._backend.load):
+                    import asyncio
+                    if asyncio.iscoroutinefunction(self._backend.load):
+                        loaded = await self._backend.load(self._storage_key)
+                    else:
+                        loaded = self._backend.load(self._storage_key)
+                    self._data = loaded if loaded is not None else self._default_data()
+            else:
+                # File-based async read
+                import aiofiles
+                try:
+                    async with aiofiles.open(self.storage_path, "r", encoding="utf-8") as f:
+                        content = await f.read()
+                        self._data = json.loads(content)
+                except (FileNotFoundError, json.JSONDecodeError):
+                    self._data = self._default_data()
+            
+            return self._data.copy()
+    
+    async def save_async(self, data: Dict[str, Any]) -> None:
+        """Save data asynchronously."""
+        async with self._get_async_lock():
+            self._data = data
+            
+            if self._backend is not None:
+                import asyncio
+                if asyncio.iscoroutinefunction(self._backend.save):
+                    await self._backend.save(self._storage_key, self._data)
+                else:
+                    self._backend.save(self._storage_key, self._data)
+            else:
+                # File-based async write
+                import aiofiles
+                self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+                async with aiofiles.open(self.storage_path, "w", encoding="utf-8") as f:
+                    await f.write(json.dumps(self._data, indent=2, default=str, ensure_ascii=False))
+    
+    async def exists_async(self) -> bool:
+        """Check if storage exists asynchronously."""
+        if self._backend is not None:
+            import asyncio
+            if asyncio.iscoroutinefunction(self._backend.exists):
+                return await self._backend.exists(self._storage_key)
+            return self._backend.exists(self._storage_key)
+        return self.storage_path.exists()
+    
+    async def delete_async(self) -> bool:
+        """Delete storage asynchronously."""
+        async with self._get_async_lock():
+            if self._backend is not None:
+                import asyncio
+                if asyncio.iscoroutinefunction(self._backend.delete):
+                    result = await self._backend.delete(self._storage_key)
+                else:
+                    result = self._backend.delete(self._storage_key)
+                if result:
+                    self._data = self._default_data()
+                return result
+            
+            if self.storage_path.exists():
+                self.storage_path.unlink()
+                self._data = self._default_data()
+                return True
+            return False
+
+
+__all__ = [
+    'FileLock',
+    'BaseJSONStore',
+    'AsyncBaseJSONStore',
+    'list_json_sessions',
+    'cleanup_old_sessions',
+]

src/praisonai-agents/praisonaiagents/storage/models.py

@@ -0,0 +1,101 @@
+"""
+Storage Models for PraisonAI Agents.
+
+Provides common dataclasses for session/trace info to eliminate duplication
+across TrainingSessionInfo, TraceInfo, and similar classes.
+
+DRY: This module extracts the common session metadata pattern.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class BaseSessionInfo:
+    """
+    Base class for session/trace info.
+    
+    Provides common fields for session metadata used by:
+    - TrainingSessionInfo (train/agents/storage.py)
+    - TraceInfo (replay/storage.py)
+    - SessionData (session/store.py)
+    
+    Attributes:
+        session_id: Unique session identifier
+        path: Path to the session file
+        size_bytes: File size in bytes
+        created_at: When the session was created
+        modified_at: When the session was last modified
+        item_count: Number of items (iterations, events, messages, etc.)
+    """
+    session_id: str
+    path: Path
+    size_bytes: int = 0
+    created_at: datetime = field(default_factory=datetime.now)
+    modified_at: datetime = field(default_factory=datetime.now)
+    item_count: int = 0
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "session_id": self.session_id,
+            "path": str(self.path),
+            "size_bytes": self.size_bytes,
+            "created_at": self.created_at.isoformat(),
+            "modified_at": self.modified_at.isoformat(),
+            "item_count": self.item_count,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "BaseSessionInfo":
+        """Create from dictionary."""
+        created_at = data.get("created_at")
+        modified_at = data.get("modified_at")
+        
+        if isinstance(created_at, str):
+            created_at = datetime.fromisoformat(created_at)
+        elif created_at is None:
+            created_at = datetime.now()
+            
+        if isinstance(modified_at, str):
+            modified_at = datetime.fromisoformat(modified_at)
+        elif modified_at is None:
+            modified_at = datetime.now()
+        
+        return cls(
+            session_id=data["session_id"],
+            path=Path(data["path"]),
+            size_bytes=data.get("size_bytes", 0),
+            created_at=created_at,
+            modified_at=modified_at,
+            item_count=data.get("item_count", 0),
+        )
+    
+    @classmethod
+    def from_path(cls, path: Path, session_id: Optional[str] = None) -> "BaseSessionInfo":
+        """
+        Create from a file path.
+        
+        Args:
+            path: Path to the session file
+            session_id: Optional session ID (defaults to file stem)
+            
+        Returns:
+            BaseSessionInfo with file metadata
+        """
+        stat = path.stat()
+        return cls(
+            session_id=session_id or path.stem,
+            path=path,
+            size_bytes=stat.st_size,
+            created_at=datetime.fromtimestamp(stat.st_ctime),
+            modified_at=datetime.fromtimestamp(stat.st_mtime),
+        )
+
+
+__all__ = [
+    'BaseSessionInfo',
+]

src/praisonai-agents/praisonaiagents/storage/protocols.py

@@ -0,0 +1,248 @@
+"""
+Storage Protocol Definitions for PraisonAI Agents.
+
+Provides Protocol interfaces for storage implementations to enable:
+- DRY: Common storage patterns shared across modules
+- Mocking: Easy testing without real file I/O
+- Extensibility: Custom storage backends
+
+These protocols are lightweight and have zero performance impact.
+"""
+from typing import Protocol, runtime_checkable, List, Dict, Any, TypeVar
+from pathlib import Path
+from datetime import datetime
+
+
+T = TypeVar('T')
+
+
+@runtime_checkable
+class SessionInfoProtocol(Protocol):
+    """
+    Protocol for session/trace info.
+    
+    Defines the minimal interface for session metadata.
+    Used by TrainingSessionInfo, TraceInfo, and similar classes.
+    """
+    session_id: str
+    path: Path
+    size_bytes: int
+    created_at: datetime
+    modified_at: datetime
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        ...
+
+
+@runtime_checkable
+class JSONStoreProtocol(Protocol):
+    """
+    Protocol for JSON-based storage.
+    
+    Defines the interface for persistent JSON storage.
+    Implementations include BaseStore, TrainingStorage, etc.
+    
+    Example:
+        ```python
+        class MockStore:
+            storage_path = Path("/tmp/mock.json")
+            
+            def load(self):
+                return {}
+            
+            def save(self, data):
+                pass
+            
+            def exists(self):
+                return True
+        
+        # Use in tests
+        store: JSONStoreProtocol = MockStore()
+        ```
+    """
+    storage_path: Path
+    
+    def load(self) -> Dict[str, Any]:
+        """Load data from storage."""
+        ...
+    
+    def save(self, data: Dict[str, Any]) -> None:
+        """Save data to storage."""
+        ...
+    
+    def exists(self) -> bool:
+        """Check if storage file exists."""
+        ...
+
+
+@runtime_checkable
+class ListableStoreProtocol(Protocol):
+    """
+    Protocol for stores that support listing sessions.
+    """
+    
+    def list_sessions(self, limit: int = 50) -> List[Any]:
+        """List all sessions with metadata."""
+        ...
+
+
+@runtime_checkable
+class CleanableStoreProtocol(Protocol):
+    """
+    Protocol for stores that support cleanup.
+    """
+    
+    def cleanup(self, max_age_days: int = 30, max_size_mb: int = 100) -> int:
+        """
+        Clean up old entries.
+        
+        Args:
+            max_age_days: Delete entries older than this
+            max_size_mb: Delete oldest if total size exceeds this
+            
+        Returns:
+            Number of entries deleted
+        """
+        ...
+
+
+@runtime_checkable
+class StorageBackendProtocol(Protocol):
+    """
+    Protocol for pluggable storage backends.
+    
+    Enables switching between file-based and database-based storage.
+    Default implementation uses JSON files, but can be swapped for:
+    - SQLite (lightweight, local)
+    - PostgreSQL (production, scalable)
+    - Redis (fast, ephemeral)
+    - MongoDB (document-oriented)
+    
+    PraisonAI supports many databases through praisonaiagents.db adapters.
+    This protocol allows storage classes to use any backend.
+    
+    Example:
+        ```python
+        class SQLiteBackend:
+            def __init__(self, db_path: str):
+                self.db_path = db_path
+            
+            def save(self, key: str, data: Dict[str, Any]) -> None:
+                # Save to SQLite
+                pass
+            
+            def load(self, key: str) -> Optional[Dict[str, Any]]:
+                # Load from SQLite
+                return {}
+            
+            def delete(self, key: str) -> bool:
+                return True
+            
+            def list_keys(self, prefix: str = "") -> List[str]:
+                return []
+            
+            def exists(self, key: str) -> bool:
+                return True
+        
+        # Use with any storage class
+        backend: StorageBackendProtocol = SQLiteBackend("data.db")
+        ```
+    """
+    
+    def save(self, key: str, data: Dict[str, Any]) -> None:
+        """
+        Save data with the given key.
+        
+        Args:
+            key: Unique identifier for the data
+            data: Dictionary to save
+        """
+        ...
+    
+    def load(self, key: str) -> Any:
+        """
+        Load data by key.
+        
+        Args:
+            key: Unique identifier for the data
+            
+        Returns:
+            The stored data, or None if not found
+        """
+        ...
+    
+    def delete(self, key: str) -> bool:
+        """
+        Delete data by key.
+        
+        Args:
+            key: Unique identifier for the data
+            
+        Returns:
+            True if deleted, False if not found
+        """
+        ...
+    
+    def list_keys(self, prefix: str = "") -> List[str]:
+        """
+        List all keys, optionally filtered by prefix.
+        
+        Args:
+            prefix: Optional prefix to filter keys
+            
+        Returns:
+            List of matching keys
+        """
+        ...
+    
+    def exists(self, key: str) -> bool:
+        """
+        Check if a key exists.
+        
+        Args:
+            key: Unique identifier to check
+            
+        Returns:
+            True if exists, False otherwise
+        """
+        ...
+
+
+@runtime_checkable
+class AsyncStorageBackendProtocol(Protocol):
+    """
+    Async Protocol for pluggable storage backends.
+    
+    Async version of StorageBackendProtocol for use in async contexts.
+    """
+    
+    async def save(self, key: str, data: Dict[str, Any]) -> None:
+        """Save data with the given key."""
+        ...
+    
+    async def load(self, key: str) -> Any:
+        """Load data by key."""
+        ...
+    
+    async def delete(self, key: str) -> bool:
+        """Delete data by key."""
+        ...
+    
+    async def list_keys(self, prefix: str = "") -> List[str]:
+        """List all keys, optionally filtered by prefix."""
+        ...
+    
+    async def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        ...
+
+
+__all__ = [
+    'SessionInfoProtocol',
+    'JSONStoreProtocol',
+    'ListableStoreProtocol',
+    'CleanableStoreProtocol',
+    'StorageBackendProtocol',
+    'AsyncStorageBackendProtocol',
+]

src/praisonai-agents/praisonaiagents/trace/__init__.py

@@ -44,6 +44,13 @@ __all__ = [
     # Redaction
     "redact_dict",
     "REDACT_KEYS",
+    # Context events (for replay)
+    "ContextEvent",
+    "ContextEventType",
+    "ContextTraceSink",
+    "ContextNoOpSink",
+    "ContextListSink",
+    "ContextTraceEmitter",
 ]
 
 
@@ -65,4 +72,17 @@ def __getattr__(name: str):
         from .redact import redact_dict, REDACT_KEYS
         return locals()[name]
     
+    # Context events (for replay)
+    if name in ("ContextEvent", "ContextEventType"):
+        from .context_events import ContextEvent, ContextEventType
+        return locals()[name]
+    
+    if name in ("ContextTraceSink", "ContextNoOpSink", "ContextListSink"):
+        from .context_events import ContextTraceSink, ContextNoOpSink, ContextListSink
+        return locals()[name]
+    
+    if name == "ContextTraceEmitter":
+        from .context_events import ContextTraceEmitter
+        return ContextTraceEmitter
+    
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai-agents/praisonaiagents/trace/context_events.py

@@ -0,0 +1,481 @@
+"""
+Context Events for PraisonAI Agents.
+
+Provides context-level tracing for replay functionality.
+Tracks context changes during agent execution for debugging and replay.
+
+Zero Performance Impact:
+- NoOpSink is the default (zero overhead when not used)
+- Disabled emitter has near-zero overhead
+- Uses __slots__ for memory efficiency
+
+Schema Version: 1.0
+"""
+
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+import json
+
+
+CONTEXT_SCHEMA_VERSION = "1.0"
+
+
+class ContextEventType(str, Enum):
+    """Types of context events for replay."""
+    SESSION_START = "session_start"
+    SESSION_END = "session_end"
+    AGENT_START = "agent_start"
+    AGENT_END = "agent_end"
+    AGENT_HANDOFF = "agent_handoff"
+    MESSAGE_ADDED = "message_added"
+    TOOL_CALL_START = "tool_call_start"
+    TOOL_CALL_END = "tool_call_end"
+    LLM_REQUEST = "llm_request"
+    LLM_RESPONSE = "llm_response"
+    CONTEXT_SNAPSHOT = "context_snapshot"
+
+
+@dataclass
+class ContextEvent:
+    """
+    A single context event for replay.
+    
+    Represents a discrete change in agent context during execution.
+    Designed for minimal memory footprint and fast serialization.
+    
+    Attributes:
+        event_type: Type of context event
+        timestamp: Unix timestamp when event occurred
+        session_id: Session identifier for grouping events
+        agent_name: Name of the agent (if applicable)
+        sequence_num: Sequential event number within session
+        messages_count: Number of messages in context at this point
+        tokens_used: Tokens used in context at this point
+        tokens_budget: Total token budget available
+        data: Event-specific data (tool args, message content, etc.)
+    """
+    event_type: ContextEventType
+    timestamp: float
+    session_id: str
+    agent_name: Optional[str] = None
+    sequence_num: int = 0
+    messages_count: int = 0
+    tokens_used: int = 0
+    tokens_budget: int = 0
+    data: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "schema_version": CONTEXT_SCHEMA_VERSION,
+            "event_type": self.event_type.value if isinstance(self.event_type, ContextEventType) else self.event_type,
+            "timestamp": self.timestamp,
+            "session_id": self.session_id,
+            "agent_name": self.agent_name,
+            "sequence_num": self.sequence_num,
+            "messages_count": self.messages_count,
+            "tokens_used": self.tokens_used,
+            "tokens_budget": self.tokens_budget,
+            "data": self.data,
+        }
+    
+    def to_json(self) -> str:
+        """Convert to JSON string."""
+        return json.dumps(self.to_dict(), default=str)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ContextEvent":
+        """Create ContextEvent from dictionary."""
+        event_type_str = data.get("event_type", "")
+        
+        # Convert string to enum
+        try:
+            event_type = ContextEventType(event_type_str)
+        except ValueError:
+            # Fallback for unknown event types
+            event_type = ContextEventType.CONTEXT_SNAPSHOT
+        
+        return cls(
+            event_type=event_type,
+            timestamp=data.get("timestamp", 0.0),
+            session_id=data.get("session_id", ""),
+            agent_name=data.get("agent_name"),
+            sequence_num=data.get("sequence_num", 0),
+            messages_count=data.get("messages_count", 0),
+            tokens_used=data.get("tokens_used", 0),
+            tokens_budget=data.get("tokens_budget", 0),
+            data=data.get("data", {}),
+        )
+
+
+@runtime_checkable
+class ContextTraceSink(Protocol):
+    """
+    Protocol for context trace event sinks.
+    
+    Implementations receive context events and handle them
+    (e.g., write to file, collect in memory, send to server).
+    """
+    
+    def emit(self, event: ContextEvent) -> None:
+        """Emit a context event."""
+        ...
+    
+    def flush(self) -> None:
+        """Flush any buffered events."""
+        ...
+    
+    def close(self) -> None:
+        """Close the sink and release resources."""
+        ...
+
+
+class ContextNoOpSink:
+    """
+    No-operation sink that discards all events.
+    
+    Used as the default sink when context tracing is disabled.
+    Has near-zero overhead.
+    """
+    
+    __slots__ = ()
+    
+    def emit(self, event: ContextEvent) -> None:
+        """Discard the event."""
+        pass
+    
+    def flush(self) -> None:
+        """No-op."""
+        pass
+    
+    def close(self) -> None:
+        """No-op."""
+        pass
+
+
+class ContextListSink:
+    """
+    Sink that collects events in a list.
+    
+    Useful for testing and programmatic access to context events.
+    """
+    
+    __slots__ = ("_events",)
+    
+    def __init__(self):
+        self._events: List[ContextEvent] = []
+    
+    def emit(self, event: ContextEvent) -> None:
+        """Add event to the list."""
+        self._events.append(event)
+    
+    def flush(self) -> None:
+        """No-op for list sink."""
+        pass
+    
+    def close(self) -> None:
+        """No-op for list sink."""
+        pass
+    
+    def get_events(self) -> List[ContextEvent]:
+        """Get all collected events."""
+        return self._events.copy()
+    
+    def clear(self) -> None:
+        """Clear all collected events."""
+        self._events.clear()
+    
+    def __len__(self) -> int:
+        """Return number of events."""
+        return len(self._events)
+    
+    def __iter__(self):
+        """Iterate over events."""
+        return iter(self._events)
+
+
+class ContextTraceEmitter:
+    """
+    Emitter for context trace events.
+    
+    Provides a high-level API for emitting context events during
+    agent execution. Handles sequence numbering and redaction.
+    
+    Usage:
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="my-session")
+        
+        emitter.session_start()
+        emitter.agent_start("researcher")
+        emitter.message_added("researcher", "user", "Hello", 1, 10)
+        emitter.agent_end("researcher")
+        emitter.session_end()
+    """
+    
+    __slots__ = ("_sink", "_session_id", "_enabled", "_redact", "_sequence")
+    
+    def __init__(
+        self,
+        sink: Optional[ContextTraceSink] = None,
+        session_id: str = "",
+        enabled: bool = True,
+        redact: bool = True,
+    ):
+        """
+        Initialize context trace emitter.
+        
+        Args:
+            sink: Sink to emit events to (default: ContextNoOpSink)
+            session_id: Session identifier for all events
+            enabled: Whether tracing is enabled
+            redact: Whether to redact sensitive data
+        """
+        self._sink = sink if sink is not None else ContextNoOpSink()
+        self._session_id = session_id
+        self._enabled = enabled
+        self._redact = redact
+        self._sequence = 0
+    
+    def _emit(self, event: ContextEvent) -> None:
+        """Internal emit with enabled check and sequence assignment."""
+        if not self._enabled:
+            return
+        event.sequence_num = self._sequence
+        self._sequence += 1
+        self._sink.emit(event)
+    
+    def _redact_dict(self, data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+        """Redact sensitive data from dictionary if enabled."""
+        if data is None or not self._redact:
+            return data
+        
+        from .redact import redact_dict
+        return redact_dict(data)
+    
+    def session_start(self, metadata: Optional[Dict[str, Any]] = None) -> None:
+        """Emit session start event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.SESSION_START,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            data=metadata or {},
+        ))
+    
+    def session_end(self, metadata: Optional[Dict[str, Any]] = None) -> None:
+        """Emit session end event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.SESSION_END,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            data=metadata or {},
+        ))
+    
+    def agent_start(
+        self,
+        agent_name: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Emit agent start event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            data=metadata or {},
+        ))
+    
+    def agent_end(
+        self,
+        agent_name: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Emit agent end event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.AGENT_END,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            data=metadata or {},
+        ))
+    
+    def message_added(
+        self,
+        agent_name: str,
+        role: str,
+        content: str,
+        messages_count: int,
+        tokens_used: int,
+        tokens_budget: int = 0,
+    ) -> None:
+        """Emit message added event."""
+        # Truncate long content for storage
+        truncated_content = content[:1000] + "..." if len(content) > 1000 else content
+        
+        self._emit(ContextEvent(
+            event_type=ContextEventType.MESSAGE_ADDED,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            messages_count=messages_count,
+            tokens_used=tokens_used,
+            tokens_budget=tokens_budget,
+            data={
+                "role": role,
+                "content": truncated_content,
+            },
+        ))
+    
+    def tool_call_start(
+        self,
+        agent_name: str,
+        tool_name: str,
+        tool_args: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Emit tool call start event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.TOOL_CALL_START,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            data={
+                "tool_name": tool_name,
+                "tool_args": self._redact_dict(tool_args) if tool_args else {},
+            },
+        ))
+    
+    def tool_call_end(
+        self,
+        agent_name: str,
+        tool_name: str,
+        result: Optional[str] = None,
+        duration_ms: float = 0.0,
+        error: Optional[str] = None,
+    ) -> None:
+        """Emit tool call end event."""
+        # Truncate long results
+        truncated_result = None
+        if result:
+            truncated_result = result[:500] + "..." if len(result) > 500 else result
+        
+        self._emit(ContextEvent(
+            event_type=ContextEventType.TOOL_CALL_END,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            data={
+                "tool_name": tool_name,
+                "result": truncated_result,
+                "duration_ms": duration_ms,
+                "error": error,
+            },
+        ))
+    
+    def llm_request(
+        self,
+        agent_name: str,
+        messages_count: int,
+        tokens_used: int,
+        tokens_budget: int = 0,
+        model: Optional[str] = None,
+    ) -> None:
+        """Emit LLM request event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.LLM_REQUEST,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            messages_count=messages_count,
+            tokens_used=tokens_used,
+            tokens_budget=tokens_budget,
+            data={
+                "model": model,
+            },
+        ))
+    
+    def llm_response(
+        self,
+        agent_name: str,
+        response_tokens: int = 0,
+        duration_ms: float = 0.0,
+        finish_reason: Optional[str] = None,
+    ) -> None:
+        """Emit LLM response event."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.LLM_RESPONSE,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            data={
+                "response_tokens": response_tokens,
+                "duration_ms": duration_ms,
+                "finish_reason": finish_reason,
+            },
+        ))
+    
+    def context_snapshot(
+        self,
+        agent_name: str,
+        messages_count: int,
+        tokens_used: int,
+        tokens_budget: int,
+        messages: Optional[List[Dict[str, Any]]] = None,
+    ) -> None:
+        """Emit context snapshot event with full context state."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.CONTEXT_SNAPSHOT,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=agent_name,
+            messages_count=messages_count,
+            tokens_used=tokens_used,
+            tokens_budget=tokens_budget,
+            data={
+                "messages": messages or [],
+            },
+        ))
+    
+    def agent_handoff(
+        self,
+        from_agent: str,
+        to_agent: str,
+        reason: Optional[str] = None,
+        context_passed: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Emit agent handoff event for tracking agent-to-agent flow."""
+        self._emit(ContextEvent(
+            event_type=ContextEventType.AGENT_HANDOFF,
+            timestamp=time.time(),
+            session_id=self._session_id,
+            agent_name=from_agent,
+            data={
+                "from_agent": from_agent,
+                "to_agent": to_agent,
+                "reason": reason,
+                "context_passed": context_passed or {},
+            },
+        ))
+    
+    def flush(self) -> None:
+        """Flush the sink."""
+        self._sink.flush()
+    
+    def close(self) -> None:
+        """Close the sink."""
+        self._sink.close()
+    
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self._session_id
+    
+    @property
+    def enabled(self) -> bool:
+        """Check if emitter is enabled."""
+        return self._enabled
+    
+    @enabled.setter
+    def enabled(self, value: bool) -> None:
+        """Set enabled state."""
+        self._enabled = value

src/praisonai-agents/praisonaiagents/utils/__init__.py

@@ -0,0 +1,49 @@
+"""
+Utils Module for PraisonAI Agents.
+
+Provides shared utilities used across the SDK:
+- Variable providers: Dynamic variable substitution ({{today}}, {{now}}, etc.)
+- Common helpers: Shared functions for Agent, Workflow, and other modules
+
+All imports are lazy-loaded for zero performance impact.
+"""
+
+from .variables import (
+    # Protocol
+    DynamicVariableProvider,
+    # Built-in providers
+    NowProvider,
+    TodayProvider,
+    DateProvider,
+    TimestampProvider,
+    UUIDProvider,
+    YearProvider,
+    MonthProvider,
+    # Registry
+    VariableProviderRegistry,
+    get_provider_registry,
+    register_variable_provider,
+    resolve_dynamic_variable,
+    # Shared utility
+    substitute_variables,
+)
+
+__all__ = [
+    # Protocol
+    "DynamicVariableProvider",
+    # Built-in providers
+    "NowProvider",
+    "TodayProvider",
+    "DateProvider",
+    "TimestampProvider",
+    "UUIDProvider",
+    "YearProvider",
+    "MonthProvider",
+    # Registry
+    "VariableProviderRegistry",
+    "get_provider_registry",
+    "register_variable_provider",
+    "resolve_dynamic_variable",
+    # Shared utility
+    "substitute_variables",
+]

src/praisonai-agents/praisonaiagents/utils/variables.py

@@ -0,0 +1,283 @@
+"""
+Dynamic Variable Providers for PraisonAI Workflows.
+
+Provides protocol-driven dynamic variable substitution for {{variable}} placeholders
+that resolve at runtime rather than requiring static values.
+
+Built-in providers:
+- now: Current datetime (ISO format)
+- today: Current date (human-readable)
+- date: Current date (YYYY-MM-DD)
+- uuid: Random UUID
+- timestamp: Unix timestamp
+
+Usage in YAML:
+    variables:
+      current_date: "{{today}}"
+      topic: "Latest AI news for {{today}}"
+
+Usage in Python:
+    from praisonaiagents.workflows.variable_providers import get_provider_registry
+    
+    registry = get_provider_registry()
+    registry.register("custom", MyCustomProvider())
+"""
+
+import threading
+import uuid as uuid_module
+from abc import ABC, abstractmethod
+from datetime import date, datetime
+from typing import Any, Callable, Dict, Optional, Union
+
+
+class DynamicVariableProvider(ABC):
+    """
+    Protocol for dynamic variable providers.
+    
+    Implement this class to create custom dynamic variables
+    that resolve at runtime.
+    
+    Example:
+        class CounterProvider(DynamicVariableProvider):
+            def __init__(self):
+                self._count = 0
+            
+            def get_value(self) -> int:
+                self._count += 1
+                return self._count
+    """
+    
+    @abstractmethod
+    def get_value(self) -> Any:
+        """Get the current value of this dynamic variable."""
+        ...
+
+
+# =============================================================================
+# Built-in Providers
+# =============================================================================
+
+class NowProvider(DynamicVariableProvider):
+    """Returns current datetime in ISO format."""
+    
+    def get_value(self) -> str:
+        return datetime.now().isoformat()
+
+
+class TodayProvider(DynamicVariableProvider):
+    """Returns current date in human-readable format (e.g., 'January 19, 2026')."""
+    
+    def get_value(self) -> str:
+        return date.today().strftime("%B %d, %Y")
+
+
+class DateProvider(DynamicVariableProvider):
+    """Returns current date in YYYY-MM-DD format."""
+    
+    def get_value(self) -> str:
+        return date.today().strftime("%Y-%m-%d")
+
+
+class TimestampProvider(DynamicVariableProvider):
+    """Returns current Unix timestamp."""
+    
+    def get_value(self) -> int:
+        return int(datetime.now().timestamp())
+
+
+class UUIDProvider(DynamicVariableProvider):
+    """Returns a new random UUID."""
+    
+    def get_value(self) -> str:
+        return str(uuid_module.uuid4())
+
+
+class YearProvider(DynamicVariableProvider):
+    """Returns current year."""
+    
+    def get_value(self) -> int:
+        return date.today().year
+
+
+class MonthProvider(DynamicVariableProvider):
+    """Returns current month name."""
+    
+    def get_value(self) -> str:
+        return date.today().strftime("%B")
+
+
+# =============================================================================
+# Provider Registry
+# =============================================================================
+
+class VariableProviderRegistry:
+    """
+    Central registry for dynamic variable providers.
+    
+    Thread-safe for multi-agent scenarios.
+    
+    Usage:
+        registry = get_provider_registry()
+        value = registry.resolve("today")  # "January 19, 2026"
+    """
+    
+    def __init__(self):
+        self._providers: Dict[str, DynamicVariableProvider] = {}
+        self._functions: Dict[str, Callable[[], Any]] = {}
+        self._lock = threading.RLock()
+        self._register_builtins()
+    
+    def _register_builtins(self) -> None:
+        """Register built-in providers."""
+        self._providers["now"] = NowProvider()
+        self._providers["today"] = TodayProvider()
+        self._providers["date"] = DateProvider()
+        self._providers["timestamp"] = TimestampProvider()
+        self._providers["uuid"] = UUIDProvider()
+        self._providers["year"] = YearProvider()
+        self._providers["month"] = MonthProvider()
+    
+    def register(
+        self,
+        name: str,
+        provider: Union[DynamicVariableProvider, Callable[[], Any]],
+        overwrite: bool = False
+    ) -> None:
+        """
+        Register a dynamic variable provider.
+        
+        Args:
+            name: Variable name (used as {{name}})
+            provider: DynamicVariableProvider instance or callable
+            overwrite: If True, overwrite existing provider
+        
+        Raises:
+            ValueError: If provider exists and overwrite=False
+        """
+        with self._lock:
+            if name in self._providers or name in self._functions:
+                if not overwrite:
+                    raise ValueError(f"Provider '{name}' already registered")
+            
+            if isinstance(provider, DynamicVariableProvider):
+                self._providers[name] = provider
+            elif callable(provider):
+                self._functions[name] = provider
+            else:
+                raise TypeError(f"Expected DynamicVariableProvider or callable, got {type(provider)}")
+    
+    def unregister(self, name: str) -> bool:
+        """Remove a provider."""
+        with self._lock:
+            if name in self._providers:
+                del self._providers[name]
+                return True
+            if name in self._functions:
+                del self._functions[name]
+                return True
+            return False
+    
+    def resolve(self, name: str) -> Optional[Any]:
+        """
+        Resolve a dynamic variable by name.
+        
+        Args:
+            name: Variable name
+            
+        Returns:
+            Resolved value, or None if not found
+        """
+        with self._lock:
+            if name in self._providers:
+                return self._providers[name].get_value()
+            if name in self._functions:
+                return self._functions[name]()
+            return None
+    
+    def has(self, name: str) -> bool:
+        """Check if a provider exists."""
+        with self._lock:
+            return name in self._providers or name in self._functions
+    
+    def list_providers(self) -> list:
+        """List all registered provider names."""
+        with self._lock:
+            return list(self._providers.keys()) + list(self._functions.keys())
+    
+    def __contains__(self, name: str) -> bool:
+        return self.has(name)
+
+
+# =============================================================================
+# Global Registry
+# =============================================================================
+
+_global_registry: Optional[VariableProviderRegistry] = None
+
+
+def get_provider_registry() -> VariableProviderRegistry:
+    """Get the global variable provider registry."""
+    global _global_registry
+    if _global_registry is None:
+        _global_registry = VariableProviderRegistry()
+    return _global_registry
+
+
+def register_variable_provider(
+    name: str,
+    provider: Union[DynamicVariableProvider, Callable[[], Any]]
+) -> None:
+    """Convenience function to register a provider with the global registry."""
+    get_provider_registry().register(name, provider)
+
+
+def resolve_dynamic_variable(name: str) -> Optional[Any]:
+    """Convenience function to resolve a dynamic variable."""
+    return get_provider_registry().resolve(name)
+
+
+# =============================================================================
+# Shared Variable Substitution Utility (DRY)
+# =============================================================================
+
+def substitute_variables(text: str, variables: Dict[str, Any]) -> str:
+    """
+    Substitute {{variable}} placeholders in text.
+    
+    This is the canonical, DRY implementation used by all workflow modules.
+    
+    Resolution order:
+    1. Dynamic variable providers ({{now}}, {{today}}, {{uuid}}, etc.)
+    2. Static variables from the provided variables dict
+    3. Keep original placeholder if not found
+    
+    Args:
+        text: Text containing {{variable}} placeholders
+        variables: Static variables dict
+        
+    Returns:
+        Text with variables substituted
+        
+    Example:
+        >>> substitute_variables("Hello {{today}} - {{name}}", {"name": "World"})
+        "Hello January 19, 2026 - World"
+    """
+    import re
+    
+    provider_registry = get_provider_registry()
+    
+    def replace(match):
+        var_name = match.group(1).strip()
+        
+        # First check dynamic providers
+        if provider_registry.has(var_name):
+            return str(provider_registry.resolve(var_name))
+        
+        # Then check static variables
+        if var_name in variables:
+            return str(variables[var_name])
+        
+        # Keep original if not found
+        return match.group(0)
+    
+    return re.sub(r'\{\{(\w+)\}\}', replace, text)

src/praisonai-agents/praisonaiagents/workflows/__init__.py

@@ -18,12 +18,14 @@ from .workflows import (
     Parallel,
     Loop,
     Repeat,
+    Include,
     
     # Convenience functions
     route,
     parallel,
     loop,
     repeat,
+    include,
     
     # Backward compatibility aliases
     StepInput,
@@ -55,8 +57,6 @@ from .workflow_configs import (
     resolve_step_routing_config,
 )
 
-from .yaml_parser import YAMLWorkflowParser
-
 __all__ = [
     # Core
     "Workflow",
@@ -74,10 +74,12 @@ __all__ = [
     "Parallel",
     "Loop",
     "Repeat",
+    "Include",
     "route",
     "parallel",
     "loop",
     "repeat",
+    "include",
     
     # Workflow Config Classes
     "WorkflowOutputConfig",
@@ -109,3 +111,20 @@ __all__ = [
     "StepInput",
     "StepOutput",
 ]
+
+
+# Lazy imports for heavy modules (YAMLWorkflowParser imports Agent which triggers pydantic/rich)
+_LAZY_IMPORTS = {
+    "YAMLWorkflowParser": "yaml_parser",
+}
+
+
+def __getattr__(name: str):
+    """Lazy import mechanism for heavy modules."""
+    if name in _LAZY_IMPORTS:
+        module_name = _LAZY_IMPORTS[name]
+        import importlib
+        module = importlib.import_module(f".{module_name}", __name__)
+        return getattr(module, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+

src/praisonai-agents/praisonaiagents/workflows/workflow_configs.py

@@ -48,25 +48,10 @@ class WorkflowOutputPreset(str, Enum):
     DEBUG = "debug"
 
 
-@dataclass
-class WorkflowOutputConfig:
-    """
-    Configuration for workflow output behavior.
-    
-    Consolidates: verbose, stream
-    
-    Usage:
-        # Preset (string)
-        Workflow(steps=[...], output="verbose")
-        
-        # Config
-        Workflow(steps=[...], output=WorkflowOutputConfig(verbose=True, stream=True))
-    """
-    verbose: bool = False
-    stream: bool = True
-    
-    def to_dict(self) -> Dict[str, Any]:
-        return {"verbose": self.verbose, "stream": self.stream}
+# WorkflowOutputConfig is now an alias to OutputConfig for DRY approach
+# Workflows use the same output configuration as Agent
+# Import OutputConfig and alias it for backward compatibility
+from ..config.feature_configs import OutputConfig as WorkflowOutputConfig
 
 
 @dataclass
@@ -319,23 +304,26 @@ def resolve_param(
     )
 
 
-def resolve_output_config(value: Any) -> WorkflowOutputConfig:
-    """Resolve output parameter to WorkflowOutputConfig using canonical resolver.
+def resolve_output_config(value: Any):
+    """Resolve output parameter using canonical OUTPUT_PRESETS (DRY approach).
     
-    Supports: None, str preset, list [preset, overrides], Config, dict
+    Supports: None, str preset, list [preset, overrides], OutputConfig, dict
+    
+    Uses the same OUTPUT_PRESETS as Agent for consistency.
     """
     from ..config.param_resolver import resolve, ArrayMode
-    from ..config.presets import WORKFLOW_OUTPUT_PRESETS
+    from ..config.presets import OUTPUT_PRESETS
+    from ..config.feature_configs import OutputConfig
     
     result = resolve(
         value=value,
         param_name="output",
-        config_class=WorkflowOutputConfig,
-        presets=WORKFLOW_OUTPUT_PRESETS,
+        config_class=OutputConfig,
+        presets=OUTPUT_PRESETS,
         array_mode=ArrayMode.PRESET_OVERRIDE,
-        default=WorkflowOutputConfig(),
+        default=OutputConfig(),  # Default is silent mode
     )
-    return result if result else WorkflowOutputConfig(verbose=False, stream=False)
+    return result if result else OutputConfig()
 
 
 def resolve_planning_config(value: Any) -> Optional[WorkflowPlanningConfig]:

src/praisonai-agents/praisonaiagents/workflows/yaml_parser.py

@@ -14,7 +14,7 @@ from typing import Any, Callable, Dict, List, Optional, Union
 from pathlib import Path
 
 from ..agent.agent import Agent
-from .workflows import Workflow, WorkflowStep, route, parallel, loop, repeat
+from .workflows import Workflow, WorkflowStep, route, parallel, loop, repeat, Include, include
 
 
 class YAMLWorkflowParser:
@@ -72,12 +72,13 @@ class YAMLWorkflowParser:
         self._agents: Dict[str, Agent] = {}
         self._callbacks: Dict[str, Callable] = {}
     
-    def parse_file(self, file_path: Union[str, Path]) -> Workflow:
+    def parse_file(self, file_path: Union[str, Path], extra_vars: Optional[Dict[str, Any]] = None) -> Workflow:
         """
         Parse a YAML workflow file.
         
         Args:
             file_path: Path to the YAML file
+            extra_vars: Optional dict of variables to override YAML defaults (e.g., from CLI --var)
             
         Returns:
             Workflow object ready for execution
@@ -89,14 +90,15 @@ class YAMLWorkflowParser:
         with open(file_path, 'r') as f:
             yaml_content = f.read()
         
-        return self.parse_string(yaml_content)
+        return self.parse_string(yaml_content, extra_vars=extra_vars)
     
-    def parse_string(self, yaml_content: str) -> Workflow:
+    def parse_string(self, yaml_content: str, extra_vars: Optional[Dict[str, Any]] = None) -> Workflow:
         """
         Parse a YAML workflow string.
         
         Args:
             yaml_content: YAML content as string
+            extra_vars: Optional dict of variables to override YAML defaults (e.g., from CLI --var)
             
         Returns:
             Workflow object ready for execution
@@ -104,7 +106,7 @@ class YAMLWorkflowParser:
         data = yaml.safe_load(yaml_content)
         # Normalize to canonical format (accept both old and new field names)
         normalized_data = self._normalize_yaml_config(data)
-        return self._parse_workflow_data(normalized_data)
+        return self._parse_workflow_data(normalized_data, extra_vars=extra_vars)
     
     def _normalize_yaml_config(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
@@ -150,6 +152,19 @@ class YAMLWorkflowParser:
             if 'steps' not in normalized:
                 normalized['steps'] = self._extract_steps_from_roles(normalized['roles'])
         
+        # 3.5. Handle 'includes:' section - append include steps after role-based steps
+        # This allows agents.yaml files to include other recipes as final steps
+        if 'includes' in normalized:
+            if 'steps' not in normalized:
+                normalized['steps'] = []
+            for include_item in normalized['includes']:
+                if isinstance(include_item, str):
+                    # Simple include: just recipe name
+                    normalized['steps'].append({'include': include_item})
+                elif isinstance(include_item, dict):
+                    # Include with configuration
+                    normalized['steps'].append({'include': include_item})
+        
         # 4. Normalize agent fields: backstory -> instructions (for agents section)
         if 'agents' in normalized:
             for agent_id, agent_config in normalized['agents'].items():
@@ -209,12 +224,13 @@ class YAMLWorkflowParser:
                     steps.append(step)
         return steps
     
-    def _parse_workflow_data(self, data: Dict[str, Any]) -> Workflow:
+    def _parse_workflow_data(self, data: Dict[str, Any], extra_vars: Optional[Dict[str, Any]] = None) -> Workflow:
         """
         Parse workflow data dictionary into a Workflow object.
         
         Args:
             data: Parsed YAML data dictionary
+            extra_vars: Optional dict of variables to override YAML defaults (e.g., from CLI --var)
             
         Returns:
             Workflow object
@@ -236,9 +252,23 @@ class YAMLWorkflowParser:
         reasoning = workflow_config.get('reasoning', False)
         verbose = workflow_config.get('verbose', False)
         memory_config = workflow_config.get('memory_config')
+        output_mode = workflow_config.get('output')  # NEW: parse output mode from YAML
         
-        # Parse variables
+        # Parse variables (YAML defaults), then merge CLI overrides (extra_vars take precedence)
         variables = data.get('variables', {})
+        if extra_vars:
+            variables = {**variables, **extra_vars}  # CLI --var overrides YAML defaults
+
+        
+        # Auto-add 'topic' to variables for include recipe propagation
+        # This ensures included recipes receive the parent's topic via {{topic}}
+        topic_value = data.get('topic')
+        if topic_value and 'topic' not in variables:
+            # Substitute template variables (e.g., {{today}}) in topic if present
+            if "{{" in str(topic_value):
+                from praisonaiagents.utils.variables import substitute_variables
+                topic_value = substitute_variables(str(topic_value), variables)
+            variables['topic'] = topic_value
         
         # Parse agents - support both 'agents' and 'roles' keys for backward compatibility
         agents_data = data.get('agents', {})
@@ -262,12 +292,56 @@ class YAMLWorkflowParser:
             from .workflow_configs import WorkflowPlanningConfig
             planning_value = WorkflowPlanningConfig(enabled=bool(planning), llm=planning_llm, reasoning=reasoning)
         
+        # Determine output mode: explicit output > verbose flag > default
+        workflow_output = output_mode
+        if workflow_output is None and verbose:
+            workflow_output = "verbose"
+        
+        # Build memory config if specified - pass raw dict for flexibility
+        memory_value = memory_config if memory_config else None
+        
+        # Parse context management config (CRITICAL for token overflow prevention)
+        # Supports: context: true OR context: {auto_compact: true, strategy: smart, ...}
+        # Also supports YAML-friendly aliases: enabled -> auto_compact, max_tool_output_tokens -> tool_output_max
+        context_config = data.get('context')
+        if context_config is True:
+            # Simple enable: context: true
+            context_value = True
+        elif isinstance(context_config, dict):
+            # Detailed config: context: {auto_compact: true, ...}
+            # Apply field aliasing for YAML-friendly names
+            normalized_config = context_config.copy()
+            
+            # Alias: enabled -> auto_compact
+            if 'enabled' in normalized_config:
+                normalized_config['auto_compact'] = normalized_config.pop('enabled')
+            
+            # Alias: max_tool_output_tokens -> tool_output_max
+            if 'max_tool_output_tokens' in normalized_config:
+                normalized_config['tool_output_max'] = normalized_config.pop('max_tool_output_tokens')
+            
+            # Alias: threshold -> compact_threshold
+            if 'threshold' in normalized_config:
+                normalized_config['compact_threshold'] = normalized_config.pop('threshold')
+            
+            try:
+                from ..context.models import ContextConfig
+                context_value = ContextConfig(**normalized_config)
+            except Exception:
+                # Fallback: just enable with True if ContextConfig fails
+                context_value = True
+        else:
+            context_value = None
+        
         workflow = Workflow(
             name=name,
             steps=steps,
             variables=variables,
             planning=planning_value,
             default_llm=default_llm,
+            output=workflow_output,  # Pass output mode to Workflow
+            memory=memory_value,  # Pass memory config to Workflow
+            context=context_value,  # Pass context management config to Workflow
         )
         
         # Store additional attributes for feature parity with agents.yaml
@@ -278,7 +352,14 @@ class YAMLWorkflowParser:
         
         # Store workflow input (from 'input' or 'topic' field)
         # This is the default input passed to workflow.start() if no input is provided
-        workflow.default_input = data.get('input', data.get('topic', ''))
+        default_input = data.get('input', data.get('topic', ''))
+        
+        # Substitute dynamic variables ({{today}}, {{now}}, etc.) in default_input
+        if default_input and "{{" in default_input:
+            from praisonaiagents.utils.variables import substitute_variables
+            default_input = substitute_variables(default_input, {})
+        
+        workflow.default_input = default_input
         
         return workflow
     
@@ -500,6 +581,8 @@ class YAMLWorkflowParser:
             return self._parse_loop_step(step_data)
         elif 'repeat' in step_data:
             return self._parse_repeat_step(step_data)
+        elif 'include' in step_data:
+            return self._parse_include_step(step_data)
         elif 'agent' in step_data:
             return self._parse_agent_step(step_data)
         else:
@@ -582,10 +665,47 @@ class YAMLWorkflowParser:
             if async_execution is not None:
                 agent._yaml_async_execution = async_execution
             
+            # Handle output_variable (store output in named variable for loop iteration)
+            output_variable = step_data.get('output_variable')
+            if output_variable:
+                agent._yaml_output_variable = output_variable
+            
             return agent
         else:
             raise ValueError(f"Agent '{agent_id}' not defined in agents section")
     
+    def _parse_include_step(self, step_data: Dict) -> Include:
+        """
+        Parse an include pattern step.
+        
+        Args:
+            step_data: Step definition with 'include' key
+            
+        Returns:
+            Include pattern object
+            
+        YAML syntax:
+            # Simple include
+            - include: wordpress-publisher
+            
+            # Include with configuration
+            - include:
+                recipe: wordpress-publisher
+                input: "{{previous_output}}"
+        """
+        include_config = step_data['include']
+        
+        if isinstance(include_config, str):
+            # Simple form: include: recipe-name
+            recipe_name = include_config
+            input_template = None
+        else:
+            # Config form: include: {recipe: ..., input: ...}
+            recipe_name = include_config.get('recipe', '')
+            input_template = include_config.get('input')
+        
+        return Include(recipe=recipe_name, input=input_template)
+    
     def _parse_route_step(self, step_data: Dict) -> Dict:
         """
         Parse a route pattern step.
@@ -635,7 +755,7 @@ class YAMLWorkflowParser:
         
         return parallel(parallel_steps)
     
-    def _parse_loop_step(self, step_data: Dict) -> Dict:
+    def _parse_loop_step(self, step_data: Dict):
         """
         Parse a loop pattern step.
         
@@ -644,24 +764,144 @@ class YAMLWorkflowParser:
             
         Returns:
             Loop pattern object
+            
+        Supports multiple syntax forms:
+            # Simple form with agent at step level
+            - loop:
+                over: items
+              agent: processor
+              
+            # Nested step form
+            - loop:
+                over: items
+                parallel: true
+              step:
+                agent: processor
+                
+            # Full nested form
+            - loop:
+                over: items
+                parallel: true
+                max_workers: 4
+                step: processor
+            
+            # Multi-step form (NEW) - steps at step_data level
+            - loop:
+                over: items
+                parallel: true
+              steps:
+                - agent: researcher
+                - agent: writer
+            
+            # Multi-step form (NEW) - steps inside loop config
+            - loop:
+                over: items
+                parallel: true
+                steps:
+                  - agent: researcher
+                  - agent: writer
         """
         loop_config = step_data['loop']
-        agent_id = step_data.get('agent')
         
-        if agent_id and agent_id in self._agents:
-            agent = self._agents[agent_id]
+        # Get loop parameters
+        over = loop_config.get('over')
+        from_csv = loop_config.get('from_csv')
+        from_file = loop_config.get('from_file')
+        var_name = loop_config.get('var_name', 'item')
+        parallel_flag = loop_config.get('parallel', False)
+        max_workers_raw = loop_config.get('max_workers')
+        # Ensure max_workers is int (YAML may parse as string or contain template)
+        max_workers = None
+        if max_workers_raw is not None:
+            try:
+                max_workers = int(max_workers_raw)
+            except (ValueError, TypeError):
+                # If it's a template variable like {{max_workers}}, leave as None
+                max_workers = None
+        
+        # Check for multi-step form (NEW) - steps: at step_data level or inside loop config
+        nested_steps = step_data.get('steps') or loop_config.get('steps')
+        if nested_steps and isinstance(nested_steps, list) and len(nested_steps) > 0:
+            # Parse each nested step
+            parsed_steps = []
+            for nested_step_data in nested_steps:
+                parsed_step = self._parse_single_step(nested_step_data)
+                if parsed_step is not None:
+                    parsed_steps.append(parsed_step)
             
+            if parsed_steps:
+                # Extract output_variable from step level
+                output_variable = step_data.get('output_variable')
+                
+                return loop(
+                    steps=parsed_steps,
+                    over=over, 
+                    from_csv=from_csv, 
+                    from_file=from_file,
+                    var_name=var_name,
+                    parallel=parallel_flag, 
+                    max_workers=max_workers,
+                    output_variable=output_variable
+                )
+        
+        # Resolve the step/agent to execute (single step - backward compat)
+        agent_or_step = None
+        
+        # First check for agent at step level (original syntax)
+        agent_id = step_data.get('agent')
+        if agent_id and agent_id in self._agents:
+            agent_or_step = self._agents[agent_id]
             # Store action if present
             if 'action' in step_data:
-                agent._yaml_action = step_data['action']
-            
-            # Get loop parameters
-            over = loop_config.get('over')
-            from_csv = loop_config.get('from_csv')
-            
-            return loop(agent, over=over, from_csv=from_csv)
+                agent_or_step._yaml_action = step_data['action']
         
-        raise ValueError("Loop step requires an agent")
+        # Check for step: syntax (nested form)
+        if 'step' in step_data:
+            step_def = step_data['step']
+            if isinstance(step_def, dict) and 'agent' in step_def:
+                agent_id = step_def['agent']
+                if agent_id in self._agents:
+                    agent_or_step = self._agents[agent_id]
+                    if 'action' in step_def:
+                        agent_or_step._yaml_action = step_def['action']
+            elif isinstance(step_def, str):
+                # step: agent_name shorthand
+                if step_def in self._agents:
+                    agent_or_step = self._agents[step_def]
+        
+        # Check for step inside loop config
+        if 'step' in loop_config:
+            step_def = loop_config['step']
+            if isinstance(step_def, str) and step_def in self._agents:
+                agent_or_step = self._agents[step_def]
+        
+        # Check for include at step level (include inside loop)
+        if agent_or_step is None and 'include' in step_data:
+            include_value = step_data['include']
+            if isinstance(include_value, str):
+                agent_or_step = Include(recipe=include_value)
+            elif isinstance(include_value, dict):
+                agent_or_step = Include(
+                    recipe=include_value.get('recipe', ''),
+                    input=include_value.get('input')
+                )
+        
+        if agent_or_step is None:
+            raise ValueError("Loop step requires an agent, include, or steps")
+        
+        # Extract output_variable from step level (where user specifies it in YAML)
+        output_variable = step_data.get('output_variable')
+        
+        return loop(
+            step=agent_or_step, 
+            over=over, 
+            from_csv=from_csv, 
+            from_file=from_file,
+            var_name=var_name,
+            parallel=parallel_flag, 
+            max_workers=max_workers,
+            output_variable=output_variable
+        )
     
     def _parse_repeat_step(self, step_data: Dict) -> Dict:
         """

src/praisonai-agents/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "praisonaiagents"
-version = "0.12.8"
+version = "0.13.4"
 description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
 readme = "README.md"
 requires-python = ">=3.10"

src/praisonai-agents/tests/__init__.py

@@ -0,0 +1 @@
+"""PraisonAI Agents test suite."""

src/praisonai-agents/tests/context/test_async_file_ops.py

@@ -0,0 +1,65 @@
+"""
+Tests for async file operations.
+"""
+
+import pytest
+import tempfile
+import asyncio
+from pathlib import Path
+
+
+@pytest.fixture
+def temp_file():
+    """Create a temporary file with test content."""
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
+        f.write("Line 1\nLine 2\nLine 3\nLine 4\nLine 5\n")
+        temp_path = f.name
+    yield temp_path
+    Path(temp_path).unlink(missing_ok=True)
+
+
+class TestAsyncFileOps:
+    """Tests for async file operations."""
+    
+    def test_aiofiles_availability_check(self):
+        """Should check if aiofiles is available."""
+        from praisonaiagents.context.fast.async_file_ops import is_aiofiles_available
+        
+        result = is_aiofiles_available()
+        assert isinstance(result, bool)
+    
+    @pytest.mark.asyncio
+    async def test_async_read_file(self, temp_file):
+        """Should read file content async."""
+        from praisonaiagents.context.fast.async_file_ops import async_read_file
+        
+        content = await async_read_file(temp_file)
+        
+        assert "Line 1" in content
+        assert "Line 5" in content
+    
+    @pytest.mark.asyncio
+    async def test_async_read_lines(self, temp_file):
+        """Should read specific lines async."""
+        from praisonaiagents.context.fast.async_file_ops import async_read_lines
+        
+        lines = await async_read_lines(temp_file, start_line=2, end_line=4)
+        
+        assert len(lines) == 3
+        assert lines[0] == "Line 2"
+        assert lines[2] == "Line 4"
+    
+    @pytest.mark.asyncio
+    async def test_fallback_works_without_aiofiles(self, temp_file, monkeypatch):
+        """Should fallback gracefully if aiofiles not available."""
+        from praisonaiagents.context.fast import async_file_ops
+        
+        # Force aiofiles to be unavailable
+        monkeypatch.setattr(async_file_ops, "_AIOFILES_AVAILABLE", False)
+        
+        content = await async_file_ops.async_read_file(temp_file)
+        assert "Line 1" in content
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/context/test_compressor.py

@@ -0,0 +1,109 @@
+"""
+Tests for context compressor.
+"""
+
+import pytest
+
+
+class TestTruncateCompressor:
+    """Tests for TruncateCompressor."""
+    
+    def test_no_compression_needed(self):
+        """Should return text unchanged if under budget."""
+        from praisonaiagents.context.fast.compressor import TruncateCompressor
+        
+        compressor = TruncateCompressor()
+        text = "Short text"
+        
+        result = compressor.compress(text, max_tokens=100)
+        
+        assert result == text
+    
+    def test_truncation_preserves_start_end(self):
+        """Should preserve start and end lines."""
+        from praisonaiagents.context.fast.compressor import TruncateCompressor
+        
+        compressor = TruncateCompressor()
+        lines = [f"Line {i}" for i in range(100)]
+        text = '\n'.join(lines)
+        
+        result = compressor.compress(text, max_tokens=50)
+        
+        # Should contain start lines
+        assert "Line 0" in result
+        # Should contain truncation marker
+        assert "truncated" in result.lower()
+    
+    def test_token_estimation(self):
+        """Should estimate tokens correctly."""
+        from praisonaiagents.context.fast.compressor import TruncateCompressor
+        
+        compressor = TruncateCompressor()
+        text = "1234" * 100  # 400 chars
+        
+        tokens = compressor.estimate_tokens(text)
+        
+        assert tokens == 100  # 4 chars per token
+
+
+class TestSmartCompressor:
+    """Tests for SmartCompressor."""
+    
+    def test_preserves_definitions(self):
+        """Should preserve function/class definitions."""
+        from praisonaiagents.context.fast.compressor import SmartCompressor
+        
+        compressor = SmartCompressor()
+        text = '''def important_function():
+    pass
+
+some_variable = 1
+another_variable = 2
+
+class ImportantClass:
+    pass
+
+random_code = 3
+more_random = 4
+'''
+        
+        result = compressor.compress(text, max_tokens=20)
+        
+        # Should preserve definitions
+        assert "def important_function" in result or len(result) < len(text)
+
+
+class TestCompressorSelection:
+    """Tests for compressor selection."""
+    
+    def test_get_truncate_compressor(self):
+        """Should return truncate compressor."""
+        from praisonaiagents.context.fast.compressor import get_compressor
+        
+        compressor = get_compressor("truncate")
+        assert compressor is not None
+    
+    def test_get_smart_compressor(self):
+        """Should return smart compressor."""
+        from praisonaiagents.context.fast.compressor import get_compressor
+        
+        compressor = get_compressor("smart")
+        assert compressor is not None
+    
+    def test_invalid_compressor_raises(self):
+        """Should raise for invalid compressor."""
+        from praisonaiagents.context.fast.compressor import get_compressor
+        
+        with pytest.raises(ValueError):
+            get_compressor("invalid")
+    
+    def test_llmlingua_availability(self):
+        """Should check llmlingua availability."""
+        from praisonaiagents.context.fast.compressor import is_llmlingua_available
+        
+        result = is_llmlingua_available()
+        assert isinstance(result, bool)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/context/test_index_manager.py

@@ -0,0 +1,119 @@
+"""
+Tests for index manager.
+"""
+
+import pytest
+import tempfile
+import os
+import time
+from pathlib import Path
+
+
+@pytest.fixture
+def temp_workspace():
+    """Create a temporary workspace."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create test files
+        (Path(tmpdir) / "file1.py").write_text("print('hello')")
+        (Path(tmpdir) / "file2.py").write_text("print('world')")
+        yield tmpdir
+
+
+class TestFileIndex:
+    """Tests for FileIndex."""
+    
+    def test_index_creation(self, temp_workspace):
+        """Should create empty index."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        index = FileIndex(workspace_path=temp_workspace)
+        
+        assert len(index) == 0
+        assert index.workspace_path == temp_workspace
+    
+    def test_index_update(self, temp_workspace):
+        """Should update index for file."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        index = FileIndex(workspace_path=temp_workspace)
+        filepath = os.path.join(temp_workspace, "file1.py")
+        
+        index.update(filepath)
+        
+        assert len(index) == 1
+        assert not index.needs_rescan(filepath)
+    
+    def test_needs_rescan_new_file(self, temp_workspace):
+        """Should need rescan for new file."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        index = FileIndex(workspace_path=temp_workspace)
+        filepath = os.path.join(temp_workspace, "file1.py")
+        
+        assert index.needs_rescan(filepath) is True
+    
+    def test_needs_rescan_modified_file(self, temp_workspace):
+        """Should need rescan for modified file."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        index = FileIndex(workspace_path=temp_workspace)
+        filepath = os.path.join(temp_workspace, "file1.py")
+        
+        index.update(filepath)
+        assert index.needs_rescan(filepath) is False
+        
+        # Modify file
+        time.sleep(0.01)  # Ensure mtime changes
+        Path(filepath).write_text("print('modified')")
+        
+        assert index.needs_rescan(filepath) is True
+    
+    def test_save_and_load(self, temp_workspace):
+        """Should save and load index."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        # Create and save
+        index = FileIndex(workspace_path=temp_workspace)
+        index.update(os.path.join(temp_workspace, "file1.py"))
+        assert index.save() is True
+        
+        # Load
+        loaded = FileIndex.load(temp_workspace)
+        assert loaded is not None
+        assert len(loaded) == 1
+    
+    def test_load_or_create_existing(self, temp_workspace):
+        """Should load existing index."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        # Create and save
+        index = FileIndex(workspace_path=temp_workspace)
+        index.update(os.path.join(temp_workspace, "file1.py"))
+        index.save()
+        
+        # Load or create
+        loaded = FileIndex.load_or_create(temp_workspace)
+        assert len(loaded) == 1
+    
+    def test_load_or_create_new(self, temp_workspace):
+        """Should create new index if none exists."""
+        from praisonaiagents.context.fast.index_manager import FileIndex
+        
+        # No existing index
+        index = FileIndex.load_or_create(temp_workspace)
+        assert len(index) == 0
+
+
+class TestIndexWatcher:
+    """Tests for IndexWatcher."""
+    
+    def test_watcher_availability(self):
+        """Should check watchfiles availability."""
+        from praisonaiagents.context.fast.index_manager import is_watchfiles_available
+        
+        result = is_watchfiles_available()
+        assert isinstance(result, bool)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/context/test_search_backends.py

@@ -0,0 +1,193 @@
+"""
+Tests for FastContext search backends.
+
+Tests the SearchBackend protocol and implementations:
+- PythonSearchBackend (always available)
+- RipgrepBackend (optional, with fallback)
+"""
+
+import pytest
+import tempfile
+import os
+from pathlib import Path
+
+
+@pytest.fixture
+def temp_workspace():
+    """Create a temporary workspace with test files."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create test files
+        src_dir = Path(tmpdir) / "src"
+        src_dir.mkdir()
+        
+        (src_dir / "main.py").write_text("""
+def main():
+    print("Hello, World!")
+    authenticate_user()
+
+def authenticate_user():
+    # Authentication logic here
+    pass
+
+class UserService:
+    def login(self, username, password):
+        return True
+""")
+        
+        (src_dir / "utils.py").write_text("""
+import os
+
+def get_config():
+    return {"debug": True}
+
+def validate_input(data):
+    if not data:
+        raise ValueError("Empty data")
+    return True
+""")
+        
+        yield tmpdir
+
+
+class TestSearchBackendProtocol:
+    """Tests for SearchBackend protocol."""
+    
+    def test_python_backend_always_available(self):
+        """Python backend should always be available."""
+        from praisonaiagents.context.fast.search_backends import PythonSearchBackend
+        
+        backend = PythonSearchBackend()
+        assert backend.is_available() is True
+    
+    def test_python_backend_grep(self, temp_workspace):
+        """Python backend grep should find matches."""
+        from praisonaiagents.context.fast.search_backends import PythonSearchBackend
+        
+        backend = PythonSearchBackend()
+        results = backend.grep(temp_workspace, "authenticate")
+        
+        assert len(results) > 0
+        assert any("main.py" in r.get("path", "") for r in results)
+    
+    def test_python_backend_glob(self, temp_workspace):
+        """Python backend glob should find files."""
+        from praisonaiagents.context.fast.search_backends import PythonSearchBackend
+        
+        backend = PythonSearchBackend()
+        results = backend.glob(temp_workspace, "**/*.py")
+        
+        assert len(results) >= 2
+
+
+class TestRipgrepBackend:
+    """Tests for RipgrepBackend."""
+    
+    def test_ripgrep_availability_check(self):
+        """RipgrepBackend should check if rg is installed."""
+        from praisonaiagents.context.fast.search_backends import RipgrepBackend
+        
+        backend = RipgrepBackend()
+        # Should return bool without error
+        result = backend.is_available()
+        assert isinstance(result, bool)
+    
+    def test_ripgrep_fallback_to_python(self, temp_workspace, monkeypatch):
+        """When ripgrep not available, should fallback to Python backend."""
+        from praisonaiagents.context.fast.search_backends import RipgrepBackend, get_search_backend
+        
+        # Force ripgrep to be unavailable
+        monkeypatch.setattr("shutil.which", lambda x: None)
+        
+        backend = RipgrepBackend()
+        assert backend.is_available() is False
+        
+        # Auto backend should fallback to python
+        auto_backend = get_search_backend("auto")
+        assert auto_backend.is_available() is True
+
+
+class TestBackendSelection:
+    """Tests for backend selection logic."""
+    
+    def test_get_python_backend(self):
+        """Should return Python backend when requested."""
+        from praisonaiagents.context.fast.search_backends import get_search_backend
+        
+        backend = get_search_backend("python")
+        assert backend.is_available() is True
+    
+    def test_get_auto_backend(self):
+        """Auto should return best available backend."""
+        from praisonaiagents.context.fast.search_backends import get_search_backend
+        
+        backend = get_search_backend("auto")
+        # Should always return a working backend
+        assert backend.is_available() is True
+    
+    def test_invalid_backend_raises(self):
+        """Invalid backend name should raise ValueError."""
+        from praisonaiagents.context.fast.search_backends import get_search_backend
+        
+        with pytest.raises(ValueError):
+            get_search_backend("invalid_backend")
+
+
+class TestSmartBackend:
+    """Tests for SmartBackend auto-selection based on codebase size."""
+    
+    def test_smart_backend_available(self):
+        """SmartBackend should always be available."""
+        from praisonaiagents.context.fast.search_backends import SmartBackend
+        
+        backend = SmartBackend()
+        assert backend.is_available() is True
+    
+    def test_small_workspace_uses_python(self, temp_workspace):
+        """Small workspace (<500 files) should use Python backend."""
+        from praisonaiagents.context.fast.search_backends import SmartBackend
+        
+        backend = SmartBackend(workspace_path=temp_workspace)
+        # Force check by doing a search
+        results = backend.grep(temp_workspace, "def")
+        
+        # Should have results and used Python (fast)
+        assert len(results) > 0
+        # Check internal state - should not have selected ripgrep
+        assert backend._use_ripgrep is False
+    
+    def test_count_files_fast(self, temp_workspace):
+        """Helper should count files quickly with limit."""
+        from praisonaiagents.context.fast.search_backends import _count_files_fast
+        
+        count = _count_files_fast(temp_workspace, limit=100)
+        assert count >= 2  # We created at least 2 files
+        assert count < 100  # Should be below limit
+    
+    def test_count_files_fast_skips_gitignore_dirs(self):
+        """Should skip .git, node_modules, etc."""
+        from praisonaiagents.context.fast.search_backends import _count_files_fast
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create files in regular dir
+            Path(tmpdir, "src").mkdir()
+            Path(tmpdir, "src", "file.py").write_text("code")
+            
+            # Create files in ignored dir
+            Path(tmpdir, "node_modules").mkdir()
+            Path(tmpdir, "node_modules", "pkg").mkdir()
+            for i in range(10):
+                Path(tmpdir, "node_modules", "pkg", f"file{i}.js").write_text("code")
+            
+            count = _count_files_fast(tmpdir)
+            # Should only count src/file.py, not node_modules files
+            assert count == 1
+    
+    def test_threshold_constant(self):
+        """Threshold should be 500 files based on profiling."""
+        from praisonaiagents.context.fast.search_backends import RIPGREP_FILE_THRESHOLD
+        
+        assert RIPGREP_FILE_THRESHOLD == 500
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/test_context_agent.py

@@ -96,7 +96,6 @@ def test_agent_inheritance():
         print(f"❌ Inheritance test failed: {e}")
         assert False, str(e)
 
-@pytest.mark.skip(reason="ContextAgent methods not yet implemented")
 def test_context_engineering_methods():
     """Test Context Engineering specific methods."""
     print("\\n🧪 Testing Context Engineering Methods...")
@@ -106,45 +105,104 @@ def test_context_engineering_methods():
         
         context_agent = create_context_agent()
         
-        # Test that all Context Engineering methods exist
+        # Test that all Context Engineering methods exist (actual method names)
         required_methods = [
-            'analyze_codebase_patterns',
-            'generate_context_document',
-            'create_validation_loop',
-            'enhance_prompt_with_context',
-            'generate_prp',
-            'extract_documentation_patterns',
+            'analyze_codebase_with_gitingest',
+            'perform_ast_analysis',
+            'extract_implementation_patterns',
             'analyze_test_patterns',
-            'create_implementation_blueprint'
+            'generate_comprehensive_prp',
+            'create_validation_framework',
+            'compile_context_documentation',
+            'analyze_integration_points',
+            'build_implementation_blueprint',
+            'create_quality_gates',
+            'generate_feature_prp',
+            'start',
         ]
         
+        missing = []
         for method_name in required_methods:
-            assert hasattr(context_agent, method_name), f"Missing method: {method_name}"
-            assert callable(getattr(context_agent, method_name)), f"Method {method_name} is not callable"
+            if not hasattr(context_agent, method_name):
+                missing.append(method_name)
+            elif not callable(getattr(context_agent, method_name)):
+                missing.append(f"{method_name} (not callable)")
         
-        print(f"✅ All {len(required_methods)} Context Engineering methods are available")
+        if missing:
+            print(f"⚠️ Missing methods: {missing}")
         
-        # Test async methods
-        async_methods = [
-            'aanalyze_codebase_patterns',
-            'agenerate_context_document',
-            'acreate_validation_loop',
-            'aenhance_prompt_with_context',
-            'agenerate_prp'
-        ]
+        available = len(required_methods) - len(missing)
+        print(f"✅ {available}/{len(required_methods)} Context Engineering methods are available")
         
-        for method_name in async_methods:
-            assert hasattr(context_agent, method_name), f"Missing async method: {method_name}"
-            assert callable(getattr(context_agent, method_name)), f"Async method {method_name} is not callable"
-        
-        print(f"✅ All {len(async_methods)} async methods are available")
-        
-        assert True
+        assert available >= 10, f"Too many missing methods: {missing}"
         
     except Exception as e:
         print(f"❌ Method availability test failed: {e}")
         assert False, str(e)
 
+def test_protocol_compliance():
+    """Test that ContextAgent implements ContextEngineerProtocol."""
+    print("\\n🧪 Testing Protocol Compliance...")
+    
+    try:
+        from praisonaiagents import ContextAgent, create_context_agent
+        from praisonaiagents.agent import ContextEngineerProtocol
+        
+        context_agent = create_context_agent(llm="gpt-4o-mini")
+        
+        # Test protocol compliance
+        assert isinstance(context_agent, ContextEngineerProtocol), \
+            "ContextAgent should implement ContextEngineerProtocol"
+        print("✅ ContextAgent implements ContextEngineerProtocol")
+        
+        # Test protocol-compatible aliases exist
+        protocol_methods = [
+            'analyze_codebase',
+            'generate_prp',
+            'create_implementation_blueprint',
+            'aanalyze_codebase',
+            'agenerate_prp',
+        ]
+        
+        for method in protocol_methods:
+            assert hasattr(context_agent, method), f"Missing protocol method: {method}"
+            assert callable(getattr(context_agent, method)), f"{method} not callable"
+        
+        print(f"✅ All {len(protocol_methods)} protocol methods available")
+        
+    except Exception as e:
+        print(f"❌ Protocol compliance test failed: {e}")
+        assert False, str(e)
+
+@pytest.mark.asyncio
+async def test_async_methods():
+    """Test async methods work correctly."""
+    print("\\n🧪 Testing Async Methods...")
+    
+    try:
+        from praisonaiagents import create_context_agent
+        import asyncio
+        
+        context_agent = create_context_agent(llm="gpt-4o-mini")
+        
+        # Test that async methods exist and are coroutines
+        async_methods = [
+            'aanalyze_codebase',
+            'agenerate_prp',
+            'acreate_implementation_blueprint',
+        ]
+        
+        for method_name in async_methods:
+            assert hasattr(context_agent, method_name), f"Missing: {method_name}"
+            method = getattr(context_agent, method_name)
+            assert asyncio.iscoroutinefunction(method), f"{method_name} not async"
+        
+        print(f"✅ All {len(async_methods)} async methods are coroutines")
+        
+    except Exception as e:
+        print(f"❌ Async methods test failed: {e}")
+        assert False, str(e)
+
 @pytest.mark.skip(reason="ContextAgent methods not yet implemented")
 def test_basic_functionality():
     """Test basic functionality of ContextAgent methods."""

src/praisonai-agents/tests/test_include_in_loop.py

@@ -0,0 +1,86 @@
+"""
+Tests for include inside loop feature.
+"""
+
+import pytest
+
+
+class TestIncludeInLoop:
+    """Tests for include steps inside loops."""
+    
+    def test_include_step_handled_in_single_step_internal(self):
+        """Test that Include steps are handled in _execute_single_step_internal."""
+        from praisonaiagents.workflows import Workflow, Include
+        
+        workflow = Workflow(steps=[])
+        
+        # Create an Include step
+        include_step = Include(recipe="test-recipe")
+        
+        # Verify Include is recognized
+        assert isinstance(include_step, Include)
+        assert include_step.recipe == "test-recipe"
+    
+    def test_loop_step_can_contain_include(self):
+        """Test that Loop can have Include as its step."""
+        from praisonaiagents.workflows import Loop, Include
+        
+        include_step = Include(recipe="wordpress-publisher")
+        loop_step = Loop(
+            step=include_step,
+            over="items",
+            parallel=True
+        )
+        
+        assert loop_step.step is include_step
+        assert isinstance(loop_step.step, Include)
+        assert loop_step.parallel is True
+
+
+class TestIncludeInLoopYAMLParsing:
+    """Tests for parsing include inside loop from YAML."""
+    
+    def test_parse_loop_with_include_step(self):
+        """Test parsing a loop that contains an include step."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Loop, Include
+        
+        yaml_content = '''
+agents:
+  writer:
+    role: Writer
+    goal: Write content
+    
+steps:
+  - name: write_posts
+    agent: writer
+    action: "Write about {{input}}"
+    output_variable: posts
+    
+  - loop:
+      over: posts
+      parallel: true
+    include: wordpress-publisher
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        assert len(workflow.steps) >= 2
+        
+        # Find the loop step
+        loop_found = False
+        for step in workflow.steps:
+            if isinstance(step, Loop):
+                loop_found = True
+                assert step.over == "posts"
+                assert step.parallel is True
+                # The step inside the loop should be an Include
+                assert isinstance(step.step, Include)
+                assert step.step.recipe == "wordpress-publisher"
+                break
+        
+        assert loop_found, "Should have a loop step with include"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/test_include_pattern.py

@@ -0,0 +1,267 @@
+"""
+TDD tests for Include pattern in YAML workflows.
+
+Tests for recipe composition:
+- Parsing include: recipe_name
+- Executing included recipes
+- Cycle detection
+- Input/output passing
+"""
+
+import pytest
+from pathlib import Path
+import tempfile
+import os
+
+
+class TestIncludePatternParsing:
+    """Tests for parsing include steps from YAML."""
+    
+    def test_include_pattern_import(self):
+        """Test that Include class can be imported."""
+        from praisonaiagents.workflows import Include
+        assert Include is not None
+    
+    def test_include_function_import(self):
+        """Test that include() convenience function can be imported."""
+        from praisonaiagents.workflows import include
+        assert include is not None
+    
+    def test_include_class_creation(self):
+        """Test creating Include object directly."""
+        from praisonaiagents.workflows import Include
+        
+        inc = Include(recipe="wordpress-publisher")
+        assert inc.recipe == "wordpress-publisher"
+        assert inc.input is None
+    
+    def test_include_with_input(self):
+        """Test Include with custom input."""
+        from praisonaiagents.workflows import Include
+        
+        inc = Include(recipe="wordpress-publisher", input="{{previous_output}}")
+        assert inc.recipe == "wordpress-publisher"
+        assert inc.input == "{{previous_output}}"
+    
+    def test_parse_simple_include_from_yaml(self):
+        """Test parsing simple include: recipe_name from YAML."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Include
+        
+        yaml_content = '''
+name: Test Workflow
+steps:
+  - include: wordpress-publisher
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Should have one step that is an Include
+        assert len(workflow.steps) == 1
+        step = workflow.steps[0]
+        assert isinstance(step, Include)
+        assert step.recipe == "wordpress-publisher"
+    
+    def test_parse_include_with_config(self):
+        """Test parsing include with configuration."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Include
+        
+        yaml_content = '''
+name: Test Workflow
+steps:
+  - include:
+      recipe: wordpress-publisher
+      input: "Custom input"
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert len(workflow.steps) == 1
+        step = workflow.steps[0]
+        assert isinstance(step, Include)
+        assert step.recipe == "wordpress-publisher"
+        assert step.input == "Custom input"
+
+
+class TestIncludePatternExecution:
+    """Tests for executing include steps."""
+    
+    def test_include_resolves_recipe_path(self):
+        """Test that include step resolves recipe path."""
+        from praisonaiagents.workflows import Include
+        
+        # Create a mock recipe in temp dir
+        with tempfile.TemporaryDirectory() as tmpdir:
+            recipe_dir = Path(tmpdir) / "test-recipe"
+            recipe_dir.mkdir()
+            
+            agents_yaml = recipe_dir / "agents.yaml"
+            agents_yaml.write_text('''
+framework: praisonai
+topic: Test
+roles:
+  agent1:
+    role: Test Agent
+    goal: Test
+    backstory: Test
+    tasks:
+      task1:
+        description: Return the input
+        expected_output: The input
+''')
+            
+            inc = Include(recipe="test-recipe")
+            # Path resolution is done at execution time
+            assert inc.recipe == "test-recipe"
+    
+    def test_execute_include_method_exists(self):
+        """Test that _execute_include method exists on Workflow."""
+        from praisonaiagents.workflows import Workflow
+        
+        workflow = Workflow(name="test", steps=[])
+        assert hasattr(workflow, "_execute_include")
+        assert callable(getattr(workflow, "_execute_include"))
+    
+    def test_execute_include_with_handler_workflow(self):
+        """Test executing a workflow with Include step using handler function."""
+        from praisonaiagents.workflows import Workflow, Include, WorkflowStep, StepResult
+        import tempfile
+        from pathlib import Path
+        import os
+        
+        # Create a temp recipe that uses a simple handler
+        with tempfile.TemporaryDirectory() as tmpdir:
+            recipe_dir = Path(tmpdir) / "simple-recipe"
+            recipe_dir.mkdir()
+            
+            # Create a minimal workflow YAML with a handler
+            agents_yaml = recipe_dir / "agents.yaml"
+            agents_yaml.write_text('''
+name: Simple Handler Recipe
+description: A recipe with a simple handler
+steps:
+  - name: echo_step
+    action: "Echo: {{input}}"
+''')
+            
+            # Create Include pointing to the temp recipe path
+            include_step = Include(recipe=str(recipe_dir))
+            
+            # Create parent workflow
+            parent = Workflow(
+                name="Parent Workflow",
+                steps=[include_step]
+            )
+            
+            # The include should resolve the local path
+            assert include_step.recipe == str(recipe_dir)
+
+
+class TestCycleDetection:
+    """Tests for cycle detection in includes."""
+    
+    def test_cycle_detection_raises_error(self):
+        """Test that circular includes raise an error."""
+        from praisonaiagents.workflows import Include
+        
+        # Simulate cycle detection logic
+        visited = {"recipe-a", "recipe-b"}
+        
+        # Attempting to include recipe-a again should fail
+        inc = Include(recipe="recipe-a")
+        
+        # The cycle detection happens in _execute_include
+        # Here we just test the Include object can track visited
+        assert "recipe-a" in visited
+
+
+class TestIncludeConvenienceFunction:
+    """Tests for the include() convenience function."""
+    
+    def test_include_function_creates_include(self):
+        """Test include() creates Include object."""
+        from praisonaiagents.workflows import include, Include
+        
+        inc = include("wordpress-publisher")
+        assert isinstance(inc, Include)
+        assert inc.recipe == "wordpress-publisher"
+    
+    def test_include_function_with_input(self):
+        """Test include() with input parameter."""
+        from praisonaiagents.workflows import include, Include
+        
+        inc = include("wordpress-publisher", input="test input")
+        assert isinstance(inc, Include)
+        assert inc.input == "test input"
+
+
+class TestIncludesInRolesFormat:
+    """Tests for 'includes:' section in roles-format YAML files."""
+    
+    def test_includes_section_simple(self):
+        """Test parsing 'includes:' section with simple recipe names."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Include
+        
+        yaml_content = '''
+framework: praisonai
+topic: Test
+
+roles:
+  writer:
+    role: Writer
+    goal: Write content
+    backstory: You are a writer
+    tasks:
+      write:
+        description: Write about AI
+
+includes:
+  - wordpress-publisher
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Should have 2 steps: 1 from roles, 1 from includes
+        assert len(workflow.steps) == 2
+        
+        # Last step should be an Include
+        last_step = workflow.steps[-1]
+        assert isinstance(last_step, Include)
+        assert last_step.recipe == "wordpress-publisher"
+    
+    def test_includes_section_with_config(self):
+        """Test parsing 'includes:' section with configuration."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Include
+        
+        yaml_content = '''
+framework: praisonai
+topic: Test
+
+roles:
+  writer:
+    role: Writer
+    goal: Write
+    backstory: Writer
+    tasks:
+      write:
+        description: Write
+
+includes:
+  - recipe: wordpress-publisher
+    input: "{{previous_output}}"
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Should have 2 steps
+        assert len(workflow.steps) == 2
+        
+        # Last step should be an Include with input
+        last_step = workflow.steps[-1]
+        assert isinstance(last_step, Include)
+        assert last_step.recipe == "wordpress-publisher"
+        assert last_step.input == "{{previous_output}}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
+

src/praisonai-agents/tests/test_include_variables.py

@@ -0,0 +1,321 @@
+"""
+Test-driven development tests for recipe variable passing.
+
+Tests the propagation of 'topic' and other YAML root-level fields
+to included recipes via the variables dict.
+"""
+
+import pytest
+import tempfile
+import os
+from pathlib import Path
+
+
+class TestTopicVariablePropagation:
+    """Tests for topic variable propagation to variables dict."""
+    
+    def test_topic_field_added_to_variables(self):
+        """Test that 'topic' field is automatically added to workflow variables."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Topic Variables Test
+topic: "AI Code & Tools"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{topic}}"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # topic should be in variables dict
+        assert "topic" in workflow.variables
+        assert workflow.variables["topic"] == "AI Code & Tools"
+    
+    def test_topic_with_template_variable_substituted(self):
+        """Test that topic containing {{var}} is substituted before adding to variables."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Topic Substitution Test
+topic: "AI Code & Tools {{today}}"
+
+variables:
+  today: "January 2026"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{topic}}"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # topic should be substituted (either with variables section value or dynamic today)
+        assert "topic" in workflow.variables
+        # Should NOT contain raw {{today}} template
+        assert "{{today}}" not in workflow.variables["topic"]
+        # Should contain "AI Code & Tools" prefix
+        assert "AI Code & Tools" in workflow.variables["topic"]
+    
+    def test_explicit_topic_in_variables_not_overwritten(self):
+        """Test that explicit topic in variables section is NOT overwritten."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Explicit Topic Test
+topic: "From topic field"
+
+variables:
+  topic: "From variables section"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{topic}}"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # Explicit topic in variables should take precedence
+        assert workflow.variables["topic"] == "From variables section"
+    
+    def test_no_topic_field_no_variable_added(self):
+        """Test that no topic variable is added if topic field is missing."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: No Topic Test
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # topic should NOT be in variables if not specified
+        assert "topic" not in workflow.variables
+
+
+class TestIncludeVariablesMerging:
+    """Tests for variables merging in included recipes."""
+    
+    def test_variables_available_in_workflow_run(self):
+        """Test that variables are available when workflow runs."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Variables Run Test
+topic: "Test Topic"
+
+variables:
+  today: "2026"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{topic}} in {{today}}"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # Both topic and today should be in variables
+        assert "topic" in workflow.variables
+        assert "today" in workflow.variables
+        assert workflow.variables["topic"] == "Test Topic"
+        assert workflow.variables["today"] == "2026"
+
+
+class TestNameFieldPropagation:
+    """Tests for name field propagation to variables."""
+    
+    def test_name_field_added_to_variables(self):
+        """Test that 'name' field is automatically added to workflow variables."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: My Workflow Name
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # name should be in variables dict for template usage
+        assert "name" in workflow.variables or workflow.name == "My Workflow Name"
+
+
+class TestInputFieldPropagation:
+    """Tests for input field propagation to variables."""
+    
+    def test_input_field_added_to_variables(self):
+        """Test that 'input' field is automatically added to workflow variables."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Input Test
+input: "Initial input text"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{input}}"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        # default_input should be set
+        assert workflow.default_input == "Initial input text"
+
+
+class TestIncludeWithVariablesSupport:
+    """Tests for enhanced include syntax with variables support."""
+    
+    def test_include_class_has_variables_attribute(self):
+        """Test that Include class supports variables attribute."""
+        from praisonaiagents.workflows import Include
+        
+        # Include should accept variables parameter
+        include = Include(recipe="test-recipe", input=None)
+        assert hasattr(include, 'recipe')
+        assert include.recipe == "test-recipe"
+    
+    def test_parse_include_with_variables(self):
+        """Test parsing include step with variables override."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Include
+        
+        yaml_content = """
+name: Include Variables Test
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work"
+  - include:
+      recipe: test-recipe
+      variables:
+        topic: "Override Topic"
+        custom_var: "Custom Value"
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        assert len(workflow.steps) == 2
+        
+        # Second step should be Include with variables
+        include_step = workflow.steps[1]
+        assert isinstance(include_step, Include)
+        if hasattr(include_step, 'variables') and include_step.variables:
+            assert include_step.variables.get("topic") == "Override Topic"
+            assert include_step.variables.get("custom_var") == "Custom Value"
+
+
+class TestPerformanceConstraints:
+    """Tests to ensure no performance regression."""
+    
+    def test_import_time_acceptable(self):
+        """Test that import time is not degraded."""
+        import time
+        
+        start = time.perf_counter()
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        elapsed = time.perf_counter() - start
+        
+        # Import should be fast (< 1 second even with all dependencies)
+        assert elapsed < 1.0, f"Import took {elapsed:.3f}s, expected < 1.0s"
+    
+    def test_parse_time_acceptable(self):
+        """Test that parsing is fast."""
+        import time
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = """
+name: Performance Test
+topic: "AI Code & Tools {{today}}"
+
+variables:
+  today: "January 2026"
+
+agents:
+  agent1:
+    name: Agent1
+    role: Worker
+    goal: Work
+    instructions: "Do work"
+
+steps:
+  - agent: agent1
+    action: "Work on {{topic}}"
+"""
+        parser = YAMLWorkflowParser()
+        
+        start = time.perf_counter()
+        for _ in range(10):
+            workflow = parser.parse_string(yaml_content)
+        elapsed = time.perf_counter() - start
+        
+        # 10 parses should be fast (< 1 second)
+        assert elapsed < 1.0, f"10 parses took {elapsed:.3f}s, expected < 1.0s"

src/praisonai-agents/tests/test_native_structured_outputs.py

@@ -0,0 +1,211 @@
+"""
+Tests for native structured outputs via response_format parameter.
+
+Tests cover:
+1. OpenAIClient.create_completion() accepts response_format param
+2. Agent._build_response_format() correctly builds format from schema
+3. Agent uses native mode for supporting models
+4. Fallback to text injection for unsupported models
+"""
+import pytest
+from unittest.mock import Mock, patch, MagicMock
+from typing import List
+
+
+class TestOpenAIClientResponseFormat:
+    """Test that OpenAIClient.create_completion() accepts response_format."""
+    
+    def test_create_completion_accepts_response_format(self):
+        """Test that create_completion passes response_format to API via kwargs."""
+        from praisonaiagents.llm.openai_client import OpenAIClient
+        
+        # Create client with mocked OpenAI - need to patch at module level before import
+        with patch('praisonaiagents.llm.openai_client._get_openai') as mock_get_openai:
+            mock_openai_module = MagicMock()
+            mock_get_openai.return_value = mock_openai_module
+            
+            # Create completion response mock
+            mock_response = MagicMock()
+            mock_choice = MagicMock()
+            mock_message = MagicMock(content="test response")
+            mock_choice.message = mock_message
+            mock_response.choices = [mock_choice]
+            
+            # Create the mock client as the actual OpenAI client would return
+            mock_client = MagicMock()
+            mock_client.chat.completions.create.return_value = mock_response
+            mock_openai_module.OpenAI.return_value = mock_client
+            
+            client = OpenAIClient(api_key="test-key")
+            
+            # Force the sync_client to use our mock
+            client._OpenAIClient__sync_client = mock_client
+            
+            response_format = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "test_output",
+                    "schema": {"type": "object"},
+                    "strict": True
+                }
+            }
+            
+            # Call create_completion with response_format
+            client.create_completion(
+                messages=[{"role": "user", "content": "test"}],
+                model="gpt-4o-mini",
+                response_format=response_format
+            )
+            
+            # Verify the create method was called
+            mock_client.chat.completions.create.assert_called_once()
+            
+            # Get the call args and verify response_format was included
+            call_args = mock_client.chat.completions.create.call_args
+            assert call_args is not None
+            # response_format comes through **kwargs
+            assert call_args.kwargs.get('response_format') == response_format
+
+
+class TestAgentBuildResponseFormat:
+    """Test Agent._build_response_format() helper method."""
+    
+    def test_build_response_format_from_pydantic(self):
+        """Test building response_format from Pydantic model."""
+        from praisonaiagents.agent.agent import Agent
+        from pydantic import BaseModel
+        
+        class TopicList(BaseModel):
+            topics: List[str]
+        
+        agent = Agent(name="test", llm="gpt-4o-mini")
+        
+        # Test the helper if it exists
+        if hasattr(agent, '_build_response_format'):
+            result = agent._build_response_format(TopicList)
+            assert result is not None
+            assert result.get("type") == "json_schema"
+            assert "json_schema" in result
+            assert result["json_schema"]["name"] == "TopicList"
+    
+    def test_build_response_format_from_dict(self):
+        """Test building response_format from dict schema."""
+        from praisonaiagents.agent.agent import Agent
+        
+        schema = {
+            "type": "array",
+            "items": {"type": "string"}
+        }
+        
+        agent = Agent(name="test", llm="gpt-4o-mini")
+        
+        if hasattr(agent, '_build_response_format'):
+            result = agent._build_response_format(schema)
+            assert result is not None
+            assert result.get("type") == "json_schema"
+
+
+class TestAgentNativeStructuredOutput:
+    """Test that Agent uses native mode for supporting models."""
+    
+    def test_agent_chat_uses_response_format_for_supported_model(self):
+        """Test that agent.chat() uses response_format for gpt-4o."""
+        from praisonaiagents.agent.agent import Agent
+        from pydantic import BaseModel
+        
+        class TopicList(BaseModel):
+            topics: List[str]
+        
+        # Mock the LLM call
+        with patch('praisonaiagents.agent.agent.Agent._chat_completion') as mock_chat:
+            mock_response = MagicMock()
+            mock_response.choices = [MagicMock()]
+            mock_response.choices[0].message.content = '{"topics": ["AI", "ML"]}'
+            mock_chat.return_value = mock_response
+            
+            with patch('praisonaiagents.llm.model_capabilities.supports_structured_outputs', return_value=True):
+                agent = Agent(name="test", llm="gpt-4o-mini")
+                agent.chat("List topics", output_pydantic=TopicList)
+                
+                # Check if response_format was passed (when method supports it)
+                if mock_chat.called:
+                    call_kwargs = mock_chat.call_args
+                    # The actual check depends on implementation
+
+
+class TestFallbackBehavior:
+    """Test fallback to text injection for unsupported models."""
+    
+    def test_text_injection_for_unsupported_model(self):
+        """Test that text injection is used when model doesn't support response_format."""
+        from praisonaiagents.agent.agent import Agent
+        
+        with patch('praisonaiagents.llm.model_capabilities.supports_structured_outputs', return_value=False):
+            agent = Agent(name="test", llm="some-unsupported-model")
+            
+            # Build messages should include text injection
+            schema = {"type": "array", "items": {"type": "string"}}
+            messages, _ = agent._build_messages(
+                "List topics",
+                output_json=schema
+            )
+            
+            # Check that schema instruction was added to message
+            user_message = messages[-1]["content"]
+            assert "JSON" in user_message or "schema" in user_message.lower()
+
+
+class TestSchemaEchoExtraction:
+    """Test that schema echo extraction still works as safety net."""
+    
+    def test_parse_json_output_extracts_data(self):
+        """Test that _parse_json_output extracts actual data from schema echo."""
+        from praisonaiagents.workflows.workflows import _parse_json_output
+        
+        # Simulated schema echo from LLM
+        echoed_output = {
+            "type": "array",
+            "items": ["topic1", "topic2", "topic3"]
+        }
+        
+        result = _parse_json_output(echoed_output, "test_step")
+        
+        # Should extract the actual array data
+        assert isinstance(result, list)
+        assert result == ["topic1", "topic2", "topic3"]
+    
+    def test_parse_json_output_handles_clean_data(self):
+        """Test that _parse_json_output passes through clean data."""
+        from praisonaiagents.workflows.workflows import _parse_json_output
+        
+        # Clean output (no schema echo)
+        clean_output = ["topic1", "topic2", "topic3"]
+        
+        result = _parse_json_output(clean_output, "test_step")
+        
+        assert result == clean_output
+
+
+class TestPerformanceImpact:
+    """Test that there's no performance impact from this change."""
+    
+    def test_no_heavy_imports_in_agent_module(self):
+        """Test that agent module doesn't import heavy dependencies at module level."""
+        import sys
+        
+        # Remove agent module if already imported
+        if 'praisonaiagents.agent.agent' in sys.modules:
+            del sys.modules['praisonaiagents.agent.agent']
+        
+        # Import and check it doesn't fail
+        from praisonaiagents.agent import agent
+        assert agent is not None
+    
+    def test_model_capabilities_lazy_import(self):
+        """Test that model_capabilities uses lazy litellm import."""
+        from praisonaiagents.llm import model_capabilities
+        
+        # The function should work even if litellm is not installed
+        # (it should gracefully return False)
+        result = model_capabilities.supports_structured_outputs("")
+        assert result == False  # Empty model should return False

src/praisonai-agents/tests/test_output_variable_and_list_parsing.py

@@ -0,0 +1,282 @@
+"""
+TDD tests for output_variable parsing and list parsing in loops.
+
+Tests for:
+1. output_variable parsing from YAML steps
+2. JSON list parsing when loop iterates over string variable
+3. Regex extraction of lists from agent text output
+4. Dynamic parallel execution based on list size
+"""
+
+import pytest
+import json
+
+
+class TestOutputVariableParsing:
+    """Tests for output_variable parsing from YAML."""
+    
+    def test_output_variable_parsed_from_yaml_step(self):
+        """Test that output_variable is parsed from YAML step definition."""
+        from praisonaiagents.workflows import YAMLWorkflowParser
+        
+        yaml_content = '''
+agents:
+  topic_finder:
+    role: Topic Finder
+    goal: Find topics
+    
+steps:
+  - name: find_topics
+    agent: topic_finder
+    action: "Find 5 topics"
+    output_variable: my_topics
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # The agent should have _yaml_output_variable set
+        assert workflow is not None
+        assert len(workflow.steps) > 0
+        
+        # Check that the agent has the output_variable attribute
+        step = workflow.steps[0]
+        output_var = getattr(step, '_yaml_output_variable', None)
+        assert output_var == 'my_topics', f"Expected 'my_topics', got {output_var}"
+    
+    def test_output_variable_used_in_workflow_execution(self):
+        """Test that output_variable stores output in the specified variable name."""
+        from praisonaiagents.workflows import WorkflowStep
+        
+        # Create a simple workflow step with output_variable
+        step = WorkflowStep(
+            name="test_step",
+            action="Return a list",
+            output={"variable": "custom_var"}
+        )
+        
+        # Verify the output_variable is set
+        assert step.output_variable == "custom_var"
+
+
+class TestListParsingInLoop:
+    """Tests for parsing lists from string variables in loops."""
+    
+    def test_json_array_string_parsed_as_list(self):
+        """Test that a JSON array string is parsed into a Python list."""
+        # Simulate what _execute_loop should do
+        json_string = '["topic1", "topic2", "topic3"]'
+        
+        # Current behavior (wrong): wraps string as single item
+        # Expected behavior: parse JSON and get list
+        try:
+            parsed = json.loads(json_string)
+            assert isinstance(parsed, list)
+            assert len(parsed) == 3
+            assert parsed == ["topic1", "topic2", "topic3"]
+        except json.JSONDecodeError:
+            pytest.fail("Should parse valid JSON array")
+    
+    def test_extract_json_array_from_text(self):
+        """Test extracting JSON array from text that contains other content."""
+        import re
+        
+        text_with_list = '''
+Here are the unique topics I found:
+UNIQUE_TOPICS: ["AI Agents", "Machine Learning", "Neural Networks"]
+
+These topics are trending this week.
+'''
+        # Extract JSON array using regex
+        pattern = r'\[(?:[^\[\]]*(?:"[^"]*")?[^\[\]]*)*\]'
+        matches = re.findall(pattern, text_with_list)
+        
+        assert len(matches) > 0, "Should find at least one array pattern"
+        
+        # Try to parse the first match as JSON
+        for match in matches:
+            try:
+                parsed = json.loads(match)
+                if isinstance(parsed, list) and len(parsed) > 0:
+                    assert parsed == ["AI Agents", "Machine Learning", "Neural Networks"]
+                    return
+            except json.JSONDecodeError:
+                continue
+        
+        pytest.fail("Should extract and parse the JSON array from text")
+    
+    def test_loop_over_variable_parses_json_list(self):
+        """Test that loop over a string variable containing JSON list works."""
+        from praisonaiagents.workflows import Loop
+        
+        # Create a loop that iterates over a variable
+        loop_step = Loop(
+            step=None,
+            over="topics",
+            parallel=True
+        )
+        
+        # The variable contains a JSON string
+        variables = {
+            "topics": '["topic1", "topic2", "topic3"]'
+        }
+        
+        # Get items - this is what _execute_loop does
+        items = variables.get(loop_step.over, [])
+        
+        # Current behavior: if string, wrap as single item
+        # We need to add JSON parsing
+        if isinstance(items, str):
+            try:
+                parsed = json.loads(items)
+                if isinstance(parsed, list):
+                    items = parsed
+            except json.JSONDecodeError:
+                items = [items]
+        
+        assert isinstance(items, list)
+        assert len(items) == 3
+        assert items == ["topic1", "topic2", "topic3"]
+
+
+class TestDynamicParallelism:
+    """Tests for dynamic parallel execution based on list size."""
+    
+    def test_parallel_workers_match_list_size(self):
+        """Test that parallel workers default to list size."""
+        from praisonaiagents.workflows import Loop
+        
+        # Create a parallel loop without max_workers
+        loop_step = Loop(
+            step=None,
+            over="items",
+            parallel=True,
+            max_workers=None  # Should default to num_items
+        )
+        
+        # Simulate items
+        items = ["a", "b", "c", "d", "e"]
+        num_items = len(items)
+        
+        # Calculate max_workers as _execute_loop does
+        max_workers = loop_step.max_workers or num_items
+        
+        assert max_workers == 5, "max_workers should equal list size when not specified"
+    
+    def test_max_workers_caps_parallelism(self):
+        """Test that max_workers caps the number of parallel workers."""
+        from praisonaiagents.workflows import Loop
+        
+        loop_step = Loop(
+            step=None,
+            over="items",
+            parallel=True,
+            max_workers=2  # Cap at 2
+        )
+        
+        items = ["a", "b", "c", "d", "e"]
+        num_items = len(items)
+        
+        max_workers = loop_step.max_workers or num_items
+        
+        assert max_workers == 2, "max_workers should be capped at specified value"
+
+
+class TestParseListFromString:
+    """Tests for the _parse_list_from_string method in Workflow."""
+    
+    def test_parse_pure_json_array(self):
+        """Test parsing a pure JSON array string."""
+        from praisonaiagents.workflows import Workflow
+        
+        workflow = Workflow(steps=[])
+        result = workflow._parse_list_from_string('["topic1", "topic2", "topic3"]')
+        
+        assert isinstance(result, list)
+        assert len(result) == 3
+        assert result == ["topic1", "topic2", "topic3"]
+    
+    def test_parse_json_array_from_text(self):
+        """Test extracting JSON array from text with other content."""
+        from praisonaiagents.workflows import Workflow
+        
+        workflow = Workflow(steps=[])
+        text = '''
+Here are the unique topics I found:
+UNIQUE_TOPICS: ["AI Agents", "Machine Learning", "Neural Networks"]
+
+These topics are trending this week.
+'''
+        result = workflow._parse_list_from_string(text)
+        
+        assert isinstance(result, list)
+        assert len(result) == 3
+        assert result == ["AI Agents", "Machine Learning", "Neural Networks"]
+    
+    def test_parse_empty_string_returns_empty_list(self):
+        """Test that empty string returns empty list."""
+        from praisonaiagents.workflows import Workflow
+        
+        workflow = Workflow(steps=[])
+        result = workflow._parse_list_from_string('')
+        
+        assert result == []
+    
+    def test_parse_non_list_string_wraps_as_single_item(self):
+        """Test that non-list string is wrapped as single item."""
+        from praisonaiagents.workflows import Workflow
+        
+        workflow = Workflow(steps=[])
+        result = workflow._parse_list_from_string('just a plain string')
+        
+        assert isinstance(result, list)
+        assert len(result) == 1
+        assert result == ['just a plain string']
+
+
+class TestBeginnerFriendlyYAML:
+    """Tests for beginner-friendly YAML syntax."""
+    
+    def test_simple_parallel_loop_syntax(self):
+        """Test that simple parallel loop syntax works."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Loop
+        
+        yaml_content = '''
+agents:
+  researcher:
+    role: Researcher
+    goal: Research topics
+    
+steps:
+  # Step 1: Find topics (output stored in topics variable)
+  - name: find_topics
+    agent: researcher
+    action: "Find 3 topics and return as JSON array"
+    output_variable: topics
+    
+  # Step 2: Process each topic in parallel
+  - loop:
+      over: topics
+      parallel: true
+    agent: researcher
+    action: "Research this topic: {{item}}"
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        assert workflow is not None
+        assert len(workflow.steps) >= 2
+        
+        # Second step should be a Loop
+        loop_found = False
+        for step in workflow.steps:
+            if isinstance(step, Loop):
+                assert step.parallel is True
+                assert step.over == "topics"
+                loop_found = True
+                break
+        
+        assert loop_found, "Should have a parallel loop step"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/test_parallel_loop.py

@@ -0,0 +1,202 @@
+"""
+TDD tests for Parallel Loop pattern in YAML workflows.
+
+Tests for parallel execution of loops:
+- Loop with parallel=True executes items concurrently
+- max_workers parameter limits parallelism  
+- All outputs are collected correctly
+- Variables are isolated between parallel executions
+"""
+
+import pytest
+import time
+import threading
+from typing import List
+
+
+class TestParallelLoopClass:
+    """Tests for Loop class with parallel parameter."""
+    
+    def test_loop_class_has_parallel_param(self):
+        """Test that Loop class accepts parallel parameter."""
+        from praisonaiagents.workflows import Loop
+        
+        loop = Loop(step=None, over="items", parallel=True)
+        assert loop.parallel is True
+    
+    def test_loop_class_has_max_workers_param(self):
+        """Test that Loop class accepts max_workers parameter."""
+        from praisonaiagents.workflows import Loop
+        
+        loop = Loop(step=None, over="items", parallel=True, max_workers=4)
+        assert loop.max_workers == 4
+    
+    def test_loop_parallel_defaults_to_false(self):
+        """Test that parallel defaults to False (sequential)."""
+        from praisonaiagents.workflows import Loop
+        
+        loop = Loop(step=None, over="items")
+        assert loop.parallel is False
+    
+    def test_loop_max_workers_defaults_to_none(self):
+        """Test that max_workers defaults to None (unlimited)."""
+        from praisonaiagents.workflows import Loop
+        
+        loop = Loop(step=None, over="items")
+        assert loop.max_workers is None
+
+
+class TestParallelLoopExecution:
+    """Tests for parallel execution behavior."""
+    
+    def test_parallel_loop_collects_all_outputs(self):
+        """Test that parallel loop collects outputs from all items."""
+        from praisonaiagents.workflows import Workflow, Loop, loop
+        from praisonaiagents import Agent
+        
+        # Create a simple agent that echoes the item
+        agent = Agent(
+            name="echo_agent",
+            role="Echo",
+            goal="Echo the input",
+            instructions="Just repeat what you're given",
+            llm="gpt-4o-mini"
+        )
+        
+        items = ["topic_a", "topic_b", "topic_c"]
+        
+        # Test that loop with parallel=True can be created
+        loop_step = loop(
+            step=agent,
+            over="items",
+            parallel=True
+        )
+        
+        # Verify the Loop object is configured correctly
+        assert loop_step.parallel is True
+        assert loop_step.over == "items"
+        assert loop_step.step == agent
+    
+    def test_parallel_loop_faster_than_sequential(self):
+        """Test that parallel loop with sleep is faster than sequential."""
+        import concurrent.futures
+        
+        # Simulate work by tracking timestamps
+        execution_times: List[float] = []
+        execution_lock = threading.Lock()
+        
+        def work_item(item: str, delay: float = 0.1):
+            """Simulated work that takes time."""
+            start = time.time()
+            time.sleep(delay)
+            with execution_lock:
+                execution_times.append(time.time() - start)
+            return f"processed_{item}"
+        
+        items = ["a", "b", "c", "d"]
+        
+        # Sequential execution
+        seq_start = time.time()
+        for item in items:
+            work_item(item)
+        seq_duration = time.time() - seq_start
+        
+        execution_times.clear()
+        
+        # Parallel execution
+        par_start = time.time()
+        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+            futures = [executor.submit(work_item, item) for item in items]
+            concurrent.futures.wait(futures)
+        par_duration = time.time() - par_start
+        
+        # Parallel should be roughly 4x faster
+        assert par_duration < seq_duration * 0.5, \
+            f"Parallel ({par_duration:.2f}s) should be faster than sequential ({seq_duration:.2f}s)"
+
+
+class TestParallelLoopYAMLParsing:
+    """Tests for parsing parallel loop from YAML."""
+    
+    def test_parse_parallel_true_from_yaml(self):
+        """Test parsing loop with parallel: true from YAML."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Loop
+        
+        # Use dict format for agents (not list)
+        yaml_content = '''
+agents:
+  processor:
+    role: Processor
+    goal: Process items
+    
+steps:
+  - loop:
+      over: items
+      parallel: true
+    agent: processor
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Check that the workflow was parsed
+        assert workflow is not None
+        # The Loop object with parallel=True should be in steps
+        found_parallel_loop = False
+        for step in workflow.steps:
+            if isinstance(step, Loop):
+                if step.parallel:
+                    found_parallel_loop = True
+                    break
+        assert found_parallel_loop, "Expected a parallel Loop step"
+
+    def test_parse_max_workers_from_yaml(self):
+        """Test parsing loop with max_workers from YAML."""
+        from praisonaiagents.workflows import YAMLWorkflowParser, Loop
+        
+        # Use dict format for agents
+        yaml_content = '''
+agents:
+  processor:
+    role: Processor
+    goal: Process items
+    
+steps:
+  - loop:
+      over: items
+      parallel: true
+      max_workers: 8
+    agent: processor
+'''
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Find the loop step with max_workers
+        found_max_workers = False
+        for step in workflow.steps:
+            if isinstance(step, Loop):
+                if step.max_workers == 8:
+                    found_max_workers = True
+                    break
+        assert found_max_workers, "Expected Loop with max_workers=8"
+
+
+class TestLoopConvenienceFunction:
+    """Tests for loop() convenience function with parallel param."""
+    
+    def test_loop_function_accepts_parallel(self):
+        """Test that loop() function accepts parallel parameter."""
+        from praisonaiagents.workflows import loop
+        
+        l = loop(step="agent_name", over="items", parallel=True)
+        assert l.parallel is True
+    
+    def test_loop_function_accepts_max_workers(self):
+        """Test that loop() function accepts max_workers parameter."""
+        from praisonaiagents.workflows import loop
+        
+        l = loop(step="agent_name", over="items", parallel=True, max_workers=4)
+        assert l.max_workers == 4
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/tests/test_parallel_loop_live.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""
+Test script for parallel loop execution.
+
+This demonstrates the parallel loop pattern works correctly by:
+1. Creating a simple workflow with a parallel loop
+2. Using a mock agent that simulates work with a delay
+3. Verifying that parallel execution is faster than sequential
+"""
+
+import time
+from praisonaiagents.workflows import Workflow, loop
+from praisonaiagents import Agent
+
+
+def test_parallel_loop_execution():
+    """Test parallel loop executes items concurrently."""
+    
+    # Create a simple agent
+    agent = Agent(
+        name="processor",
+        role="Item Processor",
+        goal="Process items",
+        instructions="Just echo back the item you received",
+        llm="gpt-4o-mini"  # Fast, cheap model
+    )
+    
+    items = ["Apple", "Google", "Microsoft", "Amazon"]
+    
+    # Create workflow with parallel loop (variables in constructor)
+    workflow = Workflow(
+        steps=[
+            loop(
+                step=agent,
+                over="items",
+                parallel=True,
+                max_workers=4
+            )
+        ],
+        variables={"items": items}
+    )
+    
+    print(f"\n🚀 Testing parallel loop with {len(items)} items...")
+    print(f"   Items: {items}")
+    
+    start = time.time()
+    result = workflow.start(input=f"Process these items")
+    duration = time.time() - start
+    
+    print(f"\n✅ Parallel loop completed in {duration:.2f}s")
+    print(f"   Output preview: {str(result)[:200]}...")
+    
+    return result
+
+
+def test_sequential_vs_parallel():
+    """Compare sequential vs parallel loop execution."""
+    
+    agent = Agent(
+        name="processor",
+        role="Item Processor",
+        goal="Process items",
+        instructions="Just echo back the item you received",
+        llm="gpt-4o-mini"
+    )
+    
+    items = ["A", "B", "C", "D"]
+    
+    # Sequential loop
+    seq_workflow = Workflow(
+        steps=[loop(step=agent, over="items", parallel=False)],
+        variables={"items": items}
+    )
+    
+    print("\n🔄 Running SEQUENTIAL loop...")
+    seq_start = time.time()
+    seq_result = seq_workflow.start(input="Process items")
+    seq_duration = time.time() - seq_start
+    print(f"   Sequential duration: {seq_duration:.2f}s")
+    
+    # Parallel loop
+    par_workflow = Workflow(
+        steps=[loop(step=agent, over="items", parallel=True, max_workers=4)],
+        variables={"items": items}
+    )
+    
+    print("\n⚡ Running PARALLEL loop...")
+    par_start = time.time()
+    par_result = par_workflow.start(input="Process items")
+    par_duration = time.time() - par_start
+    print(f"   Parallel duration: {par_duration:.2f}s")
+    
+    # Compare
+    speedup = seq_duration / par_duration if par_duration > 0 else 0
+    print(f"\n📊 Speedup: {speedup:.1f}x")
+    
+    if par_duration < seq_duration:
+        print("✅ Parallel loop is faster!")
+    else:
+        print("⚠️ Parallel loop was not faster (might be due to overhead with small items)")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("PARALLEL LOOP EXECUTION TEST")
+    print("=" * 60)
+    
+    # Run basic test
+    test_parallel_loop_execution()
+    
+    print("\n" + "=" * 60)
+    print("SEQUENTIAL vs PARALLEL COMPARISON")
+    print("=" * 60)
+    
+    # Run comparison
+    test_sequential_vs_parallel()
+    
+    print("\n✅ All tests complete!")

src/praisonai-agents/tests/test_structured_output.py

@@ -0,0 +1,241 @@
+"""
+Tests for structured output functionality in workflows.
+
+Tests cover:
+1. output_json in YAML step passes to agent.chat()
+2. Inline JSON schema (Option A) is supported
+3. Pydantic model reference (Option B) is resolved
+4. response_format is set for supporting models
+5. JSON response is validated against schema
+6. Loop receives typed list when output_json is array
+7. Backward compatibility with existing workflows
+"""
+from unittest.mock import Mock, patch
+from typing import List
+
+
+class TestStructuredOutputYAMLParsing:
+    """Test that output_json is correctly parsed from YAML."""
+    
+    def test_parse_inline_json_schema(self):
+        """Test Option A: Inline JSON schema in YAML."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  topic_finder:
+    role: Topic Finder
+    goal: Find topics
+    backstory: Expert at finding topics
+
+steps:
+  - agent: topic_finder
+    action: "Find AI topics"
+    output_json:
+      type: array
+      items:
+        type: object
+        properties:
+          title:
+            type: string
+          url:
+            type: string
+        required:
+          - title
+          - url
+    output_variable: topics
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Verify the step has output_json
+        assert len(workflow.steps) == 1
+        step = workflow.steps[0]
+        
+        # Check that _yaml_output_json was parsed (stored on agent, then transferred to step)
+        # The step is an Agent with _yaml_output_json attribute
+        assert hasattr(step, '_yaml_output_json')
+        
+    def test_parse_pydantic_reference(self):
+        """Test Option B: Pydantic model reference in YAML."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  topic_finder:
+    role: Topic Finder
+    goal: Find topics
+    backstory: Expert at finding topics
+
+steps:
+  - agent: topic_finder
+    action: "Find AI topics"
+    output_pydantic: TopicList
+    output_variable: topics
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Verify the step has output_pydantic
+        assert len(workflow.steps) == 1
+        step = workflow.steps[0]
+        
+        # Check that _yaml_output_pydantic was parsed (stored on agent)
+        assert hasattr(step, '_yaml_output_pydantic')
+
+
+class TestStructuredOutputWiring:
+    """Test that output_json is wired from YAML to WorkflowStep."""
+    
+    def test_normalize_step_includes_output_json(self):
+        """Test that _normalize_single_step transfers output_json to WorkflowStep."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep
+        
+        # Create a mock agent with _yaml_output_json
+        agent = Mock()
+        agent.name = "test_agent"
+        agent.tools = None
+        agent._yaml_action = "Test action"
+        agent._yaml_output_variable = "result"
+        agent._yaml_output_json = {"type": "array", "items": {"type": "string"}}
+        agent._yaml_step_name = "test_step"
+        
+        # Mock the chat method
+        agent.chat = Mock(return_value="test output")
+        
+        # Create workflow and normalize the step
+        workflow = Workflow(steps=[])
+        normalized = workflow._normalize_single_step(agent, 0)
+        
+        # Verify WorkflowStep has output_json
+        assert isinstance(normalized, WorkflowStep)
+        assert normalized._output_json == {"type": "array", "items": {"type": "string"}}
+        
+    def test_normalize_step_includes_output_pydantic(self):
+        """Test that _normalize_single_step transfers output_pydantic to WorkflowStep."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep
+        
+        # Create a mock agent with _yaml_output_pydantic
+        agent = Mock()
+        agent.name = "test_agent"
+        agent.tools = None
+        agent._yaml_action = "Test action"
+        agent._yaml_output_variable = "result"
+        agent._yaml_output_pydantic = "TopicList"
+        agent._yaml_step_name = "test_step"
+        agent.chat = Mock(return_value="test output")
+        
+        workflow = Workflow(steps=[])
+        normalized = workflow._normalize_single_step(agent, 0)
+        
+        assert isinstance(normalized, WorkflowStep)
+        assert normalized._output_pydantic == "TopicList"
+
+
+class TestStructuredOutputExecution:
+    """Test that output_json is passed to agent.chat() during execution."""
+    
+    def test_agent_chat_receives_output_json(self):
+        """Test that agent.chat() is called with output_json parameter."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep
+        
+        # Create mock agent
+        mock_agent = Mock()
+        mock_agent.name = "test_agent"
+        mock_agent.chat = Mock(return_value='["topic1", "topic2"]')
+        
+        # Create step with output_json
+        step = WorkflowStep(
+            name="test_step",
+            agent=mock_agent,
+            action="Find topics",
+            output={"json_model": {"type": "array", "items": {"type": "string"}}}
+        )
+        
+        # Create and run workflow
+        workflow = Workflow(steps=[step])
+        
+        # Execute the step
+        _ = workflow.run("test input")
+        
+        # Verify agent.chat was called
+        mock_agent.chat.assert_called()
+
+
+class TestResponseFormatForAllModels:
+    """Test that response_format is used for all supporting models."""
+    
+    def test_response_format_set_for_openai_models(self):
+        """Test that response_format is set for OpenAI models that support it."""
+        from praisonaiagents.llm.llm import LLM
+        from pydantic import BaseModel
+        
+        class TestSchema(BaseModel):
+            items: List[str]
+        
+        llm = LLM(model="gpt-4o-mini")
+        
+        # Build completion params with output_json
+        with patch('praisonaiagents.llm.model_capabilities.supports_structured_outputs', return_value=True):
+            params = llm._build_completion_params(
+                messages=[{"role": "user", "content": "test"}],
+                output_json=TestSchema
+            )
+        
+        # Verify response_format is set
+        assert 'response_format' in params or 'response_schema' in params
+
+
+class TestJSONValidation:
+    """Test that JSON responses are validated against schema."""
+    
+    def test_invalid_json_raises_error(self):
+        """Test that invalid JSON response raises validation error."""
+        # This test verifies that when output_json is specified,
+        # the response is validated against the schema
+        pass  # Will be implemented after core functionality
+
+
+class TestLoopWithStructuredOutput:
+    """Test that loops receive typed lists from structured output."""
+    
+    def test_loop_receives_list_not_string(self):
+        """Test that loop iteration receives a list, not a JSON string."""
+        # This test verifies that when output_json specifies an array,
+        # the loop receives an actual list to iterate over
+        pass  # Will be implemented after core functionality
+
+
+class TestBackwardCompatibility:
+    """Test backward compatibility with existing workflows."""
+    
+    def test_workflow_without_output_json_still_works(self):
+        """Test that workflows without output_json continue to work."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep
+        
+        mock_agent = Mock()
+        mock_agent.name = "test_agent"
+        mock_agent.chat = Mock(return_value="test output")
+        
+        step = WorkflowStep(
+            name="test_step",
+            agent=mock_agent,
+            action="Test action"
+        )
+        
+        workflow = Workflow(steps=[step])
+        result = workflow.run("test input")
+        
+        # Should complete without error
+        assert result is not None
+
+
+class TestPerformance:
+    """Test that there's no performance regression."""
+    
+    def test_no_heavy_imports_at_module_level(self):
+        """Test that Pydantic and other heavy deps are lazily imported."""
+        # This is a basic check - full perf testing would measure import time
+        # The workflows module should import without errors
+        from praisonaiagents.workflows import workflows as wf
+        assert wf is not None

src/praisonai-agents/tests/test_yaml_workflow_parser.py

@@ -364,7 +364,8 @@ steps:
         workflow = parser.parse_string(yaml_content)
         
         assert workflow is not None
-        assert workflow.planning == True
+        # planning field holds WorkflowPlanningConfig, check _planning_enabled for bool
+        assert workflow._planning_enabled == True
         assert workflow.reasoning == True
         assert workflow.verbose == True
     
@@ -376,9 +377,11 @@ steps:
 name: Memory Test
 workflow:
   memory_config:
-    provider: chroma
-    persist: true
-    collection: test_collection
+    backend: chroma
+    user_id: test_user
+    config:
+      persist: true
+      collection: test_collection
 
 agents:
   agent1:
@@ -396,7 +399,7 @@ steps:
         
         assert workflow is not None
         assert workflow.memory_config is not None
-        assert workflow.memory_config.get("provider") == "chroma"
+        assert workflow.memory_config.get("backend") == "chroma"
 
 
 class TestYAMLWorkflowParserCallbacks:
@@ -579,11 +582,9 @@ framework: praisonai
 
 workflow:
   planning: true
+  planning_llm: gpt-4o
   reasoning: true
   verbose: true
-  memory_config:
-    provider: chroma
-    persist: true
 
 variables:
   topic: AI trends
@@ -646,7 +647,7 @@ callbacks:
         
         assert workflow is not None
         assert workflow.name == "Complete Workflow"
-        assert workflow.planning == True
+        assert workflow._planning_enabled == True
         assert workflow.reasoning == True
         assert workflow.variables.get("topic") == "AI trends"
         assert len(workflow.steps) == 4
@@ -794,6 +795,7 @@ topic: "Research AI trends"
 
 workflow:
   planning: true
+  planning_llm: gpt-4o
   reasoning: true
   verbose: true
 
@@ -827,7 +829,7 @@ steps:
         
         assert workflow is not None
         assert workflow.name == "Extended Agents YAML"
-        assert workflow.planning == True
+        assert workflow._planning_enabled == True
         assert workflow.reasoning == True
         assert workflow.variables.get("topic") == "AI trends"
         assert len(workflow.steps) == 2

src/praisonai-agents/tests/unit/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for PraisonAI Agents."""

src/praisonai-agents/tests/unit/embedding/__init__.py

@@ -0,0 +1 @@
+# Embedding module tests

src/praisonai-agents/tests/unit/embedding/test_embedding.py

@@ -0,0 +1,313 @@
+"""
+Unit tests for praisonaiagents.embedding module.
+
+TDD: These tests define the expected behavior of the embedding module.
+"""
+
+import pytest
+from unittest.mock import Mock, patch
+
+
+class TestEmbeddingResult:
+    """Tests for EmbeddingResult dataclass."""
+    
+    def test_embedding_result_dataclass(self):
+        """Test EmbeddingResult dataclass creation."""
+        from praisonaiagents.embedding import EmbeddingResult
+        
+        result = EmbeddingResult(
+            embeddings=[[0.1, 0.2, 0.3]],
+            model="text-embedding-3-small",
+            usage={"prompt_tokens": 10, "total_tokens": 10},
+            metadata={"test": "value"}
+        )
+        
+        assert result.embeddings == [[0.1, 0.2, 0.3]]
+        assert result.model == "text-embedding-3-small"
+        assert result.usage == {"prompt_tokens": 10, "total_tokens": 10}
+        assert result.metadata == {"test": "value"}
+    
+    def test_embedding_result_defaults(self):
+        """Test EmbeddingResult with default values."""
+        from praisonaiagents.embedding import EmbeddingResult
+        
+        result = EmbeddingResult(embeddings=[[0.1, 0.2]])
+        
+        assert result.embeddings == [[0.1, 0.2]]
+        assert result.model is None
+        assert result.usage is None
+        assert result.metadata == {}
+
+
+class TestEmbeddingFunction:
+    """Tests for embedding() function."""
+    
+    @patch('litellm.embedding')
+    def test_embedding_single_text(self, mock_litellm):
+        """Test embedding() with single text input."""
+        from praisonaiagents.embedding import embedding
+        
+        # Mock litellm response
+        mock_item = Mock()
+        mock_item.embedding = [0.1, 0.2, 0.3]
+        mock_response = Mock()
+        mock_response.data = [mock_item]
+        mock_response.usage = Mock(prompt_tokens=5, total_tokens=5)
+        mock_litellm.return_value = mock_response
+        
+        result = embedding("Hello world")
+        
+        assert len(result.embeddings) == 1
+        assert result.embeddings[0] == [0.1, 0.2, 0.3]
+        mock_litellm.assert_called_once()
+    
+    @patch('litellm.embedding')
+    def test_embedding_batch_texts(self, mock_litellm):
+        """Test embedding() with batch text input."""
+        from praisonaiagents.embedding import embedding
+        
+        # Mock litellm response for batch
+        mock_item1 = Mock()
+        mock_item1.embedding = [0.1, 0.2]
+        mock_item2 = Mock()
+        mock_item2.embedding = [0.3, 0.4]
+        mock_response = Mock()
+        mock_response.data = [mock_item1, mock_item2]
+        mock_response.usage = Mock(prompt_tokens=10, total_tokens=10)
+        mock_litellm.return_value = mock_response
+        
+        result = embedding(["Hello", "World"])
+        
+        assert len(result.embeddings) == 2
+        assert result.embeddings[0] == [0.1, 0.2]
+        assert result.embeddings[1] == [0.3, 0.4]
+    
+    @patch('litellm.embedding')
+    def test_embedding_with_model(self, mock_litellm):
+        """Test embedding() with custom model."""
+        from praisonaiagents.embedding import embedding
+        
+        mock_item = Mock()
+        mock_item.embedding = [0.1]
+        mock_response = Mock()
+        mock_response.data = [mock_item]
+        mock_response.usage = Mock(prompt_tokens=5, total_tokens=5)
+        mock_litellm.return_value = mock_response
+        
+        result = embedding("test", model="text-embedding-3-large")
+        
+        call_kwargs = mock_litellm.call_args[1]
+        assert call_kwargs['model'] == "text-embedding-3-large"
+        assert result.model == "text-embedding-3-large"
+    
+    @patch('litellm.embedding')
+    def test_embedding_with_dimensions(self, mock_litellm):
+        """Test embedding() with dimensions parameter."""
+        from praisonaiagents.embedding import embedding
+        
+        mock_item = Mock()
+        mock_item.embedding = [0.1] * 256
+        mock_response = Mock()
+        mock_response.data = [mock_item]
+        mock_response.usage = Mock(prompt_tokens=5, total_tokens=5)
+        mock_litellm.return_value = mock_response
+        
+        emb_result = embedding("test", dimensions=256)
+        
+        call_kwargs = mock_litellm.call_args[1]
+        assert call_kwargs['dimensions'] == 256
+        assert emb_result is not None
+
+
+class TestAsyncEmbedding:
+    """Tests for aembedding() async function."""
+    
+    @pytest.mark.asyncio
+    @patch('litellm.aembedding')
+    async def test_aembedding_single_text(self, mock_aembedding):
+        """Test aembedding() with single text input."""
+        from praisonaiagents.embedding import aembedding
+        
+        mock_item = Mock()
+        mock_item.embedding = [0.1, 0.2, 0.3]
+        mock_response = Mock()
+        mock_response.data = [mock_item]
+        mock_response.usage = Mock(prompt_tokens=5, total_tokens=5)
+        mock_aembedding.return_value = mock_response
+        
+        result = await aembedding("Hello world")
+        
+        assert len(result.embeddings) == 1
+        assert result.embeddings[0] == [0.1, 0.2, 0.3]
+
+
+class TestGetDimensions:
+    """Tests for get_dimensions() utility."""
+    
+    def test_get_dimensions_known_models(self):
+        """Test get_dimensions() for known models."""
+        from praisonaiagents.embedding import get_dimensions
+        
+        assert get_dimensions("text-embedding-3-small") == 1536
+        assert get_dimensions("text-embedding-3-large") == 3072
+        assert get_dimensions("text-embedding-ada-002") == 1536
+    
+    def test_get_dimensions_partial_match(self):
+        """Test get_dimensions() with partial model name match."""
+        from praisonaiagents.embedding import get_dimensions
+        
+        assert get_dimensions("openai/text-embedding-3-small") == 1536
+        assert get_dimensions("azure/text-embedding-3-large") == 3072
+    
+    def test_get_dimensions_unknown_model(self):
+        """Test get_dimensions() returns default for unknown models."""
+        from praisonaiagents.embedding import get_dimensions
+        
+        assert get_dimensions("unknown-model") == 1536  # Default
+
+
+class TestAliases:
+    """Tests for function aliases."""
+    
+    def test_embed_alias(self):
+        """Test that embed is an alias for embedding."""
+        from praisonaiagents.embedding.embed import embedding, embed
+        
+        assert embed is embedding
+    
+    def test_aembed_alias(self):
+        """Test that aembed is an alias for aembedding."""
+        from praisonaiagents.embedding.embed import aembedding, aembed
+        
+        assert aembed is aembedding
+
+
+class TestLazyImports:
+    """Tests for lazy import behavior."""
+    
+    def test_no_litellm_import_on_module_load(self):
+        """Test that EmbeddingResult can be imported without litellm."""
+        # EmbeddingResult should be importable without litellm
+        from praisonaiagents.embedding import EmbeddingResult
+        
+        # This should work without litellm being imported
+        test_result = EmbeddingResult(embeddings=[[0.1]])
+        assert test_result.embeddings == [[0.1]]
+
+
+class TestBackwardsCompatibility:
+    """Tests for backwards compatibility."""
+    
+    def test_import_embedding_function_from_embedding_module(self):
+        """Test importing embedding function from praisonaiagents.embedding."""
+        from praisonaiagents.embedding import embedding
+        
+        # embedding is the function
+        assert callable(embedding)
+    
+    def test_import_embedding_result(self):
+        """Test importing EmbeddingResult from praisonaiagents.embedding."""
+        from praisonaiagents.embedding import EmbeddingResult
+        
+        assert EmbeddingResult is not None
+
+
+class TestSimplifiedImports:
+    """Tests for simplified import patterns - the main goal of this fix."""
+    
+    def test_from_praisonaiagents_import_embedding_returns_function(self):
+        """Test that 'from praisonaiagents import embedding' returns the function, not module."""
+        from praisonaiagents import embedding
+        
+        # This is the key test - embedding should be a callable function
+        assert callable(embedding), f"Expected function, got {type(embedding)}"
+        assert not isinstance(embedding, type(None))
+        # Should NOT be a module
+        import types
+        assert not isinstance(embedding, types.ModuleType), "embedding should be function, not module"
+    
+    def test_from_praisonaiagents_import_embeddings_returns_function(self):
+        """Test that 'from praisonaiagents import embeddings' (plural) works as alias."""
+        from praisonaiagents import embeddings
+        
+        assert callable(embeddings), f"Expected function, got {type(embeddings)}"
+        import types
+        assert not isinstance(embeddings, types.ModuleType), "embeddings should be function, not module"
+    
+    def test_embedding_and_embeddings_are_same_function(self):
+        """Test that embedding and embeddings are the same function."""
+        from praisonaiagents import embedding, embeddings
+        
+        assert embedding is embeddings, "embedding and embeddings should be the same function"
+    
+    def test_from_praisonaiagents_import_aembedding(self):
+        """Test that 'from praisonaiagents import aembedding' works."""
+        from praisonaiagents import aembedding
+        
+        assert callable(aembedding)
+    
+    def test_from_praisonaiagents_import_aembeddings(self):
+        """Test that 'from praisonaiagents import aembeddings' (plural) works as alias."""
+        from praisonaiagents import aembeddings
+        
+        assert callable(aembeddings)
+    
+    def test_from_praisonaiagents_import_EmbeddingResult(self):
+        """Test that 'from praisonaiagents import EmbeddingResult' works."""
+        from praisonaiagents import EmbeddingResult
+        
+        # Should be a class
+        assert isinstance(EmbeddingResult, type) or hasattr(EmbeddingResult, '__dataclass_fields__')
+    
+    def test_from_praisonaiagents_import_get_dimensions(self):
+        """Test that 'from praisonaiagents import get_dimensions' works."""
+        from praisonaiagents import get_dimensions
+        
+        assert callable(get_dimensions)
+        # Quick sanity check
+        assert get_dimensions("text-embedding-3-small") == 1536
+    
+    def test_all_embedding_imports_together(self):
+        """Test importing all embedding-related items at once."""
+        from praisonaiagents import embedding, embeddings, aembedding, aembeddings, EmbeddingResult, get_dimensions
+        
+        assert callable(embedding)
+        assert callable(embeddings)
+        assert callable(aembedding)
+        assert callable(aembeddings)
+        assert EmbeddingResult is not None
+        assert callable(get_dimensions)
+
+
+class TestLazyLoadingPreserved:
+    """Tests to ensure litellm is still lazy-loaded (no performance impact)."""
+    
+    def test_litellm_not_loaded_on_embedding_import(self):
+        """Test that litellm is NOT loaded when importing embedding function."""
+        import sys
+        
+        # Clear any cached imports
+        modules_to_clear = [k for k in sys.modules.keys() if 'litellm' in k]
+        for mod in modules_to_clear:
+            del sys.modules[mod]
+        
+        # Import the embedding function
+        from praisonaiagents import embedding
+        
+        # litellm should NOT be loaded yet (lazy loading preserved)
+        assert 'litellm' not in sys.modules, "litellm should not be loaded on import"
+    
+    def test_litellm_not_loaded_on_EmbeddingResult_import(self):
+        """Test that litellm is NOT loaded when importing EmbeddingResult."""
+        import sys
+        
+        # Clear any cached imports
+        modules_to_clear = [k for k in sys.modules.keys() if 'litellm' in k]
+        for mod in modules_to_clear:
+            del sys.modules[mod]
+        
+        # Import EmbeddingResult
+        from praisonaiagents import EmbeddingResult
+        
+        # litellm should NOT be loaded
+        assert 'litellm' not in sys.modules, "litellm should not be loaded on EmbeddingResult import"

src/praisonai-agents/tests/unit/eval/test_grader.py

@@ -0,0 +1,311 @@
+"""
+Tests for BaseLLMGrader and GradeResult.
+
+TDD: Tests for the DRY base grader implementation.
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+
+from praisonaiagents.eval.grader import BaseLLMGrader, GradeResult
+
+
+class TestGradeResult:
+    """Tests for GradeResult dataclass."""
+    
+    def test_grade_result_creation(self):
+        """Test creating a GradeResult."""
+        result = GradeResult(
+            score=8.5,
+            reasoning="Good response",
+            suggestions=["Add more detail"],
+            input_text="What is Python?",
+            output="Python is a programming language",
+            expected_output="Python is a high-level programming language",
+        )
+        
+        assert result.score == 8.5
+        assert result.reasoning == "Good response"
+        assert result.suggestions == ["Add more detail"]
+        assert result.input_text == "What is Python?"
+        assert result.output == "Python is a programming language"
+        assert result.expected_output == "Python is a high-level programming language"
+    
+    def test_grade_result_defaults(self):
+        """Test GradeResult default values."""
+        result = GradeResult(score=7.0, reasoning="OK")
+        
+        assert result.suggestions == []
+        assert result.input_text == ""
+        assert result.output == ""
+        assert result.expected_output is None
+        assert result.timestamp  # Should have a timestamp
+    
+    def test_grade_result_to_dict(self):
+        """Test GradeResult.to_dict()."""
+        result = GradeResult(
+            score=9.0,
+            reasoning="Excellent",
+            suggestions=["Minor improvement"],
+            input_text="test input",
+            output="test output",
+        )
+        
+        d = result.to_dict()
+        
+        assert d["score"] == 9.0
+        assert d["reasoning"] == "Excellent"
+        assert d["suggestions"] == ["Minor improvement"]
+        assert d["input_text"] == "test input"
+        assert d["output"] == "test output"
+        assert "timestamp" in d
+    
+    def test_grade_result_from_dict(self):
+        """Test GradeResult.from_dict()."""
+        data = {
+            "score": 8.0,
+            "reasoning": "Good",
+            "suggestions": ["Suggestion 1"],
+            "input_text": "input",
+            "output": "output",
+            "expected_output": "expected",
+            "timestamp": "2024-01-01T00:00:00",
+        }
+        
+        result = GradeResult.from_dict(data)
+        
+        assert result.score == 8.0
+        assert result.reasoning == "Good"
+        assert result.suggestions == ["Suggestion 1"]
+        assert result.input_text == "input"
+        assert result.output == "output"
+        assert result.expected_output == "expected"
+        assert result.timestamp == "2024-01-01T00:00:00"
+
+
+class TestBaseLLMGrader:
+    """Tests for BaseLLMGrader."""
+    
+    def test_grader_initialization(self):
+        """Test grader initialization with defaults."""
+        grader = BaseLLMGrader()
+        
+        assert grader.model == "gpt-4o-mini"
+        assert grader.temperature == 0.1
+        assert grader.max_tokens == 500
+    
+    def test_grader_custom_model(self):
+        """Test grader with custom model."""
+        grader = BaseLLMGrader(model="gpt-4", temperature=0.5, max_tokens=1000)
+        
+        assert grader.model == "gpt-4"
+        assert grader.temperature == 0.5
+        assert grader.max_tokens == 1000
+    
+    def test_build_prompt_without_expected(self):
+        """Test prompt building without expected output."""
+        grader = BaseLLMGrader()
+        
+        prompt = grader._build_prompt(
+            input_text="What is Python?",
+            output="Python is a programming language",
+        )
+        
+        assert "What is Python?" in prompt
+        assert "Python is a programming language" in prompt
+        assert "EXPECTED OUTPUT" not in prompt
+        assert "SCORE:" in prompt
+        assert "REASONING:" in prompt
+        assert "SUGGESTIONS:" in prompt
+    
+    def test_build_prompt_with_expected(self):
+        """Test prompt building with expected output."""
+        grader = BaseLLMGrader()
+        
+        prompt = grader._build_prompt(
+            input_text="What is Python?",
+            output="Python is a programming language",
+            expected_output="Python is a high-level programming language",
+        )
+        
+        assert "What is Python?" in prompt
+        assert "Python is a programming language" in prompt
+        assert "EXPECTED OUTPUT" in prompt
+        assert "Python is a high-level programming language" in prompt
+    
+    def test_parse_response_valid(self):
+        """Test parsing a valid LLM response."""
+        grader = BaseLLMGrader()
+        
+        response = """SCORE: 8
+REASONING: Good response with accurate information
+SUGGESTIONS:
+- Add more examples
+- Include code snippets"""
+        
+        result = grader._parse_response(
+            response,
+            input_text="test",
+            output="test output",
+            expected_output=None,
+        )
+        
+        assert result.score == 8.0
+        assert "Good response" in result.reasoning
+        assert len(result.suggestions) == 2
+        assert "Add more examples" in result.suggestions
+        assert "Include code snippets" in result.suggestions
+    
+    def test_parse_response_no_suggestions(self):
+        """Test parsing response with no suggestions."""
+        grader = BaseLLMGrader()
+        
+        response = """SCORE: 10
+REASONING: Perfect response
+SUGGESTIONS: None"""
+        
+        result = grader._parse_response(
+            response,
+            input_text="test",
+            output="test output",
+            expected_output=None,
+        )
+        
+        assert result.score == 10.0
+        assert result.suggestions == []
+    
+    def test_parse_response_clamps_score(self):
+        """Test that scores are clamped to 1-10 range."""
+        grader = BaseLLMGrader()
+        
+        # Test score > 10
+        response = """SCORE: 15
+REASONING: Over the top"""
+        
+        result = grader._parse_response(response, "", "", None)
+        assert result.score == 10.0
+        
+        # Test score < 1
+        response = """SCORE: -5
+REASONING: Way too low"""
+        
+        result = grader._parse_response(response, "", "", None)
+        assert result.score == 1.0
+    
+    def test_parse_response_invalid_score(self):
+        """Test parsing response with invalid score."""
+        grader = BaseLLMGrader()
+        
+        response = """SCORE: not a number
+REASONING: Invalid"""
+        
+        result = grader._parse_response(response, "", "", None)
+        assert result.score == 5.0  # Default
+    
+    def test_get_litellm_import_error(self):
+        """Test that missing litellm raises ImportError."""
+        grader = BaseLLMGrader()
+        
+        with patch.dict('sys.modules', {'litellm': None}):
+            with pytest.raises(ImportError) as exc_info:
+                grader._get_litellm()
+            
+            assert "litellm" in str(exc_info.value)
+    
+    @patch('praisonaiagents.eval.grader.BaseLLMGrader._get_litellm')
+    def test_grade_success(self, mock_get_litellm):
+        """Test successful grading."""
+        # Mock litellm
+        mock_litellm = MagicMock()
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message.content = """SCORE: 8
+REASONING: Good response
+SUGGESTIONS: None"""
+        mock_litellm.completion.return_value = mock_response
+        mock_get_litellm.return_value = mock_litellm
+        
+        grader = BaseLLMGrader()
+        result = grader.grade(
+            input_text="What is Python?",
+            output="Python is a programming language",
+        )
+        
+        assert result.score == 8.0
+        assert "Good response" in result.reasoning
+        mock_litellm.completion.assert_called_once()
+    
+    @patch('praisonaiagents.eval.grader.BaseLLMGrader._get_litellm')
+    def test_grade_error_handling(self, mock_get_litellm):
+        """Test grading error handling."""
+        mock_litellm = MagicMock()
+        mock_litellm.completion.side_effect = Exception("API Error")
+        mock_get_litellm.return_value = mock_litellm
+        
+        grader = BaseLLMGrader()
+        result = grader.grade(
+            input_text="test",
+            output="test output",
+        )
+        
+        assert result.score == 5.0
+        assert "Grading error" in result.reasoning
+    
+    @pytest.mark.asyncio
+    @patch('praisonaiagents.eval.grader.BaseLLMGrader._get_litellm')
+    async def test_grade_async_success(self, mock_get_litellm):
+        """Test async grading."""
+        mock_litellm = MagicMock()
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message.content = """SCORE: 9
+REASONING: Excellent
+SUGGESTIONS: None"""
+        
+        # Mock async completion
+        async def mock_acompletion(*args, **kwargs):
+            return mock_response
+        
+        mock_litellm.acompletion = mock_acompletion
+        mock_get_litellm.return_value = mock_litellm
+        
+        grader = BaseLLMGrader()
+        result = await grader.grade_async(
+            input_text="test",
+            output="test output",
+        )
+        
+        assert result.score == 9.0
+        assert "Excellent" in result.reasoning
+
+
+class TestProtocolCompliance:
+    """Test that BaseLLMGrader complies with GraderProtocol."""
+    
+    def test_grader_protocol_compliance(self):
+        """Test that BaseLLMGrader has required protocol attributes."""
+        from praisonaiagents.eval.protocols import GraderProtocol
+        
+        grader = BaseLLMGrader()
+        
+        # Check required attributes
+        assert hasattr(grader, 'model')
+        assert hasattr(grader, 'temperature')
+        assert hasattr(grader, 'grade')
+        
+        # Check it's runtime checkable
+        assert isinstance(grader, GraderProtocol)
+    
+    def test_grade_result_protocol_compliance(self):
+        """Test that GradeResult has required protocol attributes."""
+        from praisonaiagents.eval.protocols import GradeResultProtocol
+        
+        result = GradeResult(score=8.0, reasoning="Good")
+        
+        # Check required attributes
+        assert hasattr(result, 'score')
+        assert hasattr(result, 'reasoning')
+        assert hasattr(result, 'to_dict')
+        
+        # Check it's runtime checkable
+        assert isinstance(result, GradeResultProtocol)

src/praisonai-agents/tests/unit/mcp/__init__.py

@@ -0,0 +1 @@
+# MCP unit tests

src/praisonai-agents/tests/unit/mcp/test_mcp_auth_storage.py

@@ -0,0 +1,203 @@
+"""
+Unit tests for MCP auth storage module.
+
+TDD approach: These tests define the expected behavior of mcp_auth_storage.py
+"""
+import pytest
+import os
+
+
+class TestMCPAuthStorage:
+    """Tests for MCPAuthStorage class."""
+    
+    @pytest.fixture
+    def temp_auth_file(self, tmp_path):
+        """Create a temporary auth file path."""
+        return tmp_path / "mcp-auth.json"
+    
+    @pytest.fixture
+    def auth_storage(self, temp_auth_file):
+        """Create an MCPAuthStorage instance with temp file."""
+        from praisonaiagents.mcp.mcp_auth_storage import MCPAuthStorage
+        return MCPAuthStorage(filepath=str(temp_auth_file))
+    
+    def test_init_creates_empty_storage(self, auth_storage, temp_auth_file):
+        """Storage initializes without creating file until first write."""
+        # File should not exist until we write something
+        assert not temp_auth_file.exists() or temp_auth_file.read_text() == "{}"
+    
+    def test_get_nonexistent_entry_returns_none(self, auth_storage):
+        """Getting a non-existent entry returns None."""
+        result = auth_storage.get("nonexistent")
+        assert result is None
+    
+    def test_set_and_get_tokens(self, auth_storage):
+        """Can set and retrieve tokens for an MCP server."""
+        tokens = {
+            "access_token": "test_access_token",
+            "refresh_token": "test_refresh_token",
+            "expires_at": 1234567890,
+            "scope": "read write"
+        }
+        
+        auth_storage.set_tokens("github", tokens, server_url="https://api.github.com/mcp")
+        
+        entry = auth_storage.get("github")
+        assert entry is not None
+        assert entry["tokens"]["access_token"] == "test_access_token"
+        assert entry["tokens"]["refresh_token"] == "test_refresh_token"
+        assert entry["server_url"] == "https://api.github.com/mcp"
+    
+    def test_set_and_get_client_info(self, auth_storage):
+        """Can set and retrieve client info for dynamic registration."""
+        client_info = {
+            "client_id": "test_client_id",
+            "client_secret": "test_client_secret",
+            "client_id_issued_at": 1234567890,
+            "client_secret_expires_at": 1234567890 + 3600
+        }
+        
+        auth_storage.set_client_info("github", client_info, server_url="https://api.github.com/mcp")
+        
+        entry = auth_storage.get("github")
+        assert entry is not None
+        assert entry["client_info"]["client_id"] == "test_client_id"
+    
+    def test_get_for_url_validates_url(self, auth_storage):
+        """get_for_url returns None if URL has changed."""
+        tokens = {"access_token": "test"}
+        auth_storage.set_tokens("github", tokens, server_url="https://api.github.com/mcp")
+        
+        # Same URL should work
+        entry = auth_storage.get_for_url("github", "https://api.github.com/mcp")
+        assert entry is not None
+        
+        # Different URL should return None
+        entry = auth_storage.get_for_url("github", "https://different.url/mcp")
+        assert entry is None
+    
+    def test_remove_entry(self, auth_storage):
+        """Can remove an entry."""
+        tokens = {"access_token": "test"}
+        auth_storage.set_tokens("github", tokens)
+        
+        auth_storage.remove("github")
+        
+        assert auth_storage.get("github") is None
+    
+    def test_set_code_verifier(self, auth_storage):
+        """Can set and get code verifier for PKCE."""
+        auth_storage.set_code_verifier("github", "test_verifier")
+        
+        entry = auth_storage.get("github")
+        assert entry["code_verifier"] == "test_verifier"
+    
+    def test_clear_code_verifier(self, auth_storage):
+        """Can clear code verifier after use."""
+        auth_storage.set_code_verifier("github", "test_verifier")
+        auth_storage.clear_code_verifier("github")
+        
+        entry = auth_storage.get("github")
+        assert entry is not None
+        assert "code_verifier" not in entry or entry.get("code_verifier") is None
+    
+    def test_set_oauth_state(self, auth_storage):
+        """Can set and get OAuth state for CSRF protection."""
+        auth_storage.set_oauth_state("github", "random_state_123")
+        
+        state = auth_storage.get_oauth_state("github")
+        assert state == "random_state_123"
+    
+    def test_clear_oauth_state(self, auth_storage):
+        """Can clear OAuth state after use."""
+        auth_storage.set_oauth_state("github", "random_state_123")
+        auth_storage.clear_oauth_state("github")
+        
+        state = auth_storage.get_oauth_state("github")
+        assert state is None
+    
+    def test_is_token_expired_no_tokens(self, auth_storage):
+        """is_token_expired returns None if no tokens exist."""
+        result = auth_storage.is_token_expired("nonexistent")
+        assert result is None
+    
+    def test_is_token_expired_no_expiry(self, auth_storage):
+        """is_token_expired returns False if no expiry set."""
+        tokens = {"access_token": "test"}
+        auth_storage.set_tokens("github", tokens)
+        
+        result = auth_storage.is_token_expired("github")
+        assert result is False
+    
+    def test_is_token_expired_not_expired(self, auth_storage):
+        """is_token_expired returns False if token not expired."""
+        import time
+        tokens = {
+            "access_token": "test",
+            "expires_at": time.time() + 3600  # 1 hour from now
+        }
+        auth_storage.set_tokens("github", tokens)
+        
+        result = auth_storage.is_token_expired("github")
+        assert result is False
+    
+    def test_is_token_expired_expired(self, auth_storage):
+        """is_token_expired returns True if token expired."""
+        import time
+        tokens = {
+            "access_token": "test",
+            "expires_at": time.time() - 3600  # 1 hour ago
+        }
+        auth_storage.set_tokens("github", tokens)
+        
+        result = auth_storage.is_token_expired("github")
+        assert result is True
+    
+    def test_file_permissions(self, auth_storage, temp_auth_file):
+        """Auth file should have restricted permissions (0600)."""
+        tokens = {"access_token": "test"}
+        auth_storage.set_tokens("github", tokens)
+        
+        # Check file exists and has correct permissions
+        assert temp_auth_file.exists()
+        # On Unix, check permissions
+        if os.name != 'nt':
+            mode = temp_auth_file.stat().st_mode & 0o777
+            assert mode == 0o600, f"Expected 0600, got {oct(mode)}"
+    
+    def test_all_returns_all_entries(self, auth_storage):
+        """all() returns all stored entries."""
+        auth_storage.set_tokens("github", {"access_token": "gh_token"})
+        auth_storage.set_tokens("tavily", {"access_token": "tv_token"})
+        
+        all_entries = auth_storage.all()
+        assert "github" in all_entries
+        assert "tavily" in all_entries
+        assert len(all_entries) == 2
+    
+    def test_persistence_across_instances(self, temp_auth_file):
+        """Data persists across MCPAuthStorage instances."""
+        from praisonaiagents.mcp.mcp_auth_storage import MCPAuthStorage
+        
+        # First instance writes data
+        storage1 = MCPAuthStorage(filepath=str(temp_auth_file))
+        storage1.set_tokens("github", {"access_token": "test_token"})
+        
+        # Second instance reads data
+        storage2 = MCPAuthStorage(filepath=str(temp_auth_file))
+        entry = storage2.get("github")
+        
+        assert entry is not None
+        assert entry["tokens"]["access_token"] == "test_token"
+
+
+class TestMCPAuthStorageDefaults:
+    """Tests for default file path behavior."""
+    
+    def test_default_filepath(self):
+        """Default filepath is in user's praison config directory."""
+        from praisonaiagents.mcp.mcp_auth_storage import get_default_auth_filepath
+        
+        default_path = get_default_auth_filepath()
+        assert "praison" in str(default_path).lower() or ".praison" in str(default_path)
+        assert default_path.endswith("mcp-auth.json")

src/praisonai-agents/tests/unit/mcp/test_mcp_oauth_callback.py

@@ -0,0 +1,199 @@
+"""
+Unit tests for MCP OAuth callback server module.
+
+TDD approach: These tests define the expected behavior of mcp_oauth_callback.py
+"""
+import pytest
+import asyncio
+import threading
+import time
+
+
+class TestOAuthCallbackServer:
+    """Tests for OAuth callback server."""
+    
+    def test_default_port(self):
+        """Default callback port is 19876 (same as OpenCode for consistency)."""
+        from praisonaiagents.mcp.mcp_oauth_callback import OAUTH_CALLBACK_PORT
+        assert OAUTH_CALLBACK_PORT == 19876
+    
+    def test_default_path(self):
+        """Default callback path is /mcp/oauth/callback."""
+        from praisonaiagents.mcp.mcp_oauth_callback import OAUTH_CALLBACK_PATH
+        assert OAUTH_CALLBACK_PATH == "/mcp/oauth/callback"
+    
+    def test_get_redirect_url(self):
+        """get_redirect_url returns the full callback URL."""
+        from praisonaiagents.mcp.mcp_oauth_callback import get_redirect_url
+        
+        url = get_redirect_url()
+        assert url == "http://127.0.0.1:19876/mcp/oauth/callback"
+    
+    def test_get_redirect_url_custom_port(self):
+        """get_redirect_url supports custom port."""
+        from praisonaiagents.mcp.mcp_oauth_callback import get_redirect_url
+        
+        url = get_redirect_url(port=8080)
+        assert url == "http://127.0.0.1:8080/mcp/oauth/callback"
+
+
+class TestOAuthCallbackHandler:
+    """Tests for OAuthCallbackHandler class."""
+    
+    @pytest.fixture
+    def handler(self):
+        """Create an OAuthCallbackHandler instance."""
+        from praisonaiagents.mcp.mcp_oauth_callback import OAuthCallbackHandler
+        return OAuthCallbackHandler()
+    
+    def test_wait_for_callback_timeout(self, handler):
+        """wait_for_callback raises TimeoutError if no callback received."""
+        with pytest.raises(TimeoutError):
+            handler.wait_for_callback("test_state", timeout=0.1)
+    
+    def test_receive_callback_sets_code(self, handler):
+        """receive_callback stores the authorization code."""
+        handler.receive_callback("test_state", "auth_code_123")
+        
+        # Should be able to get the code now
+        code = handler.get_code("test_state")
+        assert code == "auth_code_123"
+    
+    def test_wait_for_callback_returns_code(self, handler):
+        """wait_for_callback returns code when received."""
+        # Simulate callback in a separate thread
+        def send_callback():
+            time.sleep(0.05)
+            handler.receive_callback("test_state", "auth_code_456")
+        
+        thread = threading.Thread(target=send_callback)
+        thread.start()
+        
+        code = handler.wait_for_callback("test_state", timeout=1.0)
+        assert code == "auth_code_456"
+        
+        thread.join()
+    
+    def test_state_mismatch_returns_none_for_wrong_state(self, handler):
+        """get_code with wrong state returns None."""
+        handler.receive_callback("state_a", "auth_code")
+        
+        # Getting code for wrong state returns None
+        code = handler.get_code("state_b")
+        assert code is None
+        
+        # Getting code for correct state works
+        code = handler.get_code("state_a")
+        assert code == "auth_code"
+    
+    def test_clear_state(self, handler):
+        """clear_state removes pending callback data."""
+        handler.receive_callback("test_state", "auth_code")
+        handler.clear_state("test_state")
+        
+        # Should not find the code anymore
+        code = handler.get_code("test_state")
+        assert code is None
+
+
+class TestOAuthCallbackServerAsync:
+    """Tests for async callback server functionality."""
+    
+    @pytest.fixture
+    def callback_handler(self):
+        """Create an OAuthCallbackHandler instance."""
+        from praisonaiagents.mcp.mcp_oauth_callback import OAuthCallbackHandler
+        return OAuthCallbackHandler()
+    
+    @pytest.mark.asyncio
+    async def test_async_wait_for_callback_timeout(self, callback_handler):
+        """async_wait_for_callback raises TimeoutError if no callback."""
+        with pytest.raises(asyncio.TimeoutError):
+            await callback_handler.async_wait_for_callback("test_state", timeout=0.1)
+    
+    @pytest.mark.asyncio
+    async def test_async_wait_for_callback_returns_code(self, callback_handler):
+        """async_wait_for_callback returns code when received."""
+        async def send_callback():
+            await asyncio.sleep(0.05)
+            callback_handler.receive_callback("test_state", "async_code_789")
+        
+        # Start callback sender
+        task = asyncio.create_task(send_callback())
+        
+        code = await callback_handler.async_wait_for_callback("test_state", timeout=1.0)
+        assert code == "async_code_789"
+        
+        await task
+
+
+class TestGenerateState:
+    """Tests for OAuth state generation."""
+    
+    def test_generate_state_returns_string(self):
+        """generate_state returns a random string."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_state
+        
+        state = generate_state()
+        assert isinstance(state, str)
+        assert len(state) > 0
+    
+    def test_generate_state_is_unique(self):
+        """generate_state returns unique values."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_state
+        
+        states = [generate_state() for _ in range(100)]
+        assert len(set(states)) == 100  # All unique
+
+
+class TestGenerateCodeVerifier:
+    """Tests for PKCE code verifier generation."""
+    
+    def test_generate_code_verifier_returns_string(self):
+        """generate_code_verifier returns a string."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_code_verifier
+        
+        verifier = generate_code_verifier()
+        assert isinstance(verifier, str)
+        assert len(verifier) >= 43  # PKCE spec minimum
+        assert len(verifier) <= 128  # PKCE spec maximum
+    
+    def test_generate_code_verifier_is_url_safe(self):
+        """generate_code_verifier uses URL-safe characters."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_code_verifier
+        import re
+        
+        verifier = generate_code_verifier()
+        # PKCE allows: [A-Z] / [a-z] / [0-9] / "-" / "." / "_" / "~"
+        assert re.match(r'^[A-Za-z0-9\-._~]+$', verifier)
+
+
+class TestGenerateCodeChallenge:
+    """Tests for PKCE code challenge generation."""
+    
+    def test_generate_code_challenge_returns_string(self):
+        """generate_code_challenge returns a string."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_code_challenge
+        
+        challenge = generate_code_challenge("test_verifier")
+        assert isinstance(challenge, str)
+        assert len(challenge) > 0
+    
+    def test_generate_code_challenge_is_deterministic(self):
+        """Same verifier produces same challenge."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_code_challenge
+        
+        verifier = "test_verifier_12345"
+        challenge1 = generate_code_challenge(verifier)
+        challenge2 = generate_code_challenge(verifier)
+        assert challenge1 == challenge2
+    
+    def test_generate_code_challenge_is_base64url(self):
+        """Code challenge is base64url encoded (no padding)."""
+        from praisonaiagents.mcp.mcp_oauth_callback import generate_code_challenge
+        import re
+        
+        challenge = generate_code_challenge("test_verifier")
+        # Base64url: [A-Z] / [a-z] / [0-9] / "-" / "_" (no padding =)
+        assert re.match(r'^[A-Za-z0-9\-_]+$', challenge)
+        assert '=' not in challenge

src/praisonai-agents/tests/unit/storage/test_backends.py

@@ -0,0 +1,333 @@
+"""
+Tests for storage backends (FileBackend, SQLiteBackend).
+
+Tests cover:
+- Basic CRUD operations
+- Thread safety
+- Protocol compliance
+- Backend switching
+"""
+
+import threading
+from pathlib import Path
+
+from praisonaiagents.storage.backends import FileBackend, SQLiteBackend, get_backend
+from praisonaiagents.storage.protocols import StorageBackendProtocol
+
+
+class TestFileBackend:
+    """Tests for FileBackend."""
+    
+    def test_file_backend_creation(self, tmp_path):
+        """Test FileBackend can be created."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        assert backend.storage_dir == tmp_path
+        assert backend.suffix == ".json"
+    
+    def test_file_backend_save_and_load(self, tmp_path):
+        """Test save and load operations."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        
+        data = {"key": "value", "number": 42}
+        backend.save("test_key", data)
+        
+        loaded = backend.load("test_key")
+        assert loaded == data
+    
+    def test_file_backend_exists(self, tmp_path):
+        """Test exists check."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        
+        assert not backend.exists("nonexistent")
+        
+        backend.save("test_key", {"data": "value"})
+        assert backend.exists("test_key")
+    
+    def test_file_backend_delete(self, tmp_path):
+        """Test delete operation."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        
+        backend.save("test_key", {"data": "value"})
+        assert backend.exists("test_key")
+        
+        result = backend.delete("test_key")
+        assert result is True
+        assert not backend.exists("test_key")
+        
+        # Delete nonexistent
+        result = backend.delete("nonexistent")
+        assert result is False
+    
+    def test_file_backend_list_keys(self, tmp_path):
+        """Test listing keys."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        
+        backend.save("key1", {"data": 1})
+        backend.save("key2", {"data": 2})
+        backend.save("other", {"data": 3})
+        
+        all_keys = backend.list_keys()
+        assert len(all_keys) == 3
+        assert "key1" in all_keys
+        assert "key2" in all_keys
+        assert "other" in all_keys
+        
+        # With prefix
+        filtered = backend.list_keys(prefix="key")
+        assert len(filtered) == 2
+        assert "key1" in filtered
+        assert "key2" in filtered
+    
+    def test_file_backend_clear(self, tmp_path):
+        """Test clearing all data."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        
+        backend.save("key1", {"data": 1})
+        backend.save("key2", {"data": 2})
+        
+        count = backend.clear()
+        assert count == 2
+        assert len(backend.list_keys()) == 0
+    
+    def test_file_backend_thread_safety(self, tmp_path):
+        """Test thread-safe operations."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        errors = []
+        
+        def writer(n):
+            try:
+                for i in range(10):
+                    backend.save(f"key_{n}_{i}", {"thread": n, "iteration": i})
+            except Exception as e:
+                errors.append(e)
+        
+        threads = [threading.Thread(target=writer, args=(i,)) for i in range(5)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        
+        assert len(errors) == 0
+        assert len(backend.list_keys()) == 50
+    
+    def test_file_backend_protocol_compliance(self, tmp_path):
+        """Test FileBackend implements StorageBackendProtocol."""
+        backend = FileBackend(storage_dir=str(tmp_path))
+        assert isinstance(backend, StorageBackendProtocol)
+
+
+class TestSQLiteBackend:
+    """Tests for SQLiteBackend."""
+    
+    def test_sqlite_backend_creation(self, tmp_path):
+        """Test SQLiteBackend can be created."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        assert Path(backend.db_path) == db_path
+    
+    def test_sqlite_backend_save_and_load(self, tmp_path):
+        """Test save and load operations."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        data = {"key": "value", "number": 42}
+        backend.save("test_key", data)
+        
+        loaded = backend.load("test_key")
+        assert loaded == data
+    
+    def test_sqlite_backend_exists(self, tmp_path):
+        """Test exists check."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        assert not backend.exists("nonexistent")
+        
+        backend.save("test_key", {"data": "value"})
+        assert backend.exists("test_key")
+    
+    def test_sqlite_backend_delete(self, tmp_path):
+        """Test delete operation."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        backend.save("test_key", {"data": "value"})
+        assert backend.exists("test_key")
+        
+        result = backend.delete("test_key")
+        assert result is True
+        assert not backend.exists("test_key")
+        
+        # Delete nonexistent
+        result = backend.delete("nonexistent")
+        assert result is False
+    
+    def test_sqlite_backend_list_keys(self, tmp_path):
+        """Test listing keys."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        backend.save("key1", {"data": 1})
+        backend.save("key2", {"data": 2})
+        backend.save("other", {"data": 3})
+        
+        all_keys = backend.list_keys()
+        assert len(all_keys) == 3
+        assert "key1" in all_keys
+        assert "key2" in all_keys
+        assert "other" in all_keys
+        
+        # With prefix
+        filtered = backend.list_keys(prefix="key")
+        assert len(filtered) == 2
+        assert "key1" in filtered
+        assert "key2" in filtered
+    
+    def test_sqlite_backend_clear(self, tmp_path):
+        """Test clearing all data."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        backend.save("key1", {"data": 1})
+        backend.save("key2", {"data": 2})
+        
+        count = backend.clear()
+        assert count == 2
+        assert len(backend.list_keys()) == 0
+    
+    def test_sqlite_backend_upsert(self, tmp_path):
+        """Test that save updates existing keys."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        backend.save("key1", {"version": 1})
+        backend.save("key1", {"version": 2})
+        
+        loaded = backend.load("key1")
+        assert loaded["version"] == 2
+        assert len(backend.list_keys()) == 1
+    
+    def test_sqlite_backend_thread_safety(self, tmp_path):
+        """Test thread-safe operations."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        errors = []
+        
+        def writer(n):
+            try:
+                for i in range(10):
+                    backend.save(f"key_{n}_{i}", {"thread": n, "iteration": i})
+            except Exception as e:
+                errors.append(e)
+        
+        threads = [threading.Thread(target=writer, args=(i,)) for i in range(5)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        
+        assert len(errors) == 0
+        assert len(backend.list_keys()) == 50
+    
+    def test_sqlite_backend_protocol_compliance(self, tmp_path):
+        """Test SQLiteBackend implements StorageBackendProtocol."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        assert isinstance(backend, StorageBackendProtocol)
+    
+    def test_sqlite_backend_close(self, tmp_path):
+        """Test closing the backend."""
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        
+        backend.save("key1", {"data": 1})
+        backend.close()
+        
+        # Should be able to reopen
+        backend2 = SQLiteBackend(db_path=str(db_path))
+        loaded = backend2.load("key1")
+        assert loaded["data"] == 1
+
+
+class TestGetBackend:
+    """Tests for get_backend factory function."""
+    
+    def test_get_file_backend(self, tmp_path):
+        """Test getting file backend."""
+        backend = get_backend("file", storage_dir=str(tmp_path))
+        assert isinstance(backend, FileBackend)
+    
+    def test_get_sqlite_backend(self, tmp_path):
+        """Test getting sqlite backend."""
+        db_path = tmp_path / "test.db"
+        backend = get_backend("sqlite", db_path=str(db_path))
+        assert isinstance(backend, SQLiteBackend)
+    
+    def test_get_unknown_backend(self):
+        """Test getting unknown backend raises error."""
+        try:
+            get_backend("unknown")
+            assert False, "Should have raised ValueError"
+        except ValueError as e:
+            assert "Unknown backend type" in str(e)
+
+
+class TestBaseJSONStoreWithBackend:
+    """Tests for BaseJSONStore with different backends."""
+    
+    def test_base_json_store_with_file_backend(self, tmp_path):
+        """Test BaseJSONStore with FileBackend."""
+        from praisonaiagents.storage.base import BaseJSONStore
+        
+        backend = FileBackend(storage_dir=str(tmp_path))
+        store = BaseJSONStore(
+            storage_path=tmp_path / "test.json",
+            backend=backend,
+        )
+        
+        store.save({"items": [1, 2, 3]})
+        loaded = store.load()
+        assert loaded["items"] == [1, 2, 3]
+    
+    def test_base_json_store_with_sqlite_backend(self, tmp_path):
+        """Test BaseJSONStore with SQLiteBackend."""
+        from praisonaiagents.storage.base import BaseJSONStore
+        
+        db_path = tmp_path / "test.db"
+        backend = SQLiteBackend(db_path=str(db_path))
+        store = BaseJSONStore(
+            storage_path=tmp_path / "test.json",
+            backend=backend,
+        )
+        
+        store.save({"items": [1, 2, 3]})
+        loaded = store.load()
+        assert loaded["items"] == [1, 2, 3]
+    
+    def test_base_json_store_backend_switching(self, tmp_path):
+        """Test switching between backends preserves data."""
+        from praisonaiagents.storage.base import BaseJSONStore
+        
+        # Save with file backend
+        file_backend = FileBackend(storage_dir=str(tmp_path / "files"))
+        store1 = BaseJSONStore(
+            storage_path=tmp_path / "test.json",
+            backend=file_backend,
+        )
+        store1.save({"data": "from_file"})
+        
+        # Save with sqlite backend
+        db_path = tmp_path / "test.db"
+        sqlite_backend = SQLiteBackend(db_path=str(db_path))
+        store2 = BaseJSONStore(
+            storage_path=tmp_path / "test.json",
+            backend=sqlite_backend,
+        )
+        store2.save({"data": "from_sqlite"})
+        
+        # Verify each backend has its own data
+        loaded1 = store1.load()
+        loaded2 = store2.load()
+        
+        assert loaded1["data"] == "from_file"
+        assert loaded2["data"] == "from_sqlite"

src/praisonai-agents/tests/unit/storage/test_base.py

@@ -0,0 +1,344 @@
+"""
+Tests for BaseJSONStore and BaseSessionInfo.
+
+TDD: Tests for the DRY storage base implementation.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+from datetime import datetime
+
+from praisonaiagents.storage.base import (
+    BaseJSONStore,
+    FileLock,
+    list_json_sessions,
+    cleanup_old_sessions,
+)
+from praisonaiagents.storage.models import BaseSessionInfo
+
+
+class TestBaseSessionInfo:
+    """Tests for BaseSessionInfo dataclass."""
+    
+    def test_session_info_creation(self):
+        """Test creating a BaseSessionInfo."""
+        info = BaseSessionInfo(
+            session_id="test-123",
+            path=Path("/tmp/test.json"),
+            size_bytes=1024,
+            created_at=datetime(2024, 1, 1),
+            modified_at=datetime(2024, 1, 2),
+            item_count=10,
+        )
+        
+        assert info.session_id == "test-123"
+        assert info.path == Path("/tmp/test.json")
+        assert info.size_bytes == 1024
+        assert info.item_count == 10
+    
+    def test_session_info_defaults(self):
+        """Test BaseSessionInfo default values."""
+        info = BaseSessionInfo(
+            session_id="test",
+            path=Path("/tmp/test.json"),
+        )
+        
+        assert info.size_bytes == 0
+        assert info.item_count == 0
+        assert info.created_at is not None
+        assert info.modified_at is not None
+    
+    def test_session_info_to_dict(self):
+        """Test BaseSessionInfo.to_dict()."""
+        info = BaseSessionInfo(
+            session_id="test-123",
+            path=Path("/tmp/test.json"),
+            size_bytes=1024,
+            item_count=5,
+        )
+        
+        d = info.to_dict()
+        
+        assert d["session_id"] == "test-123"
+        assert d["path"] == "/tmp/test.json"
+        assert d["size_bytes"] == 1024
+        assert d["item_count"] == 5
+        assert "created_at" in d
+        assert "modified_at" in d
+    
+    def test_session_info_from_dict(self):
+        """Test BaseSessionInfo.from_dict()."""
+        data = {
+            "session_id": "test-123",
+            "path": "/tmp/test.json",
+            "size_bytes": 2048,
+            "created_at": "2024-01-01T00:00:00",
+            "modified_at": "2024-01-02T00:00:00",
+            "item_count": 15,
+        }
+        
+        info = BaseSessionInfo.from_dict(data)
+        
+        assert info.session_id == "test-123"
+        assert info.path == Path("/tmp/test.json")
+        assert info.size_bytes == 2048
+        assert info.item_count == 15
+    
+    def test_session_info_from_path(self):
+        """Test BaseSessionInfo.from_path()."""
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            f.write(b'{"test": "data"}')
+            temp_path = Path(f.name)
+        
+        try:
+            info = BaseSessionInfo.from_path(temp_path)
+            
+            assert info.session_id == temp_path.stem
+            assert info.path == temp_path
+            assert info.size_bytes > 0
+        finally:
+            temp_path.unlink()
+
+
+class TestFileLock:
+    """Tests for FileLock."""
+    
+    def test_file_lock_acquire_release(self):
+        """Test acquiring and releasing a file lock."""
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            temp_path = Path(f.name)
+        
+        try:
+            lock_path = Path(str(temp_path) + ".lock")
+            
+            with FileLock(temp_path):
+                # Lock should exist while held
+                assert lock_path.exists()
+            
+            # Lock should be released
+            assert not lock_path.exists()
+        finally:
+            temp_path.unlink()
+    
+    def test_file_lock_context_manager(self):
+        """Test FileLock as context manager."""
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            temp_path = Path(f.name)
+        
+        try:
+            with FileLock(temp_path) as lock:
+                assert lock is not None
+        finally:
+            temp_path.unlink()
+
+
+class TestBaseJSONStore:
+    """Tests for BaseJSONStore."""
+    
+    def test_store_creation(self):
+        """Test creating a BaseJSONStore."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            assert store.storage_path == store_path
+            assert store._data == {}
+    
+    def test_store_save_and_load(self):
+        """Test saving and loading data."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            # Save data
+            store.save({"key": "value", "count": 42})
+            
+            # Create new store and load
+            store2 = BaseJSONStore(store_path)
+            data = store2.load()
+            
+            assert data["key"] == "value"
+            assert data["count"] == 42
+    
+    def test_store_exists(self):
+        """Test exists() method."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            
+            store = BaseJSONStore(store_path, create_if_missing=True)
+            assert not store.exists()  # File not created until save
+            
+            store.save({"test": True})
+            assert store.exists()
+    
+    def test_store_clear(self):
+        """Test clear() method."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            store.save({"key": "value"})
+            store.clear()
+            
+            data = store.load()
+            assert data == {}
+    
+    def test_store_delete(self):
+        """Test delete() method."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            store.save({"key": "value"})
+            assert store.exists()
+            
+            result = store.delete()
+            assert result is True
+            assert not store.exists()
+    
+    def test_store_thread_safety(self):
+        """Test thread-safe operations."""
+        import threading
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            results = []
+            
+            def writer(n):
+                for i in range(10):
+                    with store._lock:
+                        store._data[f"key_{n}_{i}"] = i
+                        store._save()
+                results.append(n)
+            
+            threads = [threading.Thread(target=writer, args=(i,)) for i in range(3)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+            
+            assert len(results) == 3
+    
+    def test_store_custom_default_data(self):
+        """Test subclass with custom default data."""
+        class CustomStore(BaseJSONStore):
+            def _default_data(self):
+                return {"items": [], "version": 1}
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = CustomStore(store_path)
+            
+            assert store._data == {"items": [], "version": 1}
+
+
+class TestListJsonSessions:
+    """Tests for list_json_sessions utility."""
+    
+    def test_list_empty_directory(self):
+        """Test listing empty directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            sessions = list_json_sessions(Path(tmpdir))
+            assert sessions == []
+    
+    def test_list_sessions(self):
+        """Test listing sessions."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
+            
+            # Create some session files
+            for i in range(3):
+                path = tmpdir / f"session_{i}.json"
+                path.write_text(json.dumps({"iterations": [1, 2, 3]}))
+            
+            sessions = list_json_sessions(tmpdir)
+            
+            assert len(sessions) == 3
+            for session in sessions:
+                assert session.session_id.startswith("session_")
+                assert session.item_count == 3
+    
+    def test_list_sessions_limit(self):
+        """Test listing with limit."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
+            
+            for i in range(10):
+                path = tmpdir / f"session_{i}.json"
+                path.write_text("{}")
+            
+            sessions = list_json_sessions(tmpdir, limit=5)
+            
+            assert len(sessions) == 5
+    
+    def test_list_jsonl_sessions(self):
+        """Test listing JSONL sessions."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
+            
+            path = tmpdir / "trace.jsonl"
+            path.write_text('{"event": 1}\n{"event": 2}\n{"event": 3}\n')
+            
+            sessions = list_json_sessions(tmpdir, suffix=".jsonl")
+            
+            assert len(sessions) == 1
+            assert sessions[0].item_count == 3
+
+
+class TestCleanupOldSessions:
+    """Tests for cleanup_old_sessions utility."""
+    
+    def test_cleanup_empty_directory(self):
+        """Test cleanup on empty directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            deleted = cleanup_old_sessions(Path(tmpdir))
+            assert deleted == 0
+    
+    def test_cleanup_nonexistent_directory(self):
+        """Test cleanup on nonexistent directory."""
+        deleted = cleanup_old_sessions(Path("/nonexistent/path"))
+        assert deleted == 0
+
+
+class TestProtocolCompliance:
+    """Test that storage classes comply with protocols."""
+    
+    def test_session_info_protocol_compliance(self):
+        """Test that BaseSessionInfo complies with SessionInfoProtocol."""
+        from praisonaiagents.storage.protocols import SessionInfoProtocol
+        
+        info = BaseSessionInfo(
+            session_id="test",
+            path=Path("/tmp/test.json"),
+        )
+        
+        # Check required attributes
+        assert hasattr(info, 'session_id')
+        assert hasattr(info, 'path')
+        assert hasattr(info, 'size_bytes')
+        assert hasattr(info, 'created_at')
+        assert hasattr(info, 'modified_at')
+        assert hasattr(info, 'to_dict')
+        
+        # Check it's runtime checkable
+        assert isinstance(info, SessionInfoProtocol)
+    
+    def test_json_store_protocol_compliance(self):
+        """Test that BaseJSONStore complies with JSONStoreProtocol."""
+        from praisonaiagents.storage.protocols import JSONStoreProtocol
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            store_path = Path(tmpdir) / "test.json"
+            store = BaseJSONStore(store_path)
+            
+            # Check required attributes
+            assert hasattr(store, 'storage_path')
+            assert hasattr(store, 'load')
+            assert hasattr(store, 'save')
+            assert hasattr(store, 'exists')
+            
+            # Check it's runtime checkable
+            assert isinstance(store, JSONStoreProtocol)

src/praisonai-agents/tests/unit/trace/test_context_events.py

@@ -0,0 +1,518 @@
+"""
+Tests for context events module.
+
+TDD tests for ContextEvent, ContextEventType, and context trace sinks.
+These tests define the expected behavior for the context replay feature.
+"""
+
+import json
+import time
+
+
+class TestContextEventType:
+    """Tests for ContextEventType enum."""
+    
+    def test_context_event_types_exist(self):
+        """Test all required context event types exist."""
+        from praisonaiagents.trace.context_events import ContextEventType
+        
+        required_types = [
+            "SESSION_START",
+            "SESSION_END",
+            "AGENT_START",
+            "AGENT_END",
+            "AGENT_HANDOFF",
+            "MESSAGE_ADDED",
+            "TOOL_CALL_START",
+            "TOOL_CALL_END",
+            "LLM_REQUEST",
+            "LLM_RESPONSE",
+            "CONTEXT_SNAPSHOT",
+        ]
+        
+        for event_type in required_types:
+            assert hasattr(ContextEventType, event_type), f"Missing event type: {event_type}"
+    
+    def test_context_event_type_values(self):
+        """Test event type string values."""
+        from praisonaiagents.trace.context_events import ContextEventType
+        
+        assert ContextEventType.SESSION_START.value == "session_start"
+        assert ContextEventType.AGENT_START.value == "agent_start"
+        assert ContextEventType.MESSAGE_ADDED.value == "message_added"
+        assert ContextEventType.TOOL_CALL_START.value == "tool_call_start"
+
+
+class TestContextEvent:
+    """Tests for ContextEvent dataclass."""
+    
+    def test_context_event_creation_minimal(self):
+        """Test creating ContextEvent with minimal fields."""
+        from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+        
+        event = ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=1736665496.123,
+            session_id="sess-123",
+        )
+        
+        assert event.event_type == ContextEventType.AGENT_START
+        assert event.timestamp == 1736665496.123
+        assert event.session_id == "sess-123"
+        assert event.sequence_num == 0
+        assert event.agent_name is None
+    
+    def test_context_event_creation_full(self):
+        """Test creating ContextEvent with all fields."""
+        from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+        
+        event = ContextEvent(
+            event_type=ContextEventType.MESSAGE_ADDED,
+            timestamp=1736665497.456,
+            session_id="sess-456",
+            agent_name="researcher",
+            sequence_num=5,
+            messages_count=10,
+            tokens_used=1500,
+            tokens_budget=128000,
+            data={"role": "user", "content": "Hello"},
+        )
+        
+        assert event.event_type == ContextEventType.MESSAGE_ADDED
+        assert event.agent_name == "researcher"
+        assert event.sequence_num == 5
+        assert event.messages_count == 10
+        assert event.tokens_used == 1500
+        assert event.tokens_budget == 128000
+        assert event.data == {"role": "user", "content": "Hello"}
+    
+    def test_context_event_to_dict(self):
+        """Test ContextEvent.to_dict() method."""
+        from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+        
+        event = ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=1736665496.123,
+            session_id="sess-123",
+            agent_name="writer",
+        )
+        
+        d = event.to_dict()
+        
+        assert d["event_type"] == "agent_start"
+        assert d["timestamp"] == 1736665496.123
+        assert d["session_id"] == "sess-123"
+        assert d["agent_name"] == "writer"
+        assert "sequence_num" in d
+    
+    def test_context_event_to_json(self):
+        """Test ContextEvent.to_json() method."""
+        from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+        
+        event = ContextEvent(
+            event_type=ContextEventType.TOOL_CALL_START,
+            timestamp=1736665497.0,
+            session_id="sess-789",
+            agent_name="coder",
+            data={"tool_name": "read_file", "args": {"path": "/tmp/test.txt"}},
+        )
+        
+        json_str = event.to_json()
+        parsed = json.loads(json_str)
+        
+        assert parsed["event_type"] == "tool_call_start"
+        assert parsed["session_id"] == "sess-789"
+        assert parsed["data"]["tool_name"] == "read_file"
+    
+    def test_context_event_from_dict(self):
+        """Test ContextEvent.from_dict() class method."""
+        from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+        
+        data = {
+            "event_type": "agent_end",
+            "timestamp": 1736665500.0,
+            "session_id": "sess-abc",
+            "agent_name": "analyst",
+            "sequence_num": 10,
+            "messages_count": 5,
+            "tokens_used": 2000,
+            "tokens_budget": 128000,
+            "data": {"status": "completed"},
+        }
+        
+        event = ContextEvent.from_dict(data)
+        
+        assert event.event_type == ContextEventType.AGENT_END
+        assert event.session_id == "sess-abc"
+        assert event.agent_name == "analyst"
+        assert event.sequence_num == 10
+        assert event.data == {"status": "completed"}
+
+
+class TestContextTraceSink:
+    """Tests for ContextTraceSink protocol."""
+    
+    def test_context_noop_sink(self):
+        """Test ContextNoOpSink does nothing."""
+        from praisonaiagents.trace.context_events import ContextNoOpSink, ContextEvent, ContextEventType
+        
+        sink = ContextNoOpSink()
+        event = ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=time.time(),
+            session_id="test",
+        )
+        
+        # Should not raise
+        sink.emit(event)
+        sink.flush()
+        sink.close()
+    
+    def test_context_list_sink(self):
+        """Test ContextListSink collects events."""
+        from praisonaiagents.trace.context_events import ContextListSink, ContextEvent, ContextEventType
+        
+        sink = ContextListSink()
+        
+        event1 = ContextEvent(
+            event_type=ContextEventType.SESSION_START,
+            timestamp=1.0,
+            session_id="test",
+        )
+        event2 = ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=2.0,
+            session_id="test",
+            agent_name="researcher",
+        )
+        
+        sink.emit(event1)
+        sink.emit(event2)
+        
+        events = sink.get_events()
+        assert len(events) == 2
+        assert events[0].event_type == ContextEventType.SESSION_START
+        assert events[1].event_type == ContextEventType.AGENT_START
+        assert events[1].agent_name == "researcher"
+    
+    def test_context_list_sink_clear(self):
+        """Test ContextListSink.clear() method."""
+        from praisonaiagents.trace.context_events import ContextListSink, ContextEvent, ContextEventType
+        
+        sink = ContextListSink()
+        sink.emit(ContextEvent(
+            event_type=ContextEventType.AGENT_START,
+            timestamp=1.0,
+            session_id="test",
+        ))
+        
+        assert len(sink.get_events()) == 1
+        
+        sink.clear()
+        
+        assert len(sink.get_events()) == 0
+
+
+class TestContextTraceEmitter:
+    """Tests for ContextTraceEmitter class."""
+    
+    def test_emitter_session_lifecycle(self):
+        """Test emitter session start/end events."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.session_start()
+        emitter.session_end()
+        
+        events = sink.get_events()
+        assert len(events) == 2
+        assert events[0].event_type == ContextEventType.SESSION_START
+        assert events[1].event_type == ContextEventType.SESSION_END
+        assert events[0].session_id == "test-session"
+    
+    def test_emitter_agent_lifecycle(self):
+        """Test emitter agent start/end events."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.agent_start("researcher")
+        emitter.agent_end("researcher")
+        
+        events = sink.get_events()
+        assert len(events) == 2
+        assert events[0].event_type == ContextEventType.AGENT_START
+        assert events[0].agent_name == "researcher"
+        assert events[1].event_type == ContextEventType.AGENT_END
+    
+    def test_emitter_message_added(self):
+        """Test emitter message added event."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.message_added(
+            agent_name="writer",
+            role="user",
+            content="Hello, world!",
+            messages_count=1,
+            tokens_used=10,
+        )
+        
+        events = sink.get_events()
+        assert len(events) == 1
+        assert events[0].event_type == ContextEventType.MESSAGE_ADDED
+        assert events[0].agent_name == "writer"
+        assert events[0].messages_count == 1
+        assert events[0].data["role"] == "user"
+        assert events[0].data["content"] == "Hello, world!"
+    
+    def test_emitter_tool_call(self):
+        """Test emitter tool call start/end events."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.tool_call_start(
+            agent_name="coder",
+            tool_name="read_file",
+            tool_args={"path": "/tmp/test.txt"},
+        )
+        emitter.tool_call_end(
+            agent_name="coder",
+            tool_name="read_file",
+            result="file contents here",
+            duration_ms=50.5,
+        )
+        
+        events = sink.get_events()
+        assert len(events) == 2
+        assert events[0].event_type == ContextEventType.TOOL_CALL_START
+        assert events[0].data["tool_name"] == "read_file"
+        assert events[1].event_type == ContextEventType.TOOL_CALL_END
+        assert events[1].data["duration_ms"] == 50.5
+    
+    def test_emitter_llm_request_response(self):
+        """Test emitter LLM request/response events."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.llm_request(
+            agent_name="analyst",
+            messages_count=5,
+            tokens_used=1000,
+            tokens_budget=128000,
+        )
+        emitter.llm_response(
+            agent_name="analyst",
+            response_tokens=500,
+            duration_ms=1500.0,
+        )
+        
+        events = sink.get_events()
+        assert len(events) == 2
+        assert events[0].event_type == ContextEventType.LLM_REQUEST
+        assert events[0].tokens_used == 1000
+        assert events[1].event_type == ContextEventType.LLM_RESPONSE
+        assert events[1].data["response_tokens"] == 500
+    
+    def test_emitter_context_snapshot(self):
+        """Test emitter context snapshot event."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.context_snapshot(
+            agent_name="researcher",
+            messages_count=10,
+            tokens_used=5000,
+            tokens_budget=128000,
+            messages=[
+                {"role": "system", "content": "You are a researcher"},
+                {"role": "user", "content": "Find info about AI"},
+            ],
+        )
+        
+        events = sink.get_events()
+        assert len(events) == 1
+        assert events[0].event_type == ContextEventType.CONTEXT_SNAPSHOT
+        assert events[0].messages_count == 10
+        assert len(events[0].data["messages"]) == 2
+    
+    def test_emitter_disabled(self):
+        """Test emitter when disabled."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test", enabled=False)
+        
+        emitter.session_start()
+        emitter.agent_start("test")
+        emitter.message_added("test", "user", "hello", 1, 10)
+        
+        # No events should be emitted
+        assert len(sink.get_events()) == 0
+    
+    def test_emitter_sequence_numbers(self):
+        """Test that emitter assigns sequential sequence numbers."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.session_start()
+        emitter.agent_start("agent1")
+        emitter.message_added("agent1", "user", "hello", 1, 10)
+        emitter.agent_end("agent1")
+        emitter.session_end()
+        
+        events = sink.get_events()
+        assert len(events) == 5
+        
+        for i, event in enumerate(events):
+            assert event.sequence_num == i, f"Event {i} has wrong sequence_num: {event.sequence_num}"
+    
+    def test_emitter_redacts_sensitive_data(self):
+        """Test that emitter redacts sensitive data in tool args."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test", redact=True)
+        
+        emitter.tool_call_start(
+            agent_name="api_caller",
+            tool_name="call_api",
+            tool_args={"api_key": "sk-secret123", "query": "test"},
+        )
+        
+        events = sink.get_events()
+        assert events[0].data["tool_args"]["api_key"] == "[REDACTED]"
+        assert events[0].data["tool_args"]["query"] == "test"
+    
+    def test_emitter_agent_handoff(self):
+        """Test emitter agent handoff event for tracking agent flow."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink, ContextEventType
+        )
+        
+        sink = ContextListSink()
+        emitter = ContextTraceEmitter(sink=sink, session_id="test-session")
+        
+        emitter.agent_handoff(
+            from_agent="researcher",
+            to_agent="writer",
+            reason="Research complete, passing to writer",
+            context_passed={"findings": ["result1", "result2"]},
+        )
+        
+        events = sink.get_events()
+        assert len(events) == 1
+        assert events[0].event_type == ContextEventType.AGENT_HANDOFF
+        assert events[0].agent_name == "researcher"
+        assert events[0].data["from_agent"] == "researcher"
+        assert events[0].data["to_agent"] == "writer"
+        assert events[0].data["reason"] == "Research complete, passing to writer"
+        assert events[0].data["context_passed"]["findings"] == ["result1", "result2"]
+
+
+class TestPerformance:
+    """Performance tests for context events module."""
+    
+    def test_noop_sink_overhead(self):
+        """Test that ContextNoOpSink has minimal overhead."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextNoOpSink
+        )
+        
+        emitter = ContextTraceEmitter(sink=ContextNoOpSink(), session_id="perf-test")
+        
+        start = time.perf_counter()
+        for _ in range(10000):
+            emitter.message_added("agent", "user", "test message", 1, 100)
+        elapsed = time.perf_counter() - start
+        
+        # Should complete 10000 iterations in under 100ms
+        assert elapsed < 0.1, f"ContextNoOpSink too slow: {elapsed:.3f}s for 10000 iterations"
+    
+    def test_disabled_emitter_overhead(self):
+        """Test that disabled emitter has near-zero overhead."""
+        from praisonaiagents.trace.context_events import (
+            ContextTraceEmitter, ContextListSink
+        )
+        
+        emitter = ContextTraceEmitter(sink=ContextListSink(), session_id="perf-test", enabled=False)
+        
+        start = time.perf_counter()
+        for _ in range(10000):
+            emitter.message_added("agent", "user", "test message", 1, 100)
+        elapsed = time.perf_counter() - start
+        
+        # Should complete 10000 iterations in under 50ms when disabled
+        assert elapsed < 0.05, f"Disabled emitter too slow: {elapsed:.3f}s"
+    
+    def test_import_time(self):
+        """Test that import time is minimal (relative to base package)."""
+        import subprocess
+        import sys
+        
+        # First measure baseline import time for the base package
+        baseline_code = """
+import time
+start = time.perf_counter()
+import praisonaiagents
+elapsed = (time.perf_counter() - start) * 1000
+print(f"{elapsed:.1f}")
+"""
+        subprocess.run(
+            [sys.executable, "-c", baseline_code],
+            capture_output=True,
+            text=True,
+            cwd="/Users/praison/praisonai-package/src/praisonai-agents",
+        )
+        
+        # Now measure context_events import (should add minimal overhead)
+        code = """
+import time
+import praisonaiagents  # Pre-import base
+start = time.perf_counter()
+from praisonaiagents.trace.context_events import ContextEvent, ContextEventType
+elapsed = (time.perf_counter() - start) * 1000
+print(f"{elapsed:.1f}")
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True,
+            cwd="/Users/praison/praisonai-package/src/praisonai-agents",
+        )
+        
+        if result.returncode == 0:
+            import_time_ms = float(result.stdout.strip())
+            # Additional import should be under 50ms (after base is loaded)
+            assert import_time_ms < 50, f"Context events import too slow: {import_time_ms:.1f}ms"

src/praisonai-agents/tests/unit/workflows/test_multi_step_loop.py

@@ -0,0 +1,468 @@
+"""
+Tests for multi-step loop support in workflows.
+
+TDD tests for:
+- Loop class accepting steps parameter (plural)
+- YAML parsing of steps: inside loop
+- Multi-step loop execution (sequential within iteration)
+- previous_output chaining between nested steps
+- Parallel multi-step loop execution
+- Context field aliasing (enabled, max_tool_output_tokens)
+"""
+
+import pytest
+from unittest.mock import MagicMock
+
+
+class TestLoopClassStepsParameter:
+    """Test that Loop class accepts steps parameter."""
+    
+    def test_loop_class_has_steps_field(self):
+        """Loop class should have a steps field for multiple steps."""
+        from praisonaiagents.workflows.workflows import Loop
+        
+        # Create loop with steps (plural)
+        mock_step1 = MagicMock()
+        mock_step2 = MagicMock()
+        
+        loop = Loop(steps=[mock_step1, mock_step2], over="items")
+        
+        assert loop.steps == [mock_step1, mock_step2]
+        assert loop.step is None  # step should be None when steps is used
+        assert loop.over == "items"
+    
+    def test_loop_class_backward_compat_single_step(self):
+        """Loop class should still work with single step parameter."""
+        from praisonaiagents.workflows.workflows import Loop
+        
+        mock_step = MagicMock()
+        loop = Loop(step=mock_step, over="items")
+        
+        assert loop.step == mock_step
+        assert loop.steps is None
+        assert loop.over == "items"
+    
+    def test_loop_class_rejects_both_step_and_steps(self):
+        """Loop class should raise error if both step and steps provided."""
+        from praisonaiagents.workflows.workflows import Loop
+        
+        mock_step = MagicMock()
+        
+        with pytest.raises(ValueError, match="Cannot specify both 'step' and 'steps'"):
+            Loop(step=mock_step, steps=[mock_step], over="items")
+    
+    def test_loop_class_requires_step_or_steps(self):
+        """Loop class should raise error if neither step nor steps provided."""
+        from praisonaiagents.workflows.workflows import Loop
+        
+        with pytest.raises(ValueError, match="Loop requires 'step' or 'steps'"):
+            Loop(over="items")
+
+
+class TestLoopFactoryFunction:
+    """Test that loop() factory function accepts steps parameter."""
+    
+    def test_loop_factory_accepts_steps(self):
+        """loop() factory should accept steps parameter."""
+        from praisonaiagents.workflows.workflows import loop
+        
+        mock_step1 = MagicMock()
+        mock_step2 = MagicMock()
+        
+        result = loop(steps=[mock_step1, mock_step2], over="items", parallel=True)
+        
+        assert result.steps == [mock_step1, mock_step2]
+        assert result.step is None
+        assert result.over == "items"
+        assert result.parallel is True
+    
+    def test_loop_factory_backward_compat(self):
+        """loop() factory should still work with single step."""
+        from praisonaiagents.workflows.workflows import loop
+        
+        mock_step = MagicMock()
+        result = loop(step=mock_step, over="items")
+        
+        assert result.step == mock_step
+        assert result.steps is None
+
+
+class TestYAMLParsingStepsInsideLoop:
+    """Test YAML parsing of steps: inside loop."""
+    
+    def test_yaml_loop_with_steps_at_step_level(self):
+        """YAML parser should parse steps: at step_data level (user's syntax)."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  researcher:
+    role: Researcher
+    goal: Research topics
+  writer:
+    role: Writer
+    goal: Write articles
+
+steps:
+  - loop:
+      over: topics
+      parallel: true
+      max_workers: 4
+    steps:
+      - name: research
+        agent: researcher
+        action: Research {{item}}
+      - name: write
+        agent: writer
+        action: Write about {{previous_output}}
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Find the loop step
+        loop_step = None
+        for step in workflow.steps:
+            from praisonaiagents.workflows.workflows import Loop
+            if isinstance(step, Loop):
+                loop_step = step
+                break
+        
+        assert loop_step is not None, "Loop step should be parsed"
+        assert loop_step.steps is not None, "Loop should have steps (plural)"
+        assert len(loop_step.steps) == 2, "Loop should have 2 nested steps"
+        assert loop_step.parallel is True
+        assert loop_step.max_workers == 4
+    
+    def test_yaml_loop_with_steps_inside_loop_block(self):
+        """YAML parser should parse steps: inside loop: block."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  researcher:
+    role: Researcher
+    goal: Research topics
+  writer:
+    role: Writer
+    goal: Write articles
+
+steps:
+  - loop:
+      over: topics
+      parallel: true
+      steps:
+        - name: research
+          agent: researcher
+        - name: write
+          agent: writer
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Find the loop step
+        loop_step = None
+        for step in workflow.steps:
+            from praisonaiagents.workflows.workflows import Loop
+            if isinstance(step, Loop):
+                loop_step = step
+                break
+        
+        assert loop_step is not None, "Loop step should be parsed"
+        assert loop_step.steps is not None, "Loop should have steps (plural)"
+        assert len(loop_step.steps) == 2, "Loop should have 2 nested steps"
+
+
+class TestMultiStepLoopExecution:
+    """Test multi-step loop execution."""
+    
+    def test_multi_step_loop_executes_all_steps_sequentially(self):
+        """Each iteration should execute all nested steps sequentially."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        
+        # Track execution order
+        execution_order = []
+        
+        def step1_handler(ctx):
+            execution_order.append(f"step1_{ctx.variables.get('loop_index')}")
+            return f"output1_{ctx.variables.get('loop_index')}"
+        
+        def step2_handler(ctx):
+            execution_order.append(f"step2_{ctx.variables.get('loop_index')}")
+            return f"output2_{ctx.variables.get('loop_index')}"
+        
+        step1 = WorkflowStep(name="step1", handler=step1_handler)
+        step2 = WorkflowStep(name="step2", handler=step2_handler)
+        
+        loop_step = Loop(
+            steps=[step1, step2],
+            over="items",
+            parallel=False
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"items": ["a", "b"]}
+        )
+        
+        result = workflow.start("")
+        
+        # Should execute step1 then step2 for each item
+        assert execution_order == [
+            "step1_0", "step2_0",  # First item
+            "step1_1", "step2_1",  # Second item
+        ]
+    
+    def test_previous_output_chains_between_nested_steps(self):
+        """previous_output should flow from step N to step N+1 within iteration."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        
+        received_previous = []
+        
+        def step1_handler(ctx):
+            return "from_step1"
+        
+        def step2_handler(ctx):
+            received_previous.append(ctx.previous_result)
+            return "from_step2"
+        
+        step1 = WorkflowStep(name="step1", handler=step1_handler)
+        step2 = WorkflowStep(name="step2", handler=step2_handler)
+        
+        loop_step = Loop(
+            steps=[step1, step2],
+            over="items",
+            parallel=False
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"items": ["a"]}
+        )
+        
+        workflow.start("")
+        
+        # step2 should receive step1's output
+        assert received_previous == ["from_step1"]
+
+
+class TestParallelMultiStepLoop:
+    """Test parallel execution of multi-step loops."""
+    
+    def test_parallel_multi_step_loop_isolates_iterations(self):
+        """Parallel iterations should be isolated from each other."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        import threading
+        
+        # Track which thread executed which step
+        thread_ids = {}
+        lock = threading.Lock()
+        
+        def step1_handler(ctx):
+            idx = ctx.variables.get('loop_index')
+            with lock:
+                thread_ids[f"step1_{idx}"] = threading.current_thread().ident
+            return f"output1_{idx}"
+        
+        def step2_handler(ctx):
+            idx = ctx.variables.get('loop_index')
+            with lock:
+                thread_ids[f"step2_{idx}"] = threading.current_thread().ident
+            return f"output2_{idx}"
+        
+        step1 = WorkflowStep(name="step1", handler=step1_handler)
+        step2 = WorkflowStep(name="step2", handler=step2_handler)
+        
+        loop_step = Loop(
+            steps=[step1, step2],
+            over="items",
+            parallel=True,
+            max_workers=2
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"items": ["a", "b"]}
+        )
+        
+        workflow.start("")
+        
+        # Within same iteration, steps should run on same thread (sequential)
+        assert thread_ids["step1_0"] == thread_ids["step2_0"]
+        assert thread_ids["step1_1"] == thread_ids["step2_1"]
+
+
+class TestContextFieldAliasing:
+    """Test context field aliasing in YAML parsing."""
+    
+    def test_yaml_context_enabled_maps_to_auto_compact(self):
+        """YAML 'enabled: true' should map to auto_compact=True."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  test:
+    role: Test
+    goal: Test
+
+steps:
+  - agent: test
+
+context:
+  enabled: true
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Context should be enabled
+        assert workflow.context is not None
+        if hasattr(workflow.context, 'auto_compact'):
+            assert workflow.context.auto_compact is True
+        else:
+            # If context is True (simple enable), that's also valid
+            assert workflow.context is True
+    
+    def test_yaml_context_max_tool_output_tokens_mapped(self):
+        """YAML 'max_tool_output_tokens' should map to tool_output_max."""
+        from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
+        
+        yaml_content = """
+agents:
+  test:
+    role: Test
+    goal: Test
+
+steps:
+  - agent: test
+
+context:
+  enabled: true
+  max_tool_output_tokens: 5000
+"""
+        parser = YAMLWorkflowParser()
+        workflow = parser.parse_string(yaml_content)
+        
+        # Context should have tool_output_max set
+        assert workflow.context is not None
+        if hasattr(workflow.context, 'tool_output_max'):
+            assert workflow.context.tool_output_max == 5000
+
+
+class TestParallelLoopItemIsolation:
+    """Test that parallel loop iterations get correct isolated items."""
+    
+    def test_parallel_loop_each_iteration_gets_correct_item(self):
+        """Each parallel iteration should process its own item, not duplicates."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        import threading
+        
+        # Track which item each iteration received
+        received_items = {}
+        lock = threading.Lock()
+        
+        def capture_item_handler(ctx):
+            idx = ctx.variables.get('loop_index')
+            item = ctx.variables.get('item')
+            with lock:
+                received_items[idx] = item
+            return f"processed_{item}"
+        
+        step = WorkflowStep(name="capture", handler=capture_item_handler)
+        
+        loop_step = Loop(
+            steps=[step],
+            over="items",
+            parallel=True,
+            max_workers=3
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"items": ["apple", "banana", "cherry"]}
+        )
+        
+        workflow.start("")
+        
+        # Each iteration should have received a DIFFERENT item
+        assert len(received_items) == 3
+        assert set(received_items.values()) == {"apple", "banana", "cherry"}
+    
+    def test_parallel_loop_dict_items_expanded_correctly(self):
+        """Dict items should have properties expanded for template access."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        import threading
+        
+        received_titles = {}
+        lock = threading.Lock()
+        
+        def capture_title_handler(ctx):
+            idx = ctx.variables.get('loop_index')
+            # Access expanded property
+            title = ctx.variables.get('item.title')
+            with lock:
+                received_titles[idx] = title
+            return f"processed_{title}"
+        
+        step = WorkflowStep(name="capture", handler=capture_title_handler)
+        
+        loop_step = Loop(
+            steps=[step],
+            over="topics",
+            parallel=True,
+            max_workers=3
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"topics": [
+                {"title": "Topic A", "url": "http://a.com"},
+                {"title": "Topic B", "url": "http://b.com"},
+                {"title": "Topic C", "url": "http://c.com"}
+            ]}
+        )
+        
+        workflow.start("")
+        
+        # Each iteration should have received a DIFFERENT title
+        assert len(received_titles) == 3
+        assert set(received_titles.values()) == {"Topic A", "Topic B", "Topic C"}
+
+
+class TestLoopOutputVariable:
+    """Test output_variable stores last step output per iteration."""
+    
+    def test_output_variable_stores_final_step_output(self):
+        """output_variable should store the last step's output for each iteration."""
+        from praisonaiagents.workflows.workflows import Workflow, WorkflowStep, Loop
+        
+        def step1_handler(ctx):
+            return f"step1_output_{ctx.variables.get('loop_index')}"
+        
+        def step2_handler(ctx):
+            return f"final_output_{ctx.variables.get('loop_index')}"
+        
+        step1 = WorkflowStep(name="step1", handler=step1_handler)
+        step2 = WorkflowStep(name="step2", handler=step2_handler)
+        
+        loop_step = Loop(
+            steps=[step1, step2],
+            over="items",
+            output_variable="results"
+        )
+        
+        workflow = Workflow(
+            name="test",
+            steps=[loop_step],
+            variables={"items": ["a", "b"]}
+        )
+        
+        result = workflow.start("")
+        
+        # results should contain final step outputs
+        results = result.get("variables", {}).get("results", [])
+        assert "final_output_0" in results
+        assert "final_output_1" in results

src/praisonai-agents/tests/unit/workflows/test_workflow_consolidated_api.py

@@ -107,25 +107,25 @@ class TestWorkflowOutputPrecedence:
     """Test output= param precedence: Instance > Config > String > Bool > Default."""
     
     def test_output_default(self):
-        """Default output should be verbose=False, stream=True."""
+        """Default output should be verbose=False, stream=False (silent mode, same as Agent)."""
         from praisonaiagents.workflows import Workflow
         w = Workflow(name="test")
         assert w.verbose == False
-        assert w.stream == True
+        assert w.stream == False  # Now uses OutputConfig default (silent mode)
     
     def test_output_bool_true(self):
-        """output=True should enable verbose with defaults."""
+        """output=True enables config with defaults (silent mode, same as Agent)."""
         from praisonaiagents.workflows import Workflow
         w = Workflow(name="test", output=True)
-        # True enables the config with defaults
-        assert w.stream == True
+        # True enables the config with defaults (OutputConfig defaults to silent)
+        assert w.stream == False  # OutputConfig default is stream=False
     
     def test_output_string_verbose(self):
         """output='verbose' should set verbose=True."""
         from praisonaiagents.workflows import Workflow
         w = Workflow(name="test", output="verbose")
         assert w.verbose == True
-        assert w.stream == True
+        assert w.stream == False  # verbose preset has stream=False
     
     def test_output_string_silent(self):
         """output='silent' should set verbose=False, stream=False."""
@@ -135,9 +135,10 @@ class TestWorkflowOutputPrecedence:
         assert w.stream == False
     
     def test_output_config(self):
-        """output=WorkflowOutputConfig should use config values."""
-        from praisonaiagents.workflows import Workflow, WorkflowOutputConfig
-        cfg = WorkflowOutputConfig(verbose=True, stream=False)
+        """output=OutputConfig should use config values (DRY - same as Agent)."""
+        from praisonaiagents.workflows import Workflow
+        from praisonaiagents.config.feature_configs import OutputConfig
+        cfg = OutputConfig(verbose=True, stream=False)
         w = Workflow(name="test", output=cfg)
         assert w.verbose == True
         assert w.stream == False