Release v3.10.4

api.md

@@ -286,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>
@@ -339,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.31" \
+    "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.31" \
+    "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.31" \
+    "praisonai>=3.10.4" \
     "praisonai[ui]" \
     "praisonai[crewai]"
 

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!")

src/praisonai-agents/praisonaiagents/__init__.py

@@ -534,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

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,8 +2841,8 @@ 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:
+            # 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'):
@@ -2920,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."""
@@ -3262,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}")
 
@@ -3349,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
 
@@ -3860,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)
@@ -3917,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]

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/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/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/workflows/__init__.py

@@ -57,8 +57,6 @@ from .workflow_configs import (
     resolve_step_routing_config,
 )
 
-from .yaml_parser import YAMLWorkflowParser
-
 __all__ = [
     # Core
     "Workflow",
@@ -114,3 +112,19 @@ __all__ = [
     "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/workflows.py

@@ -21,6 +21,7 @@ Storage Structure:
 
 import os
 import re
+import json
 import logging
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Callable, Tuple, Union
@@ -34,6 +35,108 @@ from .workflow_configs import (
 logger = logging.getLogger(__name__)
 
 
+def _parse_json_output(output: Any, step_name: str = "step") -> Any:
+    """
+    Parse JSON from LLM output if it's a string.
+    
+    Handles:
+    - Direct JSON strings: '{"key": "value"}'
+    - Markdown code blocks: ```json\n{"key": "value"}\n```
+    - LLM echoing schema: {'type': 'array', 'items': [...]} -> extract items
+    
+    Returns:
+        Parsed dict/list if successful, original output otherwise
+    """
+    if not isinstance(output, str) or not output:
+        # Handle LLM echoing schema structure even for non-string outputs
+        if isinstance(output, dict):
+            return _extract_from_schema_echo(output)
+        return output
+    
+    # Try direct JSON parse first
+    try:
+        parsed = json.loads(output)
+        # Check if LLM echoed the schema structure
+        if isinstance(parsed, dict):
+            parsed = _extract_from_schema_echo(parsed)
+        return parsed
+    except json.JSONDecodeError:
+        pass
+    
+    # Try extracting from markdown code block
+    json_match = re.search(r'```(?:json)?\s*\n?([\s\S]*?)\n?```', output)
+    if json_match:
+        try:
+            parsed = json.loads(json_match.group(1).strip())
+            if isinstance(parsed, dict):
+                parsed = _extract_from_schema_echo(parsed)
+            return parsed
+        except json.JSONDecodeError:
+            pass
+    
+    # Try finding JSON object/array in text
+    # Look for {...} or [...]
+    for pattern in [r'(\{[^}]+\})', r'(\[[^\]]+\])']:
+        match = re.search(pattern, output)
+        if match:
+            try:
+                parsed = json.loads(match.group(1))
+                if isinstance(parsed, dict):
+                    parsed = _extract_from_schema_echo(parsed)
+                return parsed
+            except json.JSONDecodeError:
+                continue
+    
+    # Return original if can't parse
+    logger.debug(f"Could not parse JSON from step '{step_name}' output")
+    return output
+
+
+def _extract_from_schema_echo(data: dict) -> Any:
+    """
+    Extract actual data when LLM echoes the JSON schema structure.
+    
+    Common patterns:
+    - {'type': 'array', 'items': [...]} -> return [...]
+    - {'type': 'object', 'properties': {...}, 'data': {...}} -> return data
+    - {'items': [...]} -> return [...] (native structured output wrapper)
+    - {'keywords': [...]} -> return as-is (normal output)
+    
+    Args:
+        data: Dict that might contain echoed schema
+        
+    Returns:
+        Extracted data or original dict
+    """
+    # Check if this looks like a schema echo with 'type' and actual data
+    if 'type' in data:
+        data_type = data.get('type')
+        
+        # Array schema echo: {'type': 'array', 'items': [...]}
+        if data_type == 'array' and 'items' in data:
+            items = data['items']
+            # If items is a list, that's our actual data
+            if isinstance(items, list):
+                logger.debug("Extracted items from schema-echoed array output")
+                return items
+        
+        # Object schema echo: {'type': 'object', 'properties': {...}, 'data': {...}}
+        if data_type == 'object':
+            # Check for actual data field
+            if 'data' in data and isinstance(data['data'], dict):
+                logger.debug("Extracted data from schema-echoed object output")
+                return data['data']
+    
+    # Handle native structured output wrapper: {'items': [...]}
+    # This is created by _build_response_format when wrapping arrays
+    if 'items' in data and len(data) == 1 and isinstance(data['items'], list):
+        logger.debug("Extracted items from native structured output array wrapper")
+        return data['items']
+    
+    # Not a schema echo, return as-is
+    return data
+
+
 @dataclass
 class WorkflowContext:
     """Context passed to step handlers. Contains all information about the current workflow state."""
@@ -101,10 +204,10 @@ class Parallel:
 @dataclass
 class Loop:
     """
-    Iterate over a list or CSV file, executing step for each item.
+    Iterate over a list or CSV file, executing step(s) for each item.
     
     Usage:
-        # Loop over list variable (sequential)
+        # Loop over list variable (sequential) - single step
         workflow = Workflow(
             steps=[loop(processor, over="items")],
             variables={"items": ["a", "b", "c"]}
@@ -126,32 +229,59 @@ class Loop:
             steps=[loop(processor, over="items", parallel=True, max_workers=4)],
             variables={"items": ["a", "b", "c"]}
         )
+        
+        # Multi-step loop (NEW) - execute multiple steps per iteration
+        workflow = Workflow(
+            steps=[loop(
+                steps=[researcher, writer, publisher],  # Multiple steps
+                over="topics",
+                parallel=True,
+                max_workers=4
+            )],
+            variables={"topics": [...]}
+        )
     """
-    step: Any = None
+    step: Any = None  # Single step (backward compat)
+    steps: Optional[List[Any]] = None  # Multiple steps (NEW)
     over: Optional[str] = None  # Variable name containing list
     from_csv: Optional[str] = None  # CSV file path
     from_file: Optional[str] = None  # Text file path (one item per line)
     var_name: str = "item"  # Variable name for current item
     parallel: bool = False  # Execute iterations in parallel
     max_workers: Optional[int] = None  # Max parallel workers (None = unlimited)
+    output_variable: Optional[str] = None  # Store loop results in this variable name
     
     def __init__(
         self, 
-        step: Any, 
+        step: Any = None,
+        steps: Optional[List[Any]] = None,
         over: Optional[str] = None,
         from_csv: Optional[str] = None,
         from_file: Optional[str] = None,
         var_name: str = "item",
         parallel: bool = False,
-        max_workers: Optional[int] = None
+        max_workers: Optional[int] = None,
+        output_variable: Optional[str] = None
     ):
+        # Validation: cannot have both step and steps
+        if step is not None and steps is not None:
+            raise ValueError("Cannot specify both 'step' and 'steps'")
+        # Validation: must have at least one
+        if step is None and steps is None:
+            raise ValueError("Loop requires 'step' or 'steps'")
+        # Validation: steps cannot be empty
+        if steps is not None and len(steps) == 0:
+            raise ValueError("Loop 'steps' cannot be empty")
+        
         self.step = step
+        self.steps = steps
         self.over = over
         self.from_csv = from_csv
         self.from_file = from_file
         self.var_name = var_name
         self.parallel = parallel
         self.max_workers = max_workers
+        self.output_variable = output_variable
 
 
 @dataclass
@@ -192,25 +322,30 @@ def parallel(steps: List) -> Parallel:
     """Execute steps in parallel."""
     return Parallel(steps=steps)
 
-def loop(step: Any, over: Optional[str] = None, from_csv: Optional[str] = None, 
+def loop(step: Any = None, steps: Optional[List[Any]] = None,
+         over: Optional[str] = None, from_csv: Optional[str] = None, 
          from_file: Optional[str] = None, var_name: str = "item",
-         parallel: bool = False, max_workers: Optional[int] = None) -> Loop:
-    """Loop over items executing step for each.
+         parallel: bool = False, max_workers: Optional[int] = None,
+         output_variable: Optional[str] = None) -> Loop:
+    """Loop over items executing step(s) for each.
     
     Args:
-        step: The step (agent/function) to execute for each item
+        step: Single step (agent/function) to execute for each item (backward compat)
+        steps: Multiple steps to execute sequentially for each item (NEW)
         over: Variable name containing the list to iterate over
         from_csv: Path to CSV file to read items from
         from_file: Path to text file with one item per line
         var_name: Variable name for current item (default: "item")
         parallel: If True, execute iterations in parallel (default: False)
         max_workers: Max parallel workers when parallel=True (default: None = unlimited)
+        output_variable: Variable name to store all loop outputs (default: None = "loop_outputs")
     
     Returns:
         Loop object configured for iteration
     """
-    return Loop(step=step, over=over, from_csv=from_csv, from_file=from_file, 
-                var_name=var_name, parallel=parallel, max_workers=max_workers)
+    return Loop(step=step, steps=steps, over=over, from_csv=from_csv, from_file=from_file, 
+                var_name=var_name, parallel=parallel, max_workers=max_workers,
+                output_variable=output_variable)
 
 
 def repeat(step: Any, until: Optional[Callable[[WorkflowContext], bool]] = None,
@@ -1149,6 +1284,10 @@ class Workflow:
                             
                     elif step.agent:
                         # Direct agent with tools
+                        # Propagate context management to existing agent if workflow has it enabled
+                        if self.context and not step.agent._context_manager_initialized:
+                            step.agent._context_param = self.context
+                        
                         # Substitute variables in action
                         action = step.action or input
                         for key, value in all_variables.items():
@@ -1196,6 +1335,10 @@ class Workflow:
                         
                         output = step.agent.chat(action, **chat_kwargs)
                         
+                        # Parse JSON output if output_json was requested and output is a string
+                        if step_output_json and output and isinstance(output, str):
+                            output = _parse_json_output(output, step.name)
+                        
                         # Handle output_pydantic if present
                         output_pydantic = getattr(step, 'output_pydantic', None)
                         if output_pydantic and output:
@@ -1223,7 +1366,8 @@ class Workflow:
                             tools=step_tools if step_tools else None,
                             output=self.output,  # Propagate output config to child agents
                             reasoning=self.reasoning,
-                            stream=stream
+                            stream=stream,
+                            context=self.context,  # Propagate context management to child agents
                         )
                         # Substitute variables in action
                         action = step.action
@@ -1246,6 +1390,11 @@ class Workflow:
                         
                         output = temp_agent.chat(action, stream=stream)
                         
+                        # Parse JSON output if output_json was requested
+                        step_output_json = getattr(step, '_output_json', None)
+                        if step_output_json and output and isinstance(output, str):
+                            output = _parse_json_output(output, step.name)
+                        
                 except Exception as e:
                     step_error = e
                     output = f"Error: {e}"
@@ -1331,6 +1480,26 @@ class Workflow:
             var_name = step.output_variable or f"{step.name}_output"
             all_variables[var_name] = output
             
+            # Validate output and warn about issues
+            if output is None:
+                logger.warning(f"⚠️  Step '{step.name}': Output is None. Agent may not have returned expected format.")
+                if verbose:
+                    print(f"⚠️  WARNING: Step '{step.name}' output is None!")
+            else:
+                # Check type against output_json schema if defined
+                expected_schema = getattr(step, '_output_json', None)
+                if expected_schema and isinstance(expected_schema, dict):
+                    expected_type = expected_schema.get('type')
+                    actual_type = type(output).__name__
+                    if expected_type == 'object' and not isinstance(output, dict):
+                        logger.warning(f"⚠️  Step '{step.name}': Expected object/dict, got {actual_type}")
+                        if verbose:
+                            print(f"⚠️  Step '{step.name}': Expected 'object', received '{actual_type}'")
+                    elif expected_type == 'array' and not isinstance(output, list):
+                        logger.warning(f"⚠️  Step '{step.name}': Expected array/list, got {actual_type}")
+                        if verbose:
+                            print(f"⚠️  Step '{step.name}': Expected 'array', received '{actual_type}'")
+            
             i += 1
         
         # Update workflow status
@@ -1388,33 +1557,12 @@ class Workflow:
         return await self.astart(input, llm, verbose)
     
     def _normalize_steps(self) -> List['WorkflowStep']:
-        """Convert mixed steps (Agent, function, WorkflowStep) to WorkflowStep list."""
-        normalized = []
+        """Convert mixed steps (Agent, function, WorkflowStep) to WorkflowStep list.
         
-        for i, step in enumerate(self.steps):
-            if isinstance(step, WorkflowStep):
-                normalized.append(step)
-            elif callable(step):
-                # It's a function - wrap as handler
-                normalized.append(WorkflowStep(
-                    name=getattr(step, '__name__', f'step_{i+1}'),
-                    handler=step
-                ))
-            elif hasattr(step, 'chat'):
-                # It's an Agent - wrap with agent reference
-                normalized.append(WorkflowStep(
-                    name=getattr(step, 'name', f'agent_{i+1}'),
-                    agent=step,
-                    action="{{input}}"
-                ))
-            else:
-                # Unknown type - try to use as string action
-                normalized.append(WorkflowStep(
-                    name=f'step_{i+1}',
-                    action=str(step)
-                ))
-        
-        return normalized
+        This method uses _normalize_single_step to ensure consistent normalization
+        and avoid duplicated code paths (DRY principle).
+        """
+        return [self._normalize_single_step(step, i) for i, step in enumerate(self.steps)]
     
     def _create_plan(self, input: str, model: str, verbose: bool) -> Optional[str]:
         """Create an execution plan for the workflow using LLM.
@@ -1581,6 +1729,10 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
                 output = f"Error: {e}"
         elif normalized.agent:
             try:
+                # Propagate context management to existing agent if workflow has it enabled
+                if self.context and not normalized.agent._context_manager_initialized:
+                    normalized.agent._context_param = self.context
+                
                 action = normalized.action or input
                 # Substitute variables
                 for key, value in all_variables.items():
@@ -1593,6 +1745,11 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
                         action = f"{action}\n\nContext from previous step:\n{previous_output}"
                 action = action.replace("{{input}}", input)
                 output = normalized.agent.chat(action, stream=stream)
+                
+                # Parse JSON output if output_json was requested
+                step_output_json = getattr(normalized, '_output_json', None)
+                if step_output_json and output and isinstance(output, str):
+                    output = _parse_json_output(output, normalized.name)
             except Exception as e:
                 output = f"Error: {e}"
         elif normalized.action:
@@ -1604,8 +1761,8 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
                     role=config.get("role", "Assistant"),
                     goal=config.get("goal", "Complete the task"),
                     llm=config.get("llm", model),
-                    # verbose parameter removed - Agent no longer accepts it
-                    stream=stream
+                    stream=stream,
+                    context=self.context,  # Propagate context management to child agents
                 )
                 action = normalized.action
                 for key, value in all_variables.items():
@@ -1618,6 +1775,11 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
                         action = f"{action}\n\nContext from previous step:\n{previous_output}"
                 
                 output = temp_agent.chat(action, stream=stream)
+                
+                # Parse JSON output if output_json was requested
+                step_output_json = getattr(normalized, '_output_json', None)
+                if step_output_json and output and isinstance(output, str):
+                    output = _parse_json_output(output, normalized.name)
             except Exception as e:
                 output = f"Error: {e}"
         
@@ -1769,24 +1931,51 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
         
         num_items = len(items)
         
+        # Determine steps to run (multi-step or single step)
+        steps_to_run = loop_step.steps if loop_step.steps else [loop_step.step]
+        is_multi_step = loop_step.steps is not None and len(loop_step.steps) > 1
+        
         if loop_step.parallel and num_items > 1:
             # Parallel execution
             max_workers = loop_step.max_workers or num_items
             if verbose:
-                print(f"⚡🔁 Parallel looping over {num_items} items (max_workers={max_workers})...")
+                step_info = f" ({len(steps_to_run)} steps each)" if is_multi_step else ""
+                print(f"⚡🔁 Parallel looping over {num_items} items{step_info} (max_workers={max_workers})...")
             
             def execute_item(idx_item_tuple):
-                """Execute step for a single item in thread."""
+                """Execute step(s) for a single item in thread."""
                 idx, item = idx_item_tuple
-                loop_vars = all_variables.copy()
+                # CRITICAL: Deep copy variables to ensure thread isolation
+                import copy
+                loop_vars = copy.deepcopy(all_variables)
                 loop_vars[loop_step.var_name] = item
                 loop_vars["loop_index"] = idx
                 
-                # Execute step (verbose=False in parallel to avoid interleaved output)
-                step_result = self._execute_single_step_internal(
-                    loop_step.step, previous_output, input, loop_vars, model, False, idx, stream=False
-                )
-                return idx, step_result
+                # Also expand nested item properties for template access (e.g., {{item.title}})
+                if isinstance(item, dict):
+                    for key, value in item.items():
+                        loop_vars[f"{loop_step.var_name}.{key}"] = value
+                
+                # Execute all steps sequentially within this iteration
+                iteration_output = previous_output
+                iteration_results = []
+                for step_idx, step in enumerate(steps_to_run):
+                    step_result = self._execute_single_step_internal(
+                        step, iteration_output, input, loop_vars, model, False, step_idx, stream=False
+                    )
+                    iteration_output = step_result.get("output")
+                    iteration_results.append(step_result)
+                    # Update variables if step set any
+                    if step_result.get("variables"):
+                        loop_vars.update(step_result["variables"])
+                
+                # Return final output (last step's output)
+                final_result = {
+                    "step": f"loop_{idx}",
+                    "output": iteration_output,
+                    "steps": iteration_results
+                }
+                return idx, final_result
             
             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                 futures = [executor.submit(execute_item, (idx, item)) for idx, item in enumerate(items)]
@@ -1815,24 +2004,71 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
         else:
             # Sequential execution (original behavior)
             if verbose:
-                print(f"🔁 Looping over {num_items} items...")
+                step_info = f" ({len(steps_to_run)} steps each)" if is_multi_step else ""
+                print(f"🔁 Looping over {num_items} items{step_info}...")
             
             for idx, item in enumerate(items):
                 # Add current item to variables
-                loop_vars = all_variables.copy()
+                import copy
+                loop_vars = copy.deepcopy(all_variables)
                 loop_vars[loop_step.var_name] = item
                 loop_vars["loop_index"] = idx
                 
-                step_result = self._execute_single_step_internal(
-                    loop_step.step, previous_output, input, loop_vars, model, verbose, idx, stream=stream
-                )
-                results.append({"step": f"{step_result['step']}_{idx}", "output": step_result["output"]})
-                outputs.append(step_result["output"])
-                previous_output = step_result["output"]
-        
-        all_variables["loop_outputs"] = outputs
+                # Also expand nested item properties for template access (e.g., {{item.title}})
+                if isinstance(item, dict):
+                    for key, value in item.items():
+                        loop_vars[f"{loop_step.var_name}.{key}"] = value
+                
+                # Execute all steps sequentially within this iteration
+                iteration_output = previous_output
+                for step_idx, step in enumerate(steps_to_run):
+                    step_result = self._execute_single_step_internal(
+                        step, iteration_output, input, loop_vars, model, verbose, step_idx, stream=stream
+                    )
+                    iteration_output = step_result.get("output")
+                    # Update variables if step set any
+                    if step_result.get("variables"):
+                        loop_vars.update(step_result["variables"])
+                
+                results.append({"step": f"loop_{idx}", "output": iteration_output})
+                outputs.append(iteration_output)
+                previous_output = iteration_output
+        # Store outputs in user-specified variable or default to loop_outputs
+        output_var_name = loop_step.output_variable or "loop_outputs"
+        all_variables[output_var_name] = outputs
+        all_variables["loop_outputs"] = outputs  # Also keep for backward compatibility
         combined_output = "\n".join(str(o) for o in outputs) if outputs else ""
         
+        # Validate outputs and warn about issues
+        none_count = sum(1 for o in outputs if o is None)
+        if none_count > 0:
+            logger.warning(f"⚠️  Loop '{output_var_name}': {none_count}/{len(outputs)} outputs are None. "
+                          f"Check if agent returned expected format.")
+            if verbose:
+                print(f"⚠️  WARNING: {none_count}/{len(outputs)} loop outputs are None!")
+        
+        # Check for type consistency
+        if outputs and len(outputs) > 0:
+            expected_type = loop_step.step._output_json if hasattr(loop_step.step, '_output_json') else None
+            if expected_type:
+                expected_schema_type = expected_type.get('type') if isinstance(expected_type, dict) else None
+                for i, o in enumerate(outputs):
+                    if o is None:
+                        continue
+                    actual_type = type(o).__name__
+                    if expected_schema_type == 'object' and not isinstance(o, dict):
+                        logger.warning(f"⚠️  Loop output[{i}]: Expected object/dict, got {actual_type}")
+                        if verbose:
+                            print(f"⚠️  Loop output[{i}]: Expected 'object', received '{actual_type}'")
+                    elif expected_schema_type == 'array' and not isinstance(o, list):
+                        logger.warning(f"⚠️  Loop output[{i}]: Expected array/list, got {actual_type}")
+                        if verbose:
+                            print(f"⚠️  Loop output[{i}]: Expected 'array', received '{actual_type}'")
+        
+        # Debug logging for output_variable
+        if verbose:
+            print(f"📦 Loop stored {len(outputs)} results in variable: '{output_var_name}'")
+        
         return {"steps": results, "output": combined_output, "variables": all_variables}
     
     def _parse_list_from_string(self, text: str) -> List[Any]:
@@ -2054,12 +2290,17 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
             # Determine input for the included recipe
             recipe_input = include_step.input or previous_output or input
             
-            # Substitute variables in the input
-            if recipe_input:
+            # Convert non-string inputs to JSON string for variable substitution
+            if recipe_input and not isinstance(recipe_input, str):
+                recipe_input = json.dumps(recipe_input)
+            
+            # Substitute variables in the input (only if it's a string)
+            if recipe_input and isinstance(recipe_input, str):
                 for key, value in all_variables.items():
                     recipe_input = recipe_input.replace(f"{{{{{key}}}}}", str(value))
                 if previous_output:
-                    recipe_input = recipe_input.replace("{{previous_output}}", str(previous_output))
+                    prev_str = previous_output if isinstance(previous_output, str) else json.dumps(previous_output)
+                    recipe_input = recipe_input.replace("{{previous_output}}", prev_str)
                 recipe_input = recipe_input.replace("{{input}}", input)
             
             # Parse and execute the included workflow

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

@@ -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]:
         """
@@ -222,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
@@ -251,8 +254,21 @@ class YAMLWorkflowParser:
         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', {})
@@ -284,6 +300,39 @@ class YAMLWorkflowParser:
         # 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,
@@ -292,6 +341,7 @@ class YAMLWorkflowParser:
             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
@@ -715,7 +765,7 @@ class YAMLWorkflowParser:
         Returns:
             Loop pattern object
             
-        Supports three syntax forms:
+        Supports multiple syntax forms:
             # Simple form with agent at step level
             - loop:
                 over: items
@@ -734,6 +784,22 @@ class YAMLWorkflowParser:
                 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']
         
@@ -743,9 +809,42 @@ class YAMLWorkflowParser:
         from_file = loop_config.get('from_file')
         var_name = loop_config.get('var_name', 'item')
         parallel_flag = loop_config.get('parallel', False)
-        max_workers = loop_config.get('max_workers')
+        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
         
-        # Resolve the step/agent to execute
+        # 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)
@@ -788,16 +887,20 @@ class YAMLWorkflowParser:
                 )
         
         if agent_or_step is None:
-            raise ValueError("Loop step requires an agent or include")
+            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(
-            agent_or_step, 
+            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
+            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.16"
+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/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/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/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_context.py

@@ -0,0 +1,314 @@
+"""
+Tests for Workflow Context Management.
+
+TDD: These tests define the expected behavior for context propagation
+from Workflow to child agents, parallel loop isolation, and tool output truncation.
+
+Tests cover:
+1. Workflow context param propagates to temp agents
+2. Workflow context param propagates to existing agents  
+3. Parallel loop workers have isolated context
+4. Tool output truncation works in workflow context
+5. Zero overhead when context=False
+"""
+
+import pytest
+from unittest.mock import Mock, patch
+
+
+class TestWorkflowContextPropagation:
+    """Test that workflow context param propagates to child agents."""
+    
+    def test_workflow_context_true_propagates_to_temp_agents(self):
+        """When workflow has context=True, temp agents should receive context=True."""
+        from praisonaiagents import Workflow, WorkflowStep
+        
+        # Create workflow with context enabled
+        workflow = Workflow(
+            name="test_workflow",
+            steps=[
+                WorkflowStep(
+                    name="step1",
+                    action="Test action {{input}}",
+                )
+            ],
+            context=True,  # Enable context management
+        )
+        
+        # Mock the Agent class at the correct import location
+        with patch('praisonaiagents.agent.agent.Agent') as MockAgent:
+            mock_agent_instance = Mock()
+            mock_agent_instance.chat = Mock(return_value="test response")
+            MockAgent.return_value = mock_agent_instance
+            
+            # Run workflow
+            try:
+                workflow.run("test input")
+            except Exception:
+                pass  # May fail due to mocking, but we want to check Agent call
+            
+            # Verify Agent was called with context param
+            if MockAgent.called:
+                call_kwargs = MockAgent.call_args[1] if MockAgent.call_args[1] else {}
+                # The key assertion: context should be passed
+                assert 'context' in call_kwargs, "Agent should receive context param from workflow"
+                assert call_kwargs['context'] is True or call_kwargs['context'] is not False, \
+                    "Agent context should be truthy when workflow context=True"
+    
+    def test_workflow_context_false_no_overhead(self):
+        """When workflow has context=False, no context manager should be created."""
+        from praisonaiagents import Workflow, WorkflowStep
+        
+        workflow = Workflow(
+            name="test_workflow",
+            steps=[
+                WorkflowStep(
+                    name="step1",
+                    action="Test action",
+                )
+            ],
+            context=False,  # Disabled (default)
+        )
+        
+        # Context manager should not be initialized
+        assert workflow._context_manager is None or not workflow._context_manager_initialized, \
+            "Context manager should not be initialized when context=False"
+    
+    def test_workflow_context_propagates_to_existing_agent(self):
+        """When workflow has context=True, existing agents should get context enabled."""
+        from praisonaiagents import Workflow, Agent
+        
+        # Create an agent without context
+        agent = Agent(
+            name="test_agent",
+            instructions="Test instructions",
+            context=False,  # Initially disabled
+        )
+        
+        # Create workflow with context enabled
+        workflow = Workflow(
+            name="test_workflow",
+            steps=[agent],
+            context=True,
+        )
+        
+        # After workflow processes the agent, it should have context enabled
+        # This tests that workflow propagates its context setting to child agents
+        # The implementation should set agent._context_param or agent.context_manager
+        
+        # For now, we just verify the workflow has context enabled
+        assert workflow.context is True, "Workflow should have context=True"
+
+
+class TestParallelLoopContextIsolation:
+    """Test that parallel loop workers have isolated context."""
+    
+    def test_parallel_loop_workers_have_isolated_chat_history(self):
+        """Each parallel worker should have its own isolated chat history."""
+        from praisonaiagents.workflows.workflows import Workflow, Loop, WorkflowStep
+        
+        # Track chat histories from each worker
+        worker_histories = []
+        
+        def capture_history_handler(context):
+            """Handler that captures the chat history state."""
+            # In a real scenario, we'd check the agent's chat_history
+            # For now, we verify the context is isolated
+            worker_histories.append(context.variables.get('loop_index', -1))
+            return f"Result for item {context.variables.get('item', 'unknown')}"
+        
+        workflow = Workflow(
+            name="parallel_test",
+            steps=[
+                Loop(
+                    step=WorkflowStep(
+                        name="worker",
+                        handler=capture_history_handler,
+                    ),
+                    over="items",
+                    parallel=True,
+                    max_workers=4,
+                )
+            ],
+            variables={"items": ["a", "b", "c", "d"]},
+        )
+        
+        workflow.run("")
+        
+        # Verify all 4 items were processed
+        assert len(worker_histories) == 4, "All 4 items should be processed"
+        # Verify each got a unique loop_index (0, 1, 2, 3)
+        assert set(worker_histories) == {0, 1, 2, 3}, "Each worker should have unique loop_index"
+    
+    def test_parallel_workers_dont_share_accumulated_context(self):
+        """Parallel workers should not accumulate context from other workers."""
+        from praisonaiagents.workflows.workflows import Workflow, Loop, WorkflowStep
+        
+        # This test verifies that when using agents in parallel loops,
+        # each agent starts with a clean chat_history, not accumulated from others
+        
+        accumulated_lengths = []
+        
+        def check_history_length(context):
+            """Check that history doesn't grow across parallel workers."""
+            # In isolated execution, each worker should start fresh
+            # We can't directly access agent.chat_history in handler,
+            # but we can verify variables don't leak
+            prev = context.variables.get('_accumulated', 0)
+            accumulated_lengths.append(prev)
+            return "done"
+        
+        workflow = Workflow(
+            name="isolation_test",
+            steps=[
+                Loop(
+                    step=WorkflowStep(
+                        name="checker",
+                        handler=check_history_length,
+                    ),
+                    over="items",
+                    parallel=True,
+                    max_workers=2,
+                )
+            ],
+            variables={"items": [1, 2, 3, 4]},
+        )
+        
+        workflow.run("")
+        
+        # All workers should see 0 accumulated (isolated)
+        assert all(x == 0 for x in accumulated_lengths), \
+            "Parallel workers should not see accumulated state from others"
+
+
+class TestToolOutputTruncation:
+    """Test that tool outputs are truncated when context management is enabled."""
+    
+    def test_tool_output_truncated_with_context_enabled(self):
+        """Tool outputs should be truncated when agent has context=True."""
+        from praisonaiagents import Agent
+        
+        # Create agent with context management
+        agent = Agent(
+            name="test_agent",
+            instructions="Test",
+            context=True,
+        )
+        
+        # Verify context manager exists
+        assert agent.context_manager is not None, \
+            "Agent with context=True should have context_manager"
+        
+        # Test truncation
+        long_output = "x" * 100000  # 100K characters
+        truncated = agent._truncate_tool_output("test_tool", long_output)
+        
+        # Should be truncated (default budget is ~10K tokens ≈ 40K chars)
+        assert len(truncated) < len(long_output), \
+            "Long tool output should be truncated"
+        assert "[truncated]" in truncated.lower() or len(truncated) < 50000, \
+            "Truncated output should indicate truncation or be shorter"
+    
+    def test_tool_output_not_truncated_without_context(self):
+        """Tool outputs should pass through unchanged when context=False."""
+        from praisonaiagents import Agent
+        
+        agent = Agent(
+            name="test_agent",
+            instructions="Test",
+            context=False,  # Disabled
+        )
+        
+        long_output = "x" * 100000
+        result = agent._truncate_tool_output("test_tool", long_output)
+        
+        # Should be unchanged
+        assert result == long_output, \
+            "Tool output should not be modified when context=False"
+
+
+class TestWorkflowContextManagerConfig:
+    """Test workflow context manager configuration."""
+    
+    def test_workflow_accepts_manager_config(self):
+        """Workflow should accept ManagerConfig for detailed configuration."""
+        from praisonaiagents import Workflow, WorkflowStep
+        from praisonaiagents.context import ManagerConfig
+        
+        config = ManagerConfig(
+            auto_compact=True,
+            compact_threshold=0.7,
+            default_tool_output_max=5000,
+        )
+        
+        workflow = Workflow(
+            name="configured_workflow",
+            steps=[
+                WorkflowStep(name="step1", action="test"),
+            ],
+            context=config,
+        )
+        
+        # Workflow should store the config
+        assert workflow.context is config or workflow._context_manager is not None, \
+            "Workflow should accept and store ManagerConfig"
+    
+    def test_workflow_context_string_preset(self):
+        """Workflow should accept string presets for context."""
+        from praisonaiagents import Workflow, WorkflowStep
+        
+        # This tests future functionality where context="auto" enables smart defaults
+        workflow = Workflow(
+            name="preset_workflow",
+            steps=[
+                WorkflowStep(name="step1", action="test"),
+            ],
+            context=True,  # Simple boolean for now
+        )
+        
+        assert workflow.context is True, "Workflow should accept boolean context"
+
+
+class TestZeroOverhead:
+    """Test that context management has zero overhead when disabled."""
+    
+    def test_no_context_manager_import_when_disabled(self):
+        """ContextManager should not be imported when context=False."""
+        from praisonaiagents import Workflow, WorkflowStep
+        
+        workflow = Workflow(
+            name="no_context",
+            steps=[
+                WorkflowStep(name="step1", action="test"),
+            ],
+            context=False,
+        )
+        
+        # _context_manager should be None and not initialized
+        assert workflow._context_manager is None, \
+            "Context manager should be None when context=False"
+        assert not workflow._context_manager_initialized, \
+            "Context manager should not be initialized when context=False"
+    
+    def test_agent_context_lazy_loading(self):
+        """Agent context manager should be lazy loaded."""
+        from praisonaiagents import Agent
+        
+        agent = Agent(
+            name="lazy_agent",
+            instructions="Test",
+            context=True,
+        )
+        
+        # Before first access, should not be initialized
+        assert not agent._context_manager_initialized, \
+            "Context manager should not be initialized until first access"
+        
+        # After access, should be initialized
+        _ = agent.context_manager
+        assert agent._context_manager_initialized, \
+            "Context manager should be initialized after first access"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/praisonai-agents/uv.lock

@@ -2422,7 +2422,7 @@ wheels = [
 
 [[package]]
 name = "praisonaiagents"
-version = "0.12.16"
+version = "0.13.4"
 source = { editable = "." }
 dependencies = [
     { name = "aiohttp" },

src/praisonai/praisonai.rb

@@ -3,8 +3,8 @@ class Praisonai < Formula
   
     desc "AI tools for various AI applications"
     homepage "https://github.com/MervinPraison/PraisonAI"
-    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.9.31.tar.gz"
-    sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.9.31.tar.gz | shasum -a 256`.split.first
+    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.10.4.tar.gz"
+    sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.10.4.tar.gz | shasum -a 256`.split.first
     license "MIT"
   
     depends_on "[email protected]"

src/praisonai/praisonai/cli/app.py

@@ -216,6 +216,7 @@ def register_commands():
     from .commands.test import app as test_app
     from .commands.examples import app as examples_app
     from .commands.batch import app as batch_app
+    from .commands.replay import app as replay_app
     
     # Import TUI and queue commands
     from .features.tui.debug import create_debug_app as create_tui_debug_app
@@ -274,6 +275,7 @@ def register_commands():
     app.add_typer(test_app, name="test", help="Run test suite with tier and provider options")
     app.add_typer(examples_app, name="examples", help="Run and manage example files")
     app.add_typer(batch_app, name="batch", help="Run all PraisonAI scripts in current folder")
+    app.add_typer(replay_app, name="replay", help="Context replay for debugging agent execution")
     
     # Register standardise command
     try:

src/praisonai/praisonai/cli/commands/__init__.py

@@ -25,6 +25,7 @@ __all__ = [
     'rag_app',
     'test_app',
     'examples_app',
+    'replay_app',
 ]
 
 
@@ -84,4 +85,7 @@ def __getattr__(name: str):
     elif name == 'examples_app':
         from .examples import app as examples_app
         return examples_app
+    elif name == 'replay_app':
+        from .replay import app as replay_app
+        return replay_app
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai/praisonai/cli/commands/recipe.py

@@ -33,8 +33,14 @@ def recipe_list():
 def recipe_run(
     name: str = typer.Argument(..., help="Recipe name"),
     model: str = typer.Option(None, "--model", "-m", help="LLM model to use"),
+    var: list[str] = typer.Option(None, "--var", help="Variable override (key=value), can be used multiple times"),
 ):
-    """Run a recipe."""
+    """Run a recipe with optional variable overrides.
+    
+    Examples:
+        praisonai recipe run ai-url-blog-generator --var url="https://example.com/article"
+        praisonai recipe run ai-dynamic-blog-generator --var topic="LangGraph 0.3" --var style="coding"
+    """
     from praisonai.cli.main import PraisonAI
     import sys
     
@@ -42,6 +48,12 @@ def recipe_run(
     if model:
         argv.extend(['--model', model])
     
+    # Pass --var directly to the recipe handler (features/recipe.py)
+    # Don't convert to --workflow-var as that gets consumed by main.py's parser
+    if var:
+        for v in var:
+            argv.extend(['--var', v])
+    
     original_argv = sys.argv
     sys.argv = ['praisonai'] + argv
     
@@ -54,6 +66,8 @@ def recipe_run(
         sys.argv = original_argv
 
 
+
+
 @app.command("install")
 def recipe_install(
     source: str = typer.Argument(..., help="Recipe source (path or URL)"),

src/praisonai/praisonai/cli/commands/replay.py

@@ -0,0 +1,492 @@
+"""
+Replay command group for PraisonAI CLI.
+
+Provides context replay functionality:
+- replay list: List available traces
+- replay context <session_id>: Interactive replay of a trace
+"""
+
+from typing import Optional
+
+import typer
+
+app = typer.Typer(help="Context replay for debugging agent execution")
+
+
+@app.command("list")
+def replay_list(
+    limit: int = typer.Option(
+        20,
+        "--limit",
+        "-n",
+        help="Maximum number of traces to show",
+    ),
+    json_output: bool = typer.Option(
+        False,
+        "--json",
+        "-j",
+        help="Output as JSON",
+    ),
+):
+    """List available context traces."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay import list_traces
+        traces = list_traces(limit=limit)
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    if json_output or output.is_json_mode:
+        output.print_json({
+            "traces": [t.to_dict() for t in traces]
+        })
+        return
+    
+    if not traces:
+        output.print_info("No traces found. Run agents with trace=True to create traces.")
+        return
+    
+    # Format as table
+    headers = ["Session ID", "Events", "Size", "Modified"]
+    rows = []
+    
+    for trace in traces:
+        modified = trace.modified_at.strftime("%Y-%m-%d %H:%M")
+        size_kb = trace.size_bytes / 1024
+        size_str = f"{size_kb:.1f} KB" if size_kb >= 1 else f"{trace.size_bytes} B"
+        
+        rows.append([
+            trace.session_id[:40] + "..." if len(trace.session_id) > 40 else trace.session_id,
+            str(trace.event_count),
+            size_str,
+            modified,
+        ])
+    
+    output.print_table(headers, rows, title="Available Context Traces")
+    output.print("\nUse 'praisonai replay context <session_id>' to replay a trace.")
+
+
+@app.command("context")
+def replay_context(
+    session_id: str = typer.Argument(
+        ...,
+        help="Session ID or path to trace file",
+    ),
+    start: int = typer.Option(
+        0,
+        "--start",
+        "-s",
+        help="Start from event number (1-based)",
+    ),
+    no_rich: bool = typer.Option(
+        False,
+        "--no-rich",
+        help="Disable Rich formatting",
+    ),
+):
+    """Interactive replay of a context trace."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay import ContextTraceReader, ReplayPlayer
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    # Load trace
+    reader = ContextTraceReader(session_id)
+    
+    if not reader.exists:
+        output.print_error(f"Trace not found: {session_id}")
+        output.print_info("Use 'praisonai replay list' to see available traces.")
+        raise typer.Exit(1)
+    
+    events = reader.get_all()
+    if not events:
+        output.print_error(f"No events in trace: {session_id}")
+        raise typer.Exit(1)
+    
+    output.print_success(f"Loaded {len(events)} events from: {session_id}")
+    
+    # Create and run player
+    player = ReplayPlayer(reader, use_rich=not no_rich)
+    
+    # Set start position if specified
+    if start > 0:
+        if player.goto(start - 1):  # Convert to 0-based
+            output.print(f"Starting from event {start}")
+        else:
+            output.print_warning(f"Invalid start position. Valid range: 1-{len(events)}")
+    
+    player.run()
+
+
+@app.command("show")
+def replay_show(
+    session_id: str = typer.Argument(
+        ...,
+        help="Session ID or path to trace file",
+    ),
+    event_num: int = typer.Option(
+        None,
+        "--event",
+        "-e",
+        help="Show specific event number (1-based)",
+    ),
+    agent: Optional[str] = typer.Option(
+        None,
+        "--agent",
+        "-a",
+        help="Filter by agent name",
+    ),
+    event_type: Optional[str] = typer.Option(
+        None,
+        "--type",
+        "-t",
+        help="Filter by event type",
+    ),
+    json_output: bool = typer.Option(
+        False,
+        "--json",
+        "-j",
+        help="Output as JSON",
+    ),
+):
+    """Show trace events without interactive mode."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay import ContextTraceReader
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    reader = ContextTraceReader(session_id)
+    
+    if not reader.exists:
+        output.print_error(f"Trace not found: {session_id}")
+        raise typer.Exit(1)
+    
+    # Get events with filters
+    if event_num is not None:
+        events = [reader[event_num - 1]] if 0 < event_num <= len(reader) else []
+    elif agent:
+        events = reader.get_by_agent(agent)
+    elif event_type:
+        events = reader.get_by_type(event_type)
+    else:
+        events = reader.get_all()
+    
+    if json_output or output.is_json_mode:
+        output.print_json({
+            "session_id": session_id,
+            "total_events": len(reader),
+            "filtered_events": len(events),
+            "events": [e.to_dict() if hasattr(e, 'to_dict') else e for e in events],
+        })
+        return
+    
+    if not events:
+        output.print_info("No events match the filter criteria.")
+        return
+    
+    output.print(f"Session: {session_id}")
+    output.print(f"Showing {len(events)} of {len(reader)} events\n")
+    
+    for i, event in enumerate(events):
+        # Get event details
+        if hasattr(event, 'event_type'):
+            et = event.event_type
+            if hasattr(et, 'value'):
+                et = et.value
+        else:
+            et = event.get('event_type', 'unknown')
+        
+        agent_name = getattr(event, 'agent_name', None) or (event.get('agent_name') if isinstance(event, dict) else None)
+        seq = getattr(event, 'sequence_num', i) or (event.get('sequence_num', i) if isinstance(event, dict) else i)
+        
+        line = f"[{seq + 1}] {et.upper()}"
+        if agent_name:
+            line += f" ({agent_name})"
+        
+        output.print(line)
+
+
+@app.command("delete")
+def replay_delete(
+    session_id: str = typer.Argument(
+        ...,
+        help="Session ID to delete",
+    ),
+    force: bool = typer.Option(
+        False,
+        "--force",
+        "-f",
+        help="Skip confirmation",
+    ),
+):
+    """Delete a context trace."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay import ContextTraceReader
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    reader = ContextTraceReader(session_id)
+    
+    if not reader.exists:
+        output.print_error(f"Trace not found: {session_id}")
+        raise typer.Exit(1)
+    
+    if not force:
+        confirm = typer.confirm(f"Delete trace '{session_id}' ({len(reader)} events)?")
+        if not confirm:
+            output.print("Cancelled.")
+            raise typer.Exit(0)
+    
+    try:
+        reader.path.unlink()
+        output.print_success(f"Deleted trace: {session_id}")
+    except Exception as e:
+        output.print_error(f"Failed to delete: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("flow")
+def replay_flow(
+    session_id: str = typer.Argument(
+        ...,
+        help="Session ID or path to trace file",
+    ),
+    format: str = typer.Option(
+        "ascii",
+        "--format",
+        "-f",
+        help="Output format: ascii, json",
+    ),
+):
+    """Visualize agent flow from a context trace."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay import ContextTraceReader
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    reader = ContextTraceReader(session_id)
+    
+    if not reader.exists:
+        output.print_error(f"Trace not found: {session_id}")
+        raise typer.Exit(1)
+    
+    events = reader.get_all()
+    if not events:
+        output.print_error(f"No events in trace: {session_id}")
+        raise typer.Exit(1)
+    
+    # Build flow structure from events
+    flow_data = _build_flow_from_events(events)
+    
+    if format == "json":
+        output.print_json(flow_data)
+        return
+    
+    # ASCII visualization
+    _print_ascii_flow(output, flow_data, session_id)
+
+
+def _build_flow_from_events(events) -> dict:
+    """Build flow structure from context events."""
+    agents = {}
+    handoffs = []
+    
+    for event in events:
+        event_type = event.event_type
+        if hasattr(event_type, 'value'):
+            event_type = event_type.value
+        
+        agent_name = event.agent_name
+        
+        if event_type == "agent_start" and agent_name:
+            if agent_name not in agents:
+                agents[agent_name] = {
+                    "name": agent_name,
+                    "start_time": event.timestamp,
+                    "end_time": None,
+                    "messages": 0,
+                    "tool_calls": [],
+                    "llm_calls": 0,
+                }
+        
+        elif event_type == "agent_end" and agent_name:
+            if agent_name in agents:
+                agents[agent_name]["end_time"] = event.timestamp
+        
+        elif event_type == "agent_handoff":
+            handoffs.append({
+                "from": event.data.get("from_agent"),
+                "to": event.data.get("to_agent"),
+                "reason": event.data.get("reason"),
+                "timestamp": event.timestamp,
+            })
+        
+        elif event_type == "message_added" and agent_name:
+            if agent_name in agents:
+                agents[agent_name]["messages"] += 1
+        
+        elif event_type == "tool_call_start" and agent_name:
+            tool_name = event.data.get("tool_name", "unknown")
+            if agent_name in agents:
+                agents[agent_name]["tool_calls"].append(tool_name)
+        
+        elif event_type == "llm_request" and agent_name:
+            if agent_name in agents:
+                agents[agent_name]["llm_calls"] += 1
+    
+    # Determine agent order based on start times
+    agent_order = sorted(
+        agents.keys(),
+        key=lambda a: agents[a].get("start_time", 0) or 0
+    )
+    
+    return {
+        "session_id": events[0].session_id if events else "",
+        "agents": agents,
+        "agent_order": agent_order,
+        "handoffs": handoffs,
+        "total_events": len(events),
+    }
+
+
+def _print_ascii_flow(output, flow_data: dict, session_id: str):
+    """Print ASCII flow visualization."""
+    agents = flow_data["agents"]
+    agent_order = flow_data["agent_order"]
+    handoffs = flow_data["handoffs"]
+    
+    output.print(f"\n{'='*60}")
+    output.print(f"  AGENT FLOW: {session_id}")
+    output.print(f"{'='*60}\n")
+    
+    if not agent_order:
+        output.print("  No agents found in trace.")
+        return
+    
+    # Print agent boxes with connections
+    for i, agent_name in enumerate(agent_order):
+        agent = agents[agent_name]
+        
+        # Agent box
+        box_width = 40
+        output.print(f"  ┌{'─' * box_width}┐")
+        output.print(f"  │ {'Agent: ' + agent_name:<{box_width-2}} │")
+        output.print(f"  │{' ' * box_width}│")
+        output.print(f"  │   Messages: {agent['messages']:<{box_width-15}} │")
+        output.print(f"  │   LLM Calls: {agent['llm_calls']:<{box_width-16}} │")
+        
+        if agent['tool_calls']:
+            tools_str = ", ".join(agent['tool_calls'][:3])
+            if len(agent['tool_calls']) > 3:
+                tools_str += f" (+{len(agent['tool_calls'])-3} more)"
+            output.print(f"  │   Tools: {tools_str:<{box_width-13}} │")
+        
+        output.print(f"  └{'─' * box_width}┘")
+        
+        # Connection to next agent
+        if i < len(agent_order) - 1:
+            next_agent = agent_order[i + 1]
+            
+            # Find handoff reason
+            handoff_reason = None
+            for h in handoffs:
+                if h["from"] == agent_name and h["to"] == next_agent:
+                    handoff_reason = h.get("reason")
+                    break
+            
+            output.print("          │")
+            output.print("          ▼")
+            if handoff_reason:
+                reason_short = handoff_reason[:35] + "..." if len(handoff_reason) > 35 else handoff_reason
+                output.print(f"    [{reason_short}]")
+            output.print("")
+    
+    # Summary
+    output.print("\n" + "─"*60)
+    output.print("  Summary:")
+    output.print(f"    Agents: {len(agent_order)}")
+    output.print(f"    Handoffs: {len(handoffs)}")
+    output.print(f"    Total Events: {flow_data['total_events']}")
+    output.print("─"*60 + "\n")
+
+
+@app.command("cleanup")
+def replay_cleanup(
+    max_age_days: int = typer.Option(
+        7,
+        "--max-age",
+        help="Delete traces older than N days",
+    ),
+    max_size_mb: int = typer.Option(
+        100,
+        "--max-size",
+        help="Delete oldest traces if total size exceeds N MB",
+    ),
+    dry_run: bool = typer.Option(
+        False,
+        "--dry-run",
+        help="Show what would be deleted without deleting",
+    ),
+):
+    """Clean up old context traces."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.replay.storage import cleanup_old_traces, list_traces
+    except ImportError:
+        output.print_error("Replay module not available")
+        raise typer.Exit(1)
+    
+    if dry_run:
+        from datetime import datetime
+        traces = list_traces(limit=10000)
+        now = datetime.now()
+        
+        to_delete = []
+        for trace in traces:
+            age_days = (now - trace.modified_at).days
+            if age_days > max_age_days:
+                to_delete.append(trace)
+        
+        if to_delete:
+            output.print(f"Would delete {len(to_delete)} traces older than {max_age_days} days:")
+            for t in to_delete[:10]:
+                output.print(f"  - {t.session_id}")
+            if len(to_delete) > 10:
+                output.print(f"  ... and {len(to_delete) - 10} more")
+        else:
+            output.print("No traces would be deleted.")
+        return
+    
+    deleted = cleanup_old_traces(max_age_days=max_age_days, max_size_mb=max_size_mb)
+    
+    if deleted > 0:
+        output.print_success(f"Deleted {deleted} old traces.")
+    else:
+        output.print_info("No traces needed cleanup.")

src/praisonai/praisonai/cli/commands/session.py

@@ -13,11 +13,30 @@ from typing import Optional
 import typer
 
 from ..output.console import get_output_controller
-from ..state.sessions import get_session_manager
+from ..state.sessions import get_session_manager, set_session_backend
 
 app = typer.Typer(help="Session management")
 
 
+def _create_backend(backend_type: str, storage_path: Optional[str]):
+    """Create storage backend from CLI options."""
+    try:
+        if backend_type == "file":
+            from praisonaiagents.storage import FileBackend
+            return FileBackend(storage_dir=storage_path or "~/.praison/sessions")
+        elif backend_type == "sqlite":
+            from praisonaiagents.storage import SQLiteBackend
+            db_path = storage_path or "~/.praison/sessions.db"
+            return SQLiteBackend(db_path=db_path)
+        elif backend_type.startswith("redis://"):
+            from praisonaiagents.storage import RedisBackend
+            return RedisBackend(url=backend_type, prefix="session:")
+        else:
+            return None
+    except Exception:
+        return None
+
+
 @app.command("list")
 def session_list(
     limit: int = typer.Option(
@@ -26,9 +45,26 @@ def session_list(
         "-n",
         help="Maximum number of sessions to show",
     ),
+    storage_backend: Optional[str] = typer.Option(
+        None,
+        "--storage-backend",
+        help="Storage backend: 'file', 'sqlite', or 'redis://url'",
+    ),
+    storage_path: Optional[str] = typer.Option(
+        None,
+        "--storage-path",
+        help="Path for storage backend",
+    ),
 ):
     """List all sessions."""
     output = get_output_controller()
+    
+    # Configure backend if specified
+    if storage_backend:
+        backend = _create_backend(storage_backend, storage_path)
+        if backend:
+            set_session_backend(backend)
+    
     manager = get_session_manager()
     
     sessions = manager.list(limit=limit)

src/praisonai/praisonai/cli/commands/train.py

@@ -1,36 +1,52 @@
 """
 Train command group for PraisonAI CLI.
 
-Provides model training commands.
+Provides two types of training:
+1. LLM Fine-tuning: praisonai train dataset.json
+2. Agent Training: praisonai train agents [OPTIONS]
 """
 
+import json
+from pathlib import Path
 from typing import Optional
 
 import typer
 
-app = typer.Typer(help="Model training and fine-tuning")
+app = typer.Typer(help="Training commands for models and agents")
 
 
-@app.callback(invoke_without_command=True)
-def train_main(
-    ctx: typer.Context,
-    dataset: Optional[str] = typer.Argument(None, help="Training dataset path"),
+@app.callback()
+def train_callback():
+    """
+    Training commands for models and agents.
+    
+    For LLM fine-tuning:
+        praisonai train llm dataset.json
+    
+    For agent training (iterative improvement):
+        praisonai train agents --input "What is Python?"
+        praisonai train agents --input "Explain AI" --human
+    """
+    pass
+
+
+@app.command("llm")
+def train_llm(
+    dataset: str = typer.Argument(..., help="Training dataset path"),
     model: Optional[str] = typer.Option(None, "--model", "-m", help="Base model to fine-tune"),
     verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
 ):
     """
-    Train or fine-tune models.
+    Fine-tune LLM models using Unsloth.
     
     Examples:
-        praisonai train dataset.json
-        praisonai train --model gpt-4o dataset.json
+        praisonai train llm dataset.json
+        praisonai train llm --model llama-3.1 dataset.json
     """
     from praisonai.cli.main import PraisonAI
     import sys
     
-    argv = ['train']
-    if dataset:
-        argv.append(dataset)
+    argv = ['train', dataset]
     if model:
         argv.extend(['--model', model])
     if verbose:
@@ -46,3 +62,376 @@ def train_main(
         pass
     finally:
         sys.argv = original_argv
+
+
+@app.command("agents")
+def train_agents(
+    agent_file: Optional[str] = typer.Argument(
+        None,
+        help="Path to agent YAML file or Python file"
+    ),
+    iterations: int = typer.Option(
+        3,
+        "--iterations", "-n",
+        help="Number of training iterations"
+    ),
+    human: bool = typer.Option(
+        False,
+        "--human", "-h",
+        help="Use human feedback instead of LLM grading"
+    ),
+    scenarios: Optional[str] = typer.Option(
+        None,
+        "--scenarios", "-s",
+        help="Path to scenarios JSON file"
+    ),
+    input_text: Optional[str] = typer.Option(
+        None,
+        "--input", "-i",
+        help="Single input text for training (alternative to scenarios file)"
+    ),
+    expected: Optional[str] = typer.Option(
+        None,
+        "--expected", "-e",
+        help="Expected output for the input (optional)"
+    ),
+    output_dir: Optional[str] = typer.Option(
+        None,
+        "--output", "-o",
+        help="Output directory for training data"
+    ),
+    model: Optional[str] = typer.Option(
+        None,
+        "--model", "-m",
+        help="LLM model for grading (default: gpt-4o-mini)"
+    ),
+    verbose: bool = typer.Option(
+        True,
+        "--verbose/--quiet", "-v/-q",
+        help="Show detailed progress"
+    ),
+    dry_run: bool = typer.Option(
+        False,
+        "--dry-run",
+        help="Show what would happen without running"
+    ),
+    storage_backend: Optional[str] = typer.Option(
+        None,
+        "--storage-backend",
+        help="Storage backend: 'file', 'sqlite', or 'redis://url'. Default: file"
+    ),
+    storage_path: Optional[str] = typer.Option(
+        None,
+        "--storage-path",
+        help="Path for storage backend (file dir or sqlite db path)"
+    ),
+):
+    """
+    Train agents through iterative feedback loops.
+    
+    Two modes available:
+    
+    1. LLM-as-Judge (default): Agent runs, LLM grades output, 
+       provides suggestions, agent re-runs with feedback.
+       Fully automated, no human intervention needed.
+    
+    2. Human-in-the-Loop (--human): Agent runs, you review
+       and provide feedback, agent re-runs with your feedback.
+    
+    Examples:
+        # Simple: Train with a single input
+        praisonai train agents --input "What is Python?"
+        
+        # With expected output
+        praisonai train agents --input "What is 2+2?" --expected "4"
+        
+        # With scenarios file
+        praisonai train agents --scenarios scenarios.json
+        
+        # Human feedback mode
+        praisonai train agents --input "Explain AI" --human
+        
+        # More iterations
+        praisonai train agents --input "Hello" --iterations 5
+        
+        # With agent file
+        praisonai train agents my_agent.yaml --scenarios scenarios.json
+    """
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    # Validate inputs
+    if not input_text and not scenarios:
+        output.print_error(
+            "No training input provided",
+            remediation="Use --input 'text' or --scenarios file.json"
+        )
+        raise typer.Exit(1)
+    
+    if dry_run:
+        output.print_panel(
+            f"Mode: {'Human-in-the-Loop' if human else 'LLM-as-Judge'}\n"
+            f"Iterations: {iterations}\n"
+            f"Input: {input_text or 'from scenarios file'}\n"
+            f"Scenarios file: {scenarios or 'N/A'}\n"
+            f"Agent file: {agent_file or 'default agent'}\n"
+            f"Output dir: {output_dir or '~/.praison/train/'}",
+            title="Dry Run - Would execute:"
+        )
+        return
+    
+    # Lazy import to avoid overhead
+    try:
+        from praisonai.train.agents import AgentTrainer, TrainingScenario
+    except ImportError as e:
+        output.print_error(f"Failed to import training module: {e}")
+        raise typer.Exit(1)
+    
+    # Create or load agent
+    agent = None
+    if agent_file:
+        agent = _load_agent_from_file(agent_file, output)
+        if agent is None:
+            raise typer.Exit(1)
+    else:
+        # Create a simple default agent
+        try:
+            from praisonaiagents import Agent
+            agent = Agent(
+                instructions="You are a helpful assistant. Provide clear, accurate, and helpful responses."
+            )
+        except ImportError:
+            output.print_error(
+                "praisonaiagents not installed",
+                remediation="pip install praisonaiagents"
+            )
+            raise typer.Exit(1)
+    
+    # Create trainer
+    storage_dir = Path(output_dir) if output_dir else None
+    
+    from praisonai.train.agents.grader import TrainingGrader
+    grader = TrainingGrader(model=model) if model else None
+    
+    # Create storage backend if specified
+    backend = None
+    if storage_backend:
+        backend = _create_storage_backend(storage_backend, storage_path, output)
+        if backend is None and storage_backend not in ("file", None):
+            raise typer.Exit(1)
+    
+    trainer = AgentTrainer(
+        agent=agent,
+        iterations=iterations,
+        human_mode=human,
+        grader=grader,
+        storage_dir=storage_dir,
+        storage_backend=backend,
+        verbose=verbose,
+    )
+    
+    # Add scenarios
+    if scenarios:
+        scenarios_list = _load_scenarios_from_file(scenarios, output)
+        if scenarios_list is None:
+            raise typer.Exit(1)
+        trainer.add_scenarios(scenarios_list)
+    elif input_text:
+        trainer.add_scenario(TrainingScenario(
+            id="cli-input",
+            input_text=input_text,
+            expected_output=expected,
+        ))
+    
+    # Run training
+    try:
+        report = trainer.run()
+        
+        if not verbose:
+            report.print_summary()
+        
+        output.print_success(
+            f"Training complete! Session: {report.session_id}\n"
+            f"Average score: {report.avg_score:.1f}/10\n"
+            f"Improvement: {report.improvement:+.1f}"
+        )
+        
+    except KeyboardInterrupt:
+        output.print_warning("Training interrupted by user")
+        raise typer.Exit(130)
+    except Exception as e:
+        output.print_error(f"Training failed: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("list")
+def train_list(
+    limit: int = typer.Option(20, "--limit", "-n", help="Max sessions to show"),
+    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
+):
+    """List all training sessions."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.train.agents.storage import list_training_sessions
+    except ImportError:
+        output.print_error("Training module not available")
+        raise typer.Exit(1)
+    
+    sessions = list_training_sessions(limit=limit)
+    
+    if json_output or output.is_json_mode:
+        output.print_json({
+            "sessions": [s.to_dict() for s in sessions]
+        })
+        return
+    
+    if not sessions:
+        output.print_info("No training sessions found")
+        return
+    
+    headers = ["Session ID", "Iterations", "Size", "Modified"]
+    rows = []
+    
+    for s in sessions:
+        modified = s.modified_at.strftime("%Y-%m-%d %H:%M")
+        size_kb = s.size_bytes / 1024
+        size_str = f"{size_kb:.1f} KB" if size_kb >= 1 else f"{s.size_bytes} B"
+        
+        rows.append([
+            s.session_id[:30] + "..." if len(s.session_id) > 30 else s.session_id,
+            str(s.iteration_count),
+            size_str,
+            modified,
+        ])
+    
+    output.print_table(headers, rows, title="Training Sessions")
+
+
+@app.command("show")
+def train_show(
+    session_id: str = typer.Argument(..., help="Session ID to show"),
+    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
+):
+    """Show details of a training session."""
+    from ..output.console import get_output_controller
+    
+    output = get_output_controller()
+    
+    try:
+        from praisonai.train.agents.storage import TrainingStorage
+    except ImportError:
+        output.print_error("Training module not available")
+        raise typer.Exit(1)
+    
+    storage = TrainingStorage(session_id=session_id)
+    
+    if not storage.storage_path.exists():
+        output.print_error(f"Session not found: {session_id}")
+        raise typer.Exit(1)
+    
+    report = storage.load_report()
+    iterations = storage.load_iterations()
+    
+    if json_output or output.is_json_mode:
+        output.print_json({
+            "session_id": session_id,
+            "report": report.to_dict() if report else None,
+            "iterations": [it.to_dict() for it in iterations],
+        })
+        return
+    
+    if report:
+        report.print_summary()
+    
+    if iterations:
+        output.print("\nIterations:")
+        for it in iterations:
+            output.print(f"  [{it.iteration_num}] Score: {it.score}/10 - {it.feedback[:50]}...")
+
+
+def _load_agent_from_file(file_path: str, output) -> Optional[object]:
+    """Load agent from YAML or Python file."""
+    path = Path(file_path)
+    
+    if not path.exists():
+        output.print_error(f"File not found: {file_path}")
+        return None
+    
+    if path.suffix in ('.yaml', '.yml'):
+        try:
+            from praisonai.cli.main import PraisonAI
+            praison = PraisonAI(agent_file=str(path))
+            # Return the agents object
+            if hasattr(praison, 'agents') and praison.agents:
+                return praison.agents
+            output.print_error("Could not load agents from YAML")
+            return None
+        except Exception as e:
+            output.print_error(f"Failed to load YAML: {e}")
+            return None
+    
+    elif path.suffix == '.py':
+        output.print_warning(
+            "Python file support coming soon. "
+            "For now, use YAML files or --input flag."
+        )
+        return None
+    
+    else:
+        output.print_error(f"Unsupported file type: {path.suffix}")
+        return None
+
+
+def _create_storage_backend(backend_type: str, storage_path: Optional[str], output):
+    """Create storage backend from CLI options."""
+    try:
+        if backend_type == "file":
+            from praisonaiagents.storage import FileBackend
+            return FileBackend(storage_dir=storage_path or "~/.praison/train")
+        elif backend_type == "sqlite":
+            from praisonaiagents.storage import SQLiteBackend
+            db_path = storage_path or "~/.praison/train.db"
+            return SQLiteBackend(db_path=db_path)
+        elif backend_type.startswith("redis://"):
+            from praisonaiagents.storage import RedisBackend
+            return RedisBackend(url=backend_type, prefix="train:")
+        else:
+            output.print_error(
+                f"Unknown storage backend: {backend_type}",
+                remediation="Use 'file', 'sqlite', or 'redis://url'"
+            )
+            return None
+    except ImportError as e:
+        output.print_error(f"Failed to import storage backend: {e}")
+        return None
+    except Exception as e:
+        output.print_error(f"Failed to create storage backend: {e}")
+        return None
+
+
+def _load_scenarios_from_file(file_path: str, output) -> Optional[list]:
+    """Load scenarios from JSON file."""
+    path = Path(file_path)
+    
+    if not path.exists():
+        output.print_error(f"Scenarios file not found: {file_path}")
+        return None
+    
+    try:
+        with open(path) as f:
+            data = json.load(f)
+        
+        if isinstance(data, list):
+            return data
+        elif isinstance(data, dict) and "scenarios" in data:
+            return data["scenarios"]
+        else:
+            output.print_error("Invalid scenarios format. Expected list or {scenarios: [...]}")
+            return None
+    except json.JSONDecodeError as e:
+        output.print_error(f"Invalid JSON: {e}")
+        return None

src/praisonai/praisonai/cli/features/recipe.py

@@ -215,6 +215,10 @@ class RecipeHandler:
     def _parse_args(self, args: List[str], spec: Dict[str, Any]) -> Dict[str, Any]:
         """Parse command arguments based on spec."""
         result = {k: v.get("default") for k, v in spec.items()}
+        # Initialize multi-value options as empty lists
+        for k, v in spec.items():
+            if v.get("multi"):
+                result[k] = []
         positional_keys = [k for k, v in spec.items() if v.get("positional")]
         positional_idx = 0
         
@@ -228,7 +232,11 @@ class RecipeHandler:
                     if spec[key].get("flag"):
                         result[key] = True
                     elif i + 1 < len(args):
-                        result[key] = args[i + 1]
+                        # Support multi-value options (--var can be used multiple times)
+                        if spec[key].get("multi"):
+                            result[key].append(args[i + 1])
+                        else:
+                            result[key] = args[i + 1]
                         i += 1
                 i += 1
             elif arg.startswith("-") and len(arg) == 2:
@@ -238,7 +246,10 @@ class RecipeHandler:
                         if val.get("flag"):
                             result[key] = True
                         elif i + 1 < len(args):
-                            result[key] = args[i + 1]
+                            if val.get("multi"):
+                                result[key].append(args[i + 1])
+                            else:
+                                result[key] = args[i + 1]
                             i += 1
                         break
                 i += 1
@@ -551,6 +562,7 @@ class RecipeHandler:
             "input": {"short": "-i", "default": None},
             "config": {"short": "-c", "default": None},
             "session": {"short": "-s", "default": None},
+            "var": {"multi": True},  # --var key=value, can be used multiple times
             "json": {"flag": True, "default": False},
             "stream": {"flag": True, "default": False},
             "background": {"flag": True, "default": False},
@@ -566,6 +578,7 @@ class RecipeHandler:
             "offline": {"flag": True, "default": False},
             "force": {"flag": True, "default": False},
             "allow_dangerous_tools": {"flag": True, "default": False},
+            "save": {"flag": True, "default": False},  # Save replay trace
         }
         parsed = self._parse_args(args, spec)
         
@@ -592,6 +605,13 @@ class RecipeHandler:
         if remaining_positional and not parsed["input"]:
             input_data = {"input": remaining_positional[0]}
         
+        # Merge --var key=value pairs into input_data (variables)
+        if parsed.get("var"):
+            for v in parsed["var"]:
+                if "=" in v:
+                    key, value = v.split("=", 1)
+                    input_data[key] = value
+        
         # Parse config
         config = {}
         if parsed["config"]:
@@ -616,6 +636,7 @@ class RecipeHandler:
             "offline": parsed["offline"],
             "force": parsed["force"],
             "allow_dangerous_tools": parsed["allow_dangerous_tools"],
+            "save_replay": parsed["save"],  # Save replay trace for debugging
         }
         
         try:

src/praisonai/praisonai/cli/legacy.py

@@ -34,6 +34,8 @@ TYPER_COMMANDS = {
     'memory', 'workflow', 'tools', 'knowledge', 'deploy', 'agents', 'skills',
     'eval', 'templates', 'recipe', 'todo', 'docs', 'commit', 'hooks', 'rules',
     'registry', 'package', 'endpoints', 'test', 'examples', 'batch',
+    # Replay commands
+    'replay',
     # Standardisation commands
     'standardise', 'standardize',
 }

src/praisonai/praisonai/cli/main.py

@@ -742,7 +742,7 @@ class PraisonAI:
             return default_args
         
         # Define special commands
-        special_commands = ['chat', 'code', 'call', 'realtime', 'train', 'ui', 'context', 'research', 'memory', 'rules', 'workflow', 'hooks', 'knowledge', 'session', 'tools', 'todo', 'docs', 'mcp', 'commit', 'serve', 'schedule', 'skills', 'profile', 'eval', 'agents', 'run', 'thinking', 'compaction', 'output', 'deploy', 'templates', 'recipe', 'endpoints', 'audio', 'embed', 'embedding', 'images', 'moderate', 'files', 'batches', 'vector-stores', 'rerank', 'ocr', 'assistants', 'fine-tuning', 'completions', 'messages', 'guardrails', 'rag', 'videos', 'a2a', 'containers', 'passthrough', 'responses', 'search', 'realtime-api', 'doctor', 'registry', 'package', 'install', 'uninstall', 'acp', 'debug', 'lsp', 'diag', 'browser']
+        special_commands = ['chat', 'code', 'call', 'realtime', 'train', 'ui', 'context', 'research', 'memory', 'rules', 'workflow', 'hooks', 'knowledge', 'session', 'tools', 'todo', 'docs', 'mcp', 'commit', 'serve', 'schedule', 'skills', 'profile', 'eval', 'agents', 'run', 'thinking', 'compaction', 'output', 'deploy', 'templates', 'recipe', 'endpoints', 'audio', 'embed', 'embedding', 'images', 'moderate', 'files', 'batches', 'vector-stores', 'rerank', 'ocr', 'assistants', 'fine-tuning', 'completions', 'messages', 'guardrails', 'rag', 'videos', 'a2a', 'containers', 'passthrough', 'responses', 'search', 'realtime-api', 'doctor', 'registry', 'package', 'install', 'uninstall', 'acp', 'debug', 'lsp', 'diag', 'browser', 'replay']
         
         parser = argparse.ArgumentParser(prog="praisonai", description="praisonAI command-line interface")
         parser.add_argument("--framework", choices=["crewai", "autogen", "praisonai"], help="Specify the framework")
@@ -1009,6 +1009,18 @@ class PraisonAI:
             from praisonai.cli.features.diag import run_diag_command
             exit_code = run_diag_command(unknown_args)
             sys.exit(exit_code)
+        
+        # Handle replay command - context replay for debugging agent execution
+        if args.command == 'replay':
+            from .app import app as typer_app, register_commands
+            register_commands()
+            import sys as _sys
+            _sys.argv = ['praisonai', 'replay'] + unknown_args
+            try:
+                typer_app()
+            except SystemExit as e:
+                sys.exit(e.code if e.code else 0)
+            sys.exit(0)
 
         # Handle both command and flag versions for call
         if args.command == 'call' or args.call:
@@ -1522,6 +1534,19 @@ class PraisonAI:
                     sys.exit(e.code if e.code else 0)
                 sys.exit(0)
             
+            elif args.command == 'replay':
+                # Replay command - context replay for debugging agent execution
+                # Routes to Typer CLI for replay commands (list, context, show, flow, delete, cleanup)
+                from .app import app as typer_app, register_commands
+                register_commands()
+                import sys as _sys
+                _sys.argv = ['praisonai', 'replay'] + unknown_args
+                try:
+                    typer_app()
+                except SystemExit as e:
+                    sys.exit(e.code if e.code else 0)
+                sys.exit(0)
+            
             elif args.command == 'registry':
                 # Registry command - manage recipe registry server
                 from .features.registry import RegistryHandler
@@ -2411,6 +2436,33 @@ class PraisonAI:
             if verbose:
                 workflow.verbose = True
             
+            # Set context management from CLI args
+            context_auto_compact = getattr(args, 'context_auto_compact', None) if args else None
+            context_strategy = getattr(args, 'context_strategy', None) if args else None
+            context_threshold = getattr(args, 'context_threshold', None) if args else None
+            
+            if context_auto_compact is True or context_strategy or context_threshold:
+                # Enable context management with CLI-specified options
+                try:
+                    from praisonaiagents.context import ManagerConfig
+                    config_kwargs = {"auto_compact": True}
+                    if context_strategy:
+                        from praisonaiagents.context import OptimizerStrategy
+                        strategy_map = {
+                            "truncate": OptimizerStrategy.TRUNCATE,
+                            "sliding_window": OptimizerStrategy.SLIDING_WINDOW,
+                            "prune_tools": OptimizerStrategy.PRUNE_TOOLS,
+                            "summarize": OptimizerStrategy.SUMMARIZE,
+                            "smart": OptimizerStrategy.SMART,
+                        }
+                        config_kwargs["strategy"] = strategy_map.get(context_strategy, OptimizerStrategy.SMART)
+                    if context_threshold:
+                        config_kwargs["compact_threshold"] = context_threshold
+                    workflow.context = ManagerConfig(**config_kwargs)
+                    print(f"[cyan]Context management enabled (strategy={context_strategy or 'smart'}, threshold={context_threshold or 0.8})[/cyan]")
+                except ImportError:
+                    print("[yellow]Warning: Context management not available[/yellow]")
+            
             # Execute
             print("\n[bold]Executing workflow...[/bold]\n")
             result = workflow.start("")

src/praisonai/praisonai/cli/state/sessions.py

@@ -5,16 +5,18 @@ Provides persistence and resume functionality for sessions.
 """
 
 import json
-import os
 import shutil
 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
 
 from ..configuration.paths import get_sessions_dir, ensure_config_dirs
 from .identifiers import RunContext
 
+if TYPE_CHECKING:
+    from praisonaiagents.storage.protocols import StorageBackendProtocol
+
 
 @dataclass
 class SessionMetadata:
@@ -70,11 +72,35 @@ class SessionManager:
     Each session contains:
     - metadata.json: Session metadata
     - events.jsonl: Stream of events (NDJSON)
+    
+    Supports pluggable backends (file, sqlite, redis) via the backend parameter.
+    
+    Example with SQLite backend:
+        ```python
+        from praisonaiagents.storage import SQLiteBackend
+        backend = SQLiteBackend("~/.praison/sessions.db")
+        manager = SessionManager(backend=backend)
+        ```
     """
     
-    def __init__(self, sessions_dir: Optional[Path] = None):
+    def __init__(
+        self,
+        sessions_dir: Optional[Path] = None,
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
+        """
+        Initialize session manager.
+        
+        Args:
+            sessions_dir: Directory for file-based storage
+            backend: Optional storage backend (file, sqlite, redis).
+                     If provided, sessions_dir is ignored.
+        """
         self.sessions_dir = sessions_dir or get_sessions_dir()
-        ensure_config_dirs()
+        self._backend = backend
+        
+        if backend is None:
+            ensure_config_dirs()
     
     def _get_session_dir(self, session_id: str) -> Path:
         """Get the directory for a session."""
@@ -125,6 +151,10 @@ class SessionManager:
     
     def _save_metadata(self, metadata: SessionMetadata) -> None:
         """Save session metadata."""
+        if self._backend is not None:
+            self._backend.save(f"session:{metadata.session_id}:meta", metadata.to_dict())
+            return
+        
         path = self._get_metadata_path(metadata.session_id)
         path.parent.mkdir(parents=True, exist_ok=True)
         with open(path, "w", encoding="utf-8") as f:
@@ -132,6 +162,15 @@ class SessionManager:
     
     def _load_metadata(self, session_id: str) -> Optional[SessionMetadata]:
         """Load session metadata."""
+        if self._backend is not None:
+            data = self._backend.load(f"session:{session_id}:meta")
+            if data:
+                try:
+                    return SessionMetadata.from_dict(data)
+                except (KeyError, TypeError):
+                    return None
+            return None
+        
         path = self._get_metadata_path(session_id)
         if not path.exists():
             return None
@@ -150,12 +189,18 @@ class SessionManager:
             session_id: Session ID
             event: Event data to append
         """
-        events_path = self._get_events_path(session_id)
-        if not events_path.parent.exists():
-            return
-        
-        with open(events_path, "a", encoding="utf-8") as f:
-            f.write(json.dumps(event, default=str) + "\n")
+        if self._backend is not None:
+            # Load existing events, append, save
+            events = self._backend.load(f"session:{session_id}:events") or []
+            events.append(event)
+            self._backend.save(f"session:{session_id}:events", events)
+        else:
+            events_path = self._get_events_path(session_id)
+            if not events_path.parent.exists():
+                return
+            
+            with open(events_path, "a", encoding="utf-8") as f:
+                f.write(json.dumps(event, default=str) + "\n")
         
         # Update metadata
         metadata = self._load_metadata(session_id)
@@ -180,14 +225,29 @@ class SessionManager:
         """
         sessions = []
         
-        if not self.sessions_dir.exists():
-            return sessions
-        
-        for session_dir in self.sessions_dir.iterdir():
-            if session_dir.is_dir():
-                metadata = self._load_metadata(session_dir.name)
+        if self._backend is not None:
+            # List all session metadata keys
+            keys = self._backend.list_keys(prefix="session:")
+            session_ids = set()
+            for key in keys:
+                # Extract session_id from "session:<id>:meta" or "session:<id>:events"
+                parts = key.split(":")
+                if len(parts) >= 2:
+                    session_ids.add(parts[1])
+            
+            for session_id in session_ids:
+                metadata = self._load_metadata(session_id)
                 if metadata:
                     sessions.append(metadata)
+        else:
+            if not self.sessions_dir.exists():
+                return sessions
+            
+            for session_dir in self.sessions_dir.iterdir():
+                if session_dir.is_dir():
+                    metadata = self._load_metadata(session_dir.name)
+                    if metadata:
+                        sessions.append(metadata)
         
         # Sort by updated_at descending
         sessions.sort(key=lambda s: s.updated_at, reverse=True)
@@ -204,6 +264,11 @@ class SessionManager:
         Returns:
             True if deleted, False if not found
         """
+        if self._backend is not None:
+            deleted_meta = self._backend.delete(f"session:{session_id}:meta")
+            deleted_events = self._backend.delete(f"session:{session_id}:events")
+            return deleted_meta or deleted_events
+        
         session_dir = self._get_session_dir(session_id)
         if not session_dir.exists():
             return False
@@ -221,6 +286,9 @@ class SessionManager:
         Returns:
             List of events
         """
+        if self._backend is not None:
+            return self._backend.load(f"session:{session_id}:events") or []
+        
         events_path = self._get_events_path(session_id)
         if not events_path.exists():
             return []
@@ -303,11 +371,20 @@ class SessionManager:
 
 # Global session manager
 _session_manager: Optional[SessionManager] = None
+_session_backend: Optional["StorageBackendProtocol"] = None
+
+
+def set_session_backend(backend: "StorageBackendProtocol") -> None:
+    """Set the storage backend for session manager."""
+    global _session_backend, _session_manager
+    _session_backend = backend
+    # Reset manager to use new backend
+    _session_manager = None
 
 
 def get_session_manager() -> SessionManager:
     """Get the global session manager."""
     global _session_manager
     if _session_manager is None:
-        _session_manager = SessionManager()
+        _session_manager = SessionManager(backend=_session_backend)
     return _session_manager

src/praisonai/praisonai/deploy.py

@@ -57,7 +57,7 @@ class CloudDeployer:
             file.write("FROM python:3.11-slim\n")
             file.write("WORKDIR /app\n")
             file.write("COPY . .\n")
-            file.write("RUN pip install flask praisonai==3.9.31 gunicorn markdown\n")
+            file.write("RUN pip install flask praisonai==3.10.4 gunicorn markdown\n")
             file.write("EXPOSE 8080\n")
             file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
             

src/praisonai/praisonai/mcp_server/tool_index.py

@@ -23,9 +23,12 @@ import json
 import time
 import logging
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 from dataclasses import dataclass
 
+if TYPE_CHECKING:
+    from praisonaiagents.storage.protocols import StorageBackendProtocol
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,16 +91,32 @@ class MCPToolIndex:
     - Offline-first with sync
     """
     
-    def __init__(self, base_dir: str = "~/.praison/mcp"):
+    def __init__(
+        self,
+        base_dir: str = "~/.praison/mcp",
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
         """
         Initialize the tool index.
         
         Args:
-            base_dir: Base directory for MCP tool index
+            base_dir: Base directory for MCP tool index (file-based storage)
+            backend: Optional storage backend (file, sqlite, redis).
+                     If provided, base_dir is ignored and backend is used.
+        
+        Example with SQLite backend:
+            ```python
+            from praisonaiagents.storage import SQLiteBackend
+            backend = SQLiteBackend("~/.praison/mcp.db")
+            index = MCPToolIndex(backend=backend)
+            ```
         """
         self.base_dir = Path(base_dir).expanduser().resolve()
         self.servers_dir = self.base_dir / "servers"
-        self._ensure_dirs()
+        self._backend = backend
+        
+        if backend is None:
+            self._ensure_dirs()
     
     def _ensure_dirs(self) -> None:
         """Create necessary directories."""
@@ -151,8 +170,11 @@ class MCPToolIndex:
             schema_path.write_text(json.dumps(tool, indent=2))
         
         # Save index
-        index_path = server_dir / "_index.json"
-        index_path.write_text(json.dumps(index_data, indent=2))
+        if self._backend is not None:
+            self._backend.save(f"mcp:{server}:index", index_data)
+        else:
+            index_path = server_dir / "_index.json"
+            index_path.write_text(json.dumps(index_data, indent=2))
         
         # Update status
         self._update_status(
@@ -191,9 +213,6 @@ class MCPToolIndex:
         error: Optional[str] = None,
     ) -> None:
         """Update server status file."""
-        server_dir = self._get_server_dir(server)
-        status_path = server_dir / "_status.json"
-        
         status = ServerStatus(
             server=server,
             available=available,
@@ -203,7 +222,12 @@ class MCPToolIndex:
             error=error,
         )
         
-        status_path.write_text(json.dumps(status.to_dict(), indent=2))
+        if self._backend is not None:
+            self._backend.save(f"mcp:{server}:status", status.to_dict())
+        else:
+            server_dir = self._get_server_dir(server)
+            status_path = server_dir / "_status.json"
+            status_path.write_text(json.dumps(status.to_dict(), indent=2))
     
     def get_status(self, server: str) -> Optional[ServerStatus]:
         """
@@ -215,6 +239,15 @@ class MCPToolIndex:
         Returns:
             ServerStatus or None if not found
         """
+        if self._backend is not None:
+            data = self._backend.load(f"mcp:{server}:status")
+            if data:
+                try:
+                    return ServerStatus(**data)
+                except TypeError:
+                    return None
+            return None
+        
         server_dir = self._get_server_dir(server)
         status_path = server_dir / "_status.json"
         
@@ -229,6 +262,16 @@ class MCPToolIndex:
     
     def list_servers(self) -> List[str]:
         """List all indexed servers."""
+        if self._backend is not None:
+            keys = self._backend.list_keys(prefix="mcp:")
+            servers = set()
+            for key in keys:
+                # Extract server from "mcp:<server>:index" or "mcp:<server>:status"
+                parts = key.split(":")
+                if len(parts) >= 2:
+                    servers.add(parts[1])
+            return sorted(servers)
+        
         servers = []
         for path in self.servers_dir.iterdir():
             if path.is_dir():

src/praisonai/praisonai/recipe/core.py

@@ -189,6 +189,23 @@ def run(
         "session_id": session_id,
     }
     
+    # Initialize replay trace writer if --save flag is set
+    trace_writer = None
+    trace_emitter = None
+    if options.get("save_replay", False):
+        try:
+            from praisonai.replay import ContextTraceWriter
+            from praisonaiagents.trace.context_events import ContextTraceEmitter
+            trace_writer = ContextTraceWriter(session_id=run_id)
+            trace_emitter = ContextTraceEmitter(sink=trace_writer, session_id=run_id)
+            trace_emitter.session_start({"recipe": name, "run_id": run_id})
+        except ImportError as e:
+            import logging
+            logging.debug(f"Replay module not available: {e}")
+        except Exception as e:
+            import logging
+            logging.warning(f"Failed to initialize trace writer: {e}")
+    
     try:
         # Load template
         recipe_config = _load_recipe(name, offline=options.get("offline", False))
@@ -235,6 +252,11 @@ def run(
         
         # Dry run mode
         if options.get("dry_run", False):
+            # Close trace writer for dry-run
+            if trace_emitter:
+                trace_emitter.session_end()
+            if trace_writer:
+                trace_writer.close()
             return RecipeResult(
                 run_id=run_id,
                 recipe=recipe_config.name,
@@ -255,16 +277,23 @@ def run(
         else:
             merged_config = {**recipe_config.defaults, **input, **config}
         
-        # Execute recipe
+        # Execute recipe (pass trace_emitter for event tracking)
         output = _execute_recipe(
             recipe_config,
             merged_config,
             session_id,
             options,
+            trace_emitter=trace_emitter,
         )
         
         duration = time.time() - start_time
         
+        # Close trace writer on success
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
+        
         return RecipeResult(
             run_id=run_id,
             recipe=recipe_config.name,
@@ -276,6 +305,10 @@ def run(
         )
         
     except RecipeNotFoundError as e:
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
         return RecipeResult(
             run_id=run_id,
             recipe=name,
@@ -286,6 +319,10 @@ def run(
             trace=trace,
         )
     except RecipeDependencyError as e:
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
         return RecipeResult(
             run_id=run_id,
             recipe=e.recipe or name,
@@ -296,6 +333,10 @@ def run(
             trace=trace,
         )
     except RecipePolicyError as e:
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
         return RecipeResult(
             run_id=run_id,
             recipe=e.recipe or name,
@@ -306,6 +347,10 @@ def run(
             trace=trace,
         )
     except RecipeTimeoutError as e:
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
         return RecipeResult(
             run_id=run_id,
             recipe=name,
@@ -316,6 +361,10 @@ def run(
             trace=trace,
         )
     except Exception as e:
+        if trace_emitter:
+            trace_emitter.session_end()
+        if trace_writer:
+            trace_writer.close()
         return RecipeResult(
             run_id=run_id,
             recipe=name,
@@ -660,6 +709,43 @@ def _load_recipe(name: str, offline: bool = False) -> Optional[RecipeConfig]:
         return None
 
 
+def _check_package_installed(package_name: str) -> bool:
+    """
+    Check if a Python package is installed using importlib.metadata.
+    
+    This is more reliable than trying to import the package because:
+    1. Package names often differ from import names (e.g., tavily-python -> tavily)
+    2. importlib.metadata checks the actual installed packages
+    3. No side effects from importing modules
+    
+    Args:
+        package_name: The pip package name (e.g., 'tavily-python', 'pillow')
+        
+    Returns:
+        True if package is installed, False otherwise
+    """
+    try:
+        from importlib.metadata import distributions
+        
+        # Normalize package name for comparison (pip normalizes - and _ and case)
+        normalized = package_name.lower().replace("-", "_").replace(".", "_")
+        
+        for dist in distributions():
+            dist_name = dist.metadata.get("Name", "").lower().replace("-", "_").replace(".", "_")
+            if dist_name == normalized:
+                return True
+        
+        return False
+    except Exception:
+        # Fallback: try to import with common name transformations
+        try:
+            import_name = package_name.replace("-", "_")
+            __import__(import_name)
+            return True
+        except ImportError:
+            return False
+
+
 def _check_dependencies(recipe_config: RecipeConfig) -> Dict[str, Any]:
     """Check if recipe dependencies are satisfied."""
     result = {
@@ -670,28 +756,11 @@ def _check_dependencies(recipe_config: RecipeConfig) -> Dict[str, Any]:
         "external": [],
     }
     
-    # Common package name to import name mappings
-    PACKAGE_IMPORT_MAP = {
-        "tavily-python": "tavily",
-        "google-generativeai": "google.generativeai",
-        "openai-whisper": "whisper",
-        "python-dotenv": "dotenv",
-        "pillow": "PIL",
-        "opencv-python": "cv2",
-        "scikit-learn": "sklearn",
-        "beautifulsoup4": "bs4",
-        "pyyaml": "yaml",
-    }
-    
     # Check Python packages
     for pkg in recipe_config.get_required_packages():
-        # Get the correct import name
-        import_name = PACKAGE_IMPORT_MAP.get(pkg, pkg.replace("-", "_"))
-        try:
-            __import__(import_name)
-            result["packages"].append({"name": pkg, "available": True})
-        except ImportError:
-            result["packages"].append({"name": pkg, "available": False})
+        available = _check_package_installed(pkg)
+        result["packages"].append({"name": pkg, "available": available})
+        if not available:
             result["all_satisfied"] = False
     
     # Check environment variables
@@ -751,6 +820,7 @@ def _execute_recipe(
     merged_config: Dict[str, Any],
     session_id: str,
     options: Dict[str, Any],
+    trace_emitter: Any = None,
 ) -> Any:
     """Execute the recipe workflow."""
     try:
@@ -808,6 +878,9 @@ def _execute_recipe(
                 workflow_config, merged_config, tool_registry, options
             )
         elif "steps" in workflow_config:
+            # Pass the workflow file path for proper execution
+            workflow_file_path = template_path / workflow_file if template_path else None
+            merged_config["_workflow_file"] = str(workflow_file_path) if workflow_file_path else None
             return _execute_steps_workflow(
                 workflow_config, merged_config, tool_registry, options
             )
@@ -900,9 +973,45 @@ def _execute_steps_workflow(
     tool_registry: Any,
     options: Dict[str, Any],
 ) -> Any:
-    """Execute a steps-based workflow."""
-    # For now, convert to simple execution
-    return {"message": "Steps workflow executed", "config": config}
+    """
+    Execute a steps-based workflow using praisonaiagents Workflow.
+    
+    This properly executes the workflow with all features:
+    - include steps (modular recipes)
+    - loop steps (parallel/sequential)
+    - output_variable
+    - agent execution
+    - CLI variable overrides via --var
+    """
+    from praisonaiagents.workflows import YAMLWorkflowParser
+    
+    # Get the workflow file path from config if available
+    workflow_file = config.get("_workflow_file")
+    
+    # Extract CLI variable overrides (everything except internal keys)
+    extra_vars = {k: v for k, v in config.items() if not k.startswith("_")}
+    
+    if workflow_file:
+        # Use YAMLWorkflowParser to properly parse and execute the workflow
+        # This handles include steps, loops, variables, tools, etc.
+        parser = YAMLWorkflowParser(tool_registry=tool_registry)
+        workflow = parser.parse_file(workflow_file, extra_vars=extra_vars)
+        return workflow.start()
+    
+    # Fallback: Create workflow from config dict using parser
+    from praisonaiagents import Workflow
+    
+    # Merge YAML variables with CLI overrides (CLI takes precedence)
+    yaml_vars = workflow_config.get("variables", {})
+    merged_vars = {**yaml_vars, **extra_vars}
+    
+    workflow = Workflow(
+        name=workflow_config.get("name", "RecipeWorkflow"),
+        steps=workflow_config.get("steps", []),
+        variables=merged_vars,
+    )
+    return workflow.start()
+
 
 
 def _execute_simple_agent(

src/praisonai/praisonai/recipe/history.py

@@ -11,10 +11,13 @@ import json
 import shutil
 from datetime import datetime, timezone, timedelta
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
 from .models import RecipeResult
 
+if TYPE_CHECKING:
+    from praisonaiagents.storage.protocols import StorageBackendProtocol
+
 
 # Default storage path
 DEFAULT_RUNS_PATH = Path.home() / ".praison" / "runs"
@@ -39,7 +42,7 @@ class RunHistory:
     """
     Local run history storage.
     
-    Storage structure:
+    Storage structure (file-based):
     ~/.praison/runs/
     ├── index.json           # Run index for fast lookup
     └── <run_id>/
@@ -47,13 +50,36 @@ class RunHistory:
         ├── input.json       # Input data
         ├── output.json      # Output data
         └── events.jsonl     # Event stream (if streaming)
+    
+    Supports pluggable backends (file, sqlite, redis) via the backend parameter.
+    
+    Example with SQLite backend:
+        ```python
+        from praisonaiagents.storage import SQLiteBackend
+        backend = SQLiteBackend("~/.praison/runs.db")
+        history = RunHistory(backend=backend)
+        ```
     """
     
-    def __init__(self, path: Optional[Path] = None):
-        """Initialize run history storage."""
+    def __init__(
+        self,
+        path: Optional[Path] = None,
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
+        """
+        Initialize run history storage.
+        
+        Args:
+            path: Path for file-based storage (default: ~/.praison/runs)
+            backend: Optional storage backend (file, sqlite, redis).
+                     If provided, path is ignored and backend is used.
+        """
         self.path = Path(path) if path else DEFAULT_RUNS_PATH
         self.index_path = self.path / "index.json"
-        self._ensure_structure()
+        self._backend = backend
+        
+        if backend is None:
+            self._ensure_structure()
     
     def _ensure_structure(self):
         """Ensure storage directory structure exists."""
@@ -63,6 +89,10 @@ class RunHistory:
     
     def _load_index(self) -> Dict[str, Any]:
         """Load run index."""
+        if self._backend is not None:
+            loaded = self._backend.load("_index")
+            return loaded if loaded else {"runs": {}, "updated": _get_timestamp()}
+        
         if self.index_path.exists():
             with open(self.index_path) as f:
                 return json.load(f)
@@ -71,6 +101,11 @@ class RunHistory:
     def _save_index(self, index: Dict[str, Any]):
         """Save run index."""
         index["updated"] = _get_timestamp()
+        
+        if self._backend is not None:
+            self._backend.save("_index", index)
+            return
+        
         with open(self.index_path, "w") as f:
             json.dump(index, f, indent=2)
     
@@ -119,24 +154,35 @@ class RunHistory:
             "retention_days": retention_days,
         }
         
-        with open(run_dir / "run.json", "w") as f:
-            json.dump(run_data, f, indent=2)
-        
-        # Store input (if allowed)
-        if input_data and export_allowed:
-            with open(run_dir / "input.json", "w") as f:
-                json.dump(input_data, f, indent=2)
-        
-        # Store output (if allowed)
-        if result.output and export_allowed:
-            with open(run_dir / "output.json", "w") as f:
-                json.dump(result.output, f, indent=2)
-        
-        # Store events (if any)
-        if events:
-            with open(run_dir / "events.jsonl", "w") as f:
-                for event in events:
-                    f.write(json.dumps(event) + "\n")
+        if self._backend is not None:
+            # Backend-based storage
+            full_data = {
+                "run": run_data,
+                "input": input_data if input_data and export_allowed else None,
+                "output": result.output if result.output and export_allowed else None,
+                "events": events,
+            }
+            self._backend.save(f"run:{run_id}", full_data)
+        else:
+            # File-based storage
+            with open(run_dir / "run.json", "w") as f:
+                json.dump(run_data, f, indent=2)
+            
+            # Store input (if allowed)
+            if input_data and export_allowed:
+                with open(run_dir / "input.json", "w") as f:
+                    json.dump(input_data, f, indent=2)
+            
+            # Store output (if allowed)
+            if result.output and export_allowed:
+                with open(run_dir / "output.json", "w") as f:
+                    json.dump(result.output, f, indent=2)
+            
+            # Store events (if any)
+            if events:
+                with open(run_dir / "events.jsonl", "w") as f:
+                    for event in events:
+                        f.write(json.dumps(event) + "\n")
         
         # Update index
         index = self._load_index()
@@ -164,6 +210,22 @@ class RunHistory:
         Raises:
             RunNotFoundError: If run not found
         """
+        if self._backend is not None:
+            # Backend-based retrieval
+            full_data = self._backend.load(f"run:{run_id}")
+            if not full_data:
+                raise RunNotFoundError(f"Run not found: {run_id}")
+            
+            run_data = full_data.get("run", {})
+            if full_data.get("input"):
+                run_data["input"] = full_data["input"]
+            if full_data.get("output"):
+                run_data["output"] = full_data["output"]
+            if full_data.get("events"):
+                run_data["events"] = full_data["events"]
+            return run_data
+        
+        # File-based retrieval
         run_dir = self.path / run_id
         if not run_dir.exists():
             raise RunNotFoundError(f"Run not found: {run_id}")
@@ -294,9 +356,12 @@ class RunHistory:
         Returns:
             True if deleted
         """
-        run_dir = self.path / run_id
-        if run_dir.exists():
-            shutil.rmtree(run_dir)
+        if self._backend is not None:
+            self._backend.delete(f"run:{run_id}")
+        else:
+            run_dir = self.path / run_id
+            if run_dir.exists():
+                shutil.rmtree(run_dir)
         
         index = self._load_index()
         if run_id in index["runs"]:

src/praisonai/praisonai/replay/__init__.py

@@ -0,0 +1,56 @@
+"""
+Replay Module for PraisonAI.
+
+Provides context replay functionality for debugging and analysis.
+Allows stepping through agent execution context changes.
+
+Usage:
+    from praisonai.replay import ContextTraceWriter, ContextTraceReader, ReplayPlayer
+    
+    # Write traces during execution
+    writer = ContextTraceWriter(session_id="my-session")
+    writer.emit(event)
+    writer.close()
+    
+    # Read and replay traces
+    reader = ContextTraceReader("~/.praison/traces/my-session.jsonl")
+    for event in reader:
+        print(event)
+    
+    # Interactive replay
+    player = ReplayPlayer(reader)
+    player.run()
+"""
+
+__all__ = [
+    "ContextTraceWriter",
+    "ContextTraceReader",
+    "ReplayPlayer",
+    "get_traces_dir",
+    "list_traces",
+]
+
+
+def __getattr__(name: str):
+    """Lazy load module components."""
+    if name == "ContextTraceWriter":
+        from .writer import ContextTraceWriter
+        return ContextTraceWriter
+    
+    if name == "ContextTraceReader":
+        from .reader import ContextTraceReader
+        return ContextTraceReader
+    
+    if name == "ReplayPlayer":
+        from .player import ReplayPlayer
+        return ReplayPlayer
+    
+    if name == "get_traces_dir":
+        from .storage import get_traces_dir
+        return get_traces_dir
+    
+    if name == "list_traces":
+        from .storage import list_traces
+        return list_traces
+    
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai/praisonai/replay/player.py

@@ -0,0 +1,276 @@
+"""
+Interactive Replay Player for PraisonAI.
+
+Provides step-through replay of context traces with keyboard navigation.
+"""
+
+from datetime import datetime
+
+from .reader import ContextTraceReader
+
+
+class ReplayPlayer:
+    """
+    Interactive replay player for context traces.
+    
+    Allows stepping through context events with keyboard navigation.
+    
+    Usage:
+        reader = ContextTraceReader("my-session")
+        player = ReplayPlayer(reader)
+        player.run()
+        
+    Navigation:
+        Enter/Down: Next event
+        Up: Previous event
+        g <num>: Go to event number
+        q/Ctrl+C: Quit
+    """
+    
+    def __init__(
+        self,
+        reader: ContextTraceReader,
+        use_rich: bool = True,
+    ):
+        """
+        Initialize the player.
+        
+        Args:
+            reader: ContextTraceReader with loaded events
+            use_rich: Whether to use Rich for formatting (if available)
+        """
+        self._reader = reader
+        self._use_rich = use_rich
+        self._current_index = 0
+        self._events = reader.get_all()
+        
+        # Try to import Rich for better formatting
+        self._console = None
+        if use_rich:
+            try:
+                from rich.console import Console
+                from rich.panel import Panel
+                from rich.table import Table
+                self._console = Console()
+                self._Panel = Panel
+                self._Table = Table
+            except ImportError:
+                self._use_rich = False
+    
+    def _format_timestamp(self, timestamp: float) -> str:
+        """Format timestamp for display."""
+        try:
+            dt = datetime.fromtimestamp(timestamp)
+            return dt.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        except Exception:
+            return str(timestamp)
+    
+    def _get_event_type_display(self, event) -> str:
+        """Get display string for event type."""
+        if hasattr(event, 'event_type'):
+            et = event.event_type
+            if hasattr(et, 'value'):
+                return et.value.upper().replace('_', ' ')
+            return str(et).upper().replace('_', ' ')
+        elif isinstance(event, dict):
+            return event.get('event_type', 'UNKNOWN').upper().replace('_', ' ')
+        return 'UNKNOWN'
+    
+    def _display_event_rich(self, event, index: int) -> None:
+        """Display event using Rich formatting."""
+        if not self._console:
+            return self._display_event_plain(event, index)
+        
+        # Build header
+        event_type = self._get_event_type_display(event)
+        agent_name = getattr(event, 'agent_name', None) or (event.get('agent_name') if isinstance(event, dict) else None)
+        timestamp = getattr(event, 'timestamp', 0) or (event.get('timestamp', 0) if isinstance(event, dict) else 0)
+        
+        header = f"[bold cyan]Event {index + 1}/{len(self._events)}[/bold cyan]"
+        if agent_name:
+            header += f" | [bold green]Agent: {agent_name}[/bold green]"
+        
+        # Build content
+        content_lines = [
+            f"[bold]Type:[/bold] {event_type}",
+            f"[bold]Time:[/bold] {self._format_timestamp(timestamp)}",
+            f"[bold]Session:[/bold] {self._reader.session_id}",
+        ]
+        
+        # Add context state if available
+        messages_count = getattr(event, 'messages_count', 0) or (event.get('messages_count', 0) if isinstance(event, dict) else 0)
+        tokens_used = getattr(event, 'tokens_used', 0) or (event.get('tokens_used', 0) if isinstance(event, dict) else 0)
+        tokens_budget = getattr(event, 'tokens_budget', 0) or (event.get('tokens_budget', 0) if isinstance(event, dict) else 0)
+        
+        if messages_count or tokens_used:
+            content_lines.append("")
+            content_lines.append("[bold]Context State:[/bold]")
+            if messages_count:
+                content_lines.append(f"  Messages: {messages_count}")
+            if tokens_used:
+                budget_str = f" / {tokens_budget:,}" if tokens_budget else ""
+                content_lines.append(f"  Tokens: {tokens_used:,}{budget_str}")
+        
+        # Add event data if available
+        data = getattr(event, 'data', {}) or (event.get('data', {}) if isinstance(event, dict) else {})
+        if data:
+            content_lines.append("")
+            content_lines.append("[bold]Event Data:[/bold]")
+            for key, value in data.items():
+                if value is not None:
+                    # Truncate long values
+                    value_str = str(value)
+                    if len(value_str) > 100:
+                        value_str = value_str[:100] + "..."
+                    content_lines.append(f"  {key}: {value_str}")
+        
+        content = "\n".join(content_lines)
+        
+        # Display panel
+        self._console.clear()
+        self._console.print(self._Panel(
+            content,
+            title=header,
+            border_style="blue",
+        ))
+        
+        # Navigation help
+        self._console.print("\n[dim][Enter/↓] Next  [↑] Previous  [g N] Go to N  [q] Quit[/dim]")
+    
+    def _display_event_plain(self, event, index: int) -> None:
+        """Display event using plain text."""
+        print("\n" + "=" * 60)
+        print(f"Event {index + 1}/{len(self._events)}")
+        print("=" * 60)
+        
+        event_type = self._get_event_type_display(event)
+        agent_name = getattr(event, 'agent_name', None) or (event.get('agent_name') if isinstance(event, dict) else None)
+        timestamp = getattr(event, 'timestamp', 0) or (event.get('timestamp', 0) if isinstance(event, dict) else 0)
+        
+        print(f"Type: {event_type}")
+        print(f"Time: {self._format_timestamp(timestamp)}")
+        if agent_name:
+            print(f"Agent: {agent_name}")
+        print(f"Session: {self._reader.session_id}")
+        
+        # Context state
+        messages_count = getattr(event, 'messages_count', 0) or (event.get('messages_count', 0) if isinstance(event, dict) else 0)
+        tokens_used = getattr(event, 'tokens_used', 0) or (event.get('tokens_used', 0) if isinstance(event, dict) else 0)
+        tokens_budget = getattr(event, 'tokens_budget', 0) or (event.get('tokens_budget', 0) if isinstance(event, dict) else 0)
+        
+        if messages_count or tokens_used:
+            print("\nContext State:")
+            if messages_count:
+                print(f"  Messages: {messages_count}")
+            if tokens_used:
+                budget_str = f" / {tokens_budget:,}" if tokens_budget else ""
+                print(f"  Tokens: {tokens_used:,}{budget_str}")
+        
+        # Event data
+        data = getattr(event, 'data', {}) or (event.get('data', {}) if isinstance(event, dict) else {})
+        if data:
+            print("\nEvent Data:")
+            for key, value in data.items():
+                if value is not None:
+                    value_str = str(value)
+                    if len(value_str) > 100:
+                        value_str = value_str[:100] + "..."
+                    print(f"  {key}: {value_str}")
+        
+        print("\n[Enter/Down] Next  [Up] Previous  [g N] Go to N  [q] Quit")
+    
+    def display_current(self) -> None:
+        """Display the current event."""
+        if not self._events:
+            print("No events to display.")
+            return
+        
+        event = self._events[self._current_index]
+        
+        if self._use_rich and self._console:
+            self._display_event_rich(event, self._current_index)
+        else:
+            self._display_event_plain(event, self._current_index)
+    
+    def next(self) -> bool:
+        """Move to next event. Returns False if at end."""
+        if self._current_index < len(self._events) - 1:
+            self._current_index += 1
+            return True
+        return False
+    
+    def previous(self) -> bool:
+        """Move to previous event. Returns False if at start."""
+        if self._current_index > 0:
+            self._current_index -= 1
+            return True
+        return False
+    
+    def goto(self, index: int) -> bool:
+        """Go to specific event index (0-based). Returns False if invalid."""
+        if 0 <= index < len(self._events):
+            self._current_index = index
+            return True
+        return False
+    
+    def run(self) -> None:
+        """Run the interactive replay loop."""
+        if not self._events:
+            print(f"No events found in trace: {self._reader.session_id}")
+            return
+        
+        print(f"Replaying {len(self._events)} events from session: {self._reader.session_id}")
+        print("Press Enter to start...")
+        
+        try:
+            input()
+        except (KeyboardInterrupt, EOFError):
+            return
+        
+        while True:
+            self.display_current()
+            
+            try:
+                user_input = input("\n> ").strip().lower()
+            except (KeyboardInterrupt, EOFError):
+                print("\nExiting replay.")
+                break
+            
+            if user_input in ('q', 'quit', 'exit'):
+                print("Exiting replay.")
+                break
+            elif user_input in ('', 'n', 'next', 'down'):
+                if not self.next():
+                    print("End of trace reached.")
+            elif user_input in ('p', 'prev', 'previous', 'up'):
+                if not self.previous():
+                    print("Already at start of trace.")
+            elif user_input.startswith('g ') or user_input.startswith('goto '):
+                try:
+                    parts = user_input.split()
+                    if len(parts) >= 2:
+                        index = int(parts[1]) - 1  # Convert to 0-based
+                        if not self.goto(index):
+                            print(f"Invalid event number. Valid range: 1-{len(self._events)}")
+                except ValueError:
+                    print("Invalid number. Usage: g <event_number>")
+            elif user_input == 'h' or user_input == 'help':
+                print("\nCommands:")
+                print("  Enter/n/down - Next event")
+                print("  p/up - Previous event")
+                print("  g <N> - Go to event N")
+                print("  q - Quit")
+            else:
+                # Default to next on any other input
+                if not self.next():
+                    print("End of trace reached.")
+    
+    @property
+    def current_index(self) -> int:
+        """Get current event index."""
+        return self._current_index
+    
+    @property
+    def total_events(self) -> int:
+        """Get total number of events."""
+        return len(self._events)

src/praisonai/praisonai/replay/reader.py

@@ -0,0 +1,151 @@
+"""
+Context Trace Reader for PraisonAI.
+
+Reads context events from JSONL files for replay.
+"""
+
+import json
+from pathlib import Path
+from typing import Iterator, List, Optional
+
+from .storage import get_trace_path
+
+
+class ContextTraceReader:
+    """
+    JSONL reader for context trace events.
+    
+    Reads events from a trace file and converts them to ContextEvent objects.
+    Supports iteration and random access.
+    
+    Usage:
+        reader = ContextTraceReader("my-session")
+        
+        # Iterate over events
+        for event in reader:
+            print(event.event_type)
+        
+        # Random access
+        event = reader[5]
+        
+        # Get all events
+        events = reader.get_all()
+    """
+    
+    def __init__(
+        self,
+        session_id_or_path: str,
+        traces_dir: Optional[Path] = None,
+    ):
+        """
+        Initialize the reader.
+        
+        Args:
+            session_id_or_path: Session ID or full path to trace file
+            traces_dir: Directory for trace files (if session_id provided)
+        """
+        # Determine path
+        if "/" in session_id_or_path or session_id_or_path.endswith(".jsonl"):
+            self._path = Path(session_id_or_path).expanduser()
+            self._session_id = self._path.stem
+        else:
+            self._session_id = session_id_or_path
+            self._path = get_trace_path(session_id_or_path, traces_dir)
+        
+        self._events: Optional[List] = None
+    
+    def _load_events(self) -> List:
+        """Load all events from file."""
+        if self._events is not None:
+            return self._events
+        
+        # Lazy import to avoid circular dependency
+        try:
+            from praisonaiagents.trace.context_events import ContextEvent
+        except ImportError:
+            # Fallback: return raw dicts if praisonaiagents not available
+            ContextEvent = None
+        
+        events = []
+        
+        if not self._path.exists():
+            self._events = events
+            return events
+        
+        try:
+            with open(self._path, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    
+                    try:
+                        data = json.loads(line)
+                        if ContextEvent is not None:
+                            event = ContextEvent.from_dict(data)
+                        else:
+                            event = data
+                        events.append(event)
+                    except (json.JSONDecodeError, Exception):
+                        continue
+        except Exception:
+            pass
+        
+        self._events = events
+        return events
+    
+    def __iter__(self) -> Iterator:
+        """Iterate over events."""
+        return iter(self._load_events())
+    
+    def __len__(self) -> int:
+        """Get number of events."""
+        return len(self._load_events())
+    
+    def __getitem__(self, index: int):
+        """Get event by index."""
+        return self._load_events()[index]
+    
+    def get_all(self) -> List:
+        """Get all events as a list."""
+        return self._load_events().copy()
+    
+    def get_by_agent(self, agent_name: str) -> List:
+        """Get events for a specific agent."""
+        return [e for e in self._load_events() if getattr(e, 'agent_name', None) == agent_name]
+    
+    def get_by_type(self, event_type: str) -> List:
+        """Get events of a specific type."""
+        events = self._load_events()
+        result = []
+        for e in events:
+            # Handle both ContextEvent objects and dicts
+            if hasattr(e, 'event_type'):
+                et = e.event_type
+                if hasattr(et, 'value'):
+                    et = et.value
+                if et == event_type:
+                    result.append(e)
+            elif isinstance(e, dict) and e.get('event_type') == event_type:
+                result.append(e)
+        return result
+    
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self._session_id
+    
+    @property
+    def path(self) -> Path:
+        """Get the trace file path."""
+        return self._path
+    
+    @property
+    def exists(self) -> bool:
+        """Check if trace file exists."""
+        return self._path.exists()
+    
+    def reload(self) -> None:
+        """Reload events from file."""
+        self._events = None
+        self._load_events()

src/praisonai/praisonai/replay/storage.py

@@ -0,0 +1,132 @@
+"""
+Storage utilities for context traces.
+
+Provides functions for managing trace storage locations and listing traces.
+"""
+
+import os
+from pathlib import Path
+from typing import List, Optional
+
+# DRY: Import base classes from praisonaiagents.storage
+from praisonaiagents.storage.models import BaseSessionInfo
+from praisonaiagents.storage.base import list_json_sessions, cleanup_old_sessions as _cleanup_old_sessions
+
+
+DEFAULT_TRACES_DIR = "~/.praison/traces"
+
+
+def get_traces_dir() -> Path:
+    """
+    Get the traces directory path.
+    
+    Returns:
+        Path to traces directory (creates if not exists)
+    """
+    traces_dir = Path(os.path.expanduser(DEFAULT_TRACES_DIR))
+    traces_dir.mkdir(parents=True, exist_ok=True)
+    return traces_dir
+
+
+class TraceInfo(BaseSessionInfo):
+    """
+    Information about a trace file.
+    
+    DRY: Inherits from BaseSessionInfo which provides:
+    - session_id, path, size_bytes, created_at, modified_at, item_count
+    - to_dict(), from_dict(), from_path() methods
+    """
+    
+    @property
+    def event_count(self) -> int:
+        """Alias for item_count for backward compatibility."""
+        return self.item_count
+    
+    def to_dict(self):
+        """Override to include event_count for backward compatibility."""
+        d = super().to_dict()
+        d["event_count"] = self.event_count
+        return d
+
+
+def list_traces(
+    traces_dir: Optional[Path] = None,
+    limit: int = 50,
+) -> List[TraceInfo]:
+    """
+    List available trace files.
+    
+    DRY: Uses list_json_sessions from praisonaiagents.storage.base.
+    
+    Args:
+        traces_dir: Directory to search (default: ~/.praison/traces)
+        limit: Maximum number of traces to return
+        
+    Returns:
+        List of TraceInfo objects, sorted by modification time (newest first)
+    """
+    if traces_dir is None:
+        traces_dir = get_traces_dir()
+    
+    # DRY: Use common list_json_sessions function with .jsonl suffix
+    base_sessions = list_json_sessions(Path(traces_dir), suffix=".jsonl", limit=limit)
+    
+    # Convert BaseSessionInfo to TraceInfo for backward compatibility
+    return [
+        TraceInfo(
+            session_id=s.session_id,
+            path=s.path,
+            size_bytes=s.size_bytes,
+            created_at=s.created_at,
+            modified_at=s.modified_at,
+            item_count=s.item_count,
+        )
+        for s in base_sessions
+    ]
+
+
+def get_trace_path(session_id: str, traces_dir: Optional[Path] = None) -> Path:
+    """
+    Get the path for a trace file.
+    
+    Args:
+        session_id: Session identifier
+        traces_dir: Directory for traces (default: ~/.praison/traces)
+        
+    Returns:
+        Path to the trace file
+    """
+    if traces_dir is None:
+        traces_dir = get_traces_dir()
+    
+    return traces_dir / f"{session_id}.jsonl"
+
+
+def cleanup_old_traces(
+    traces_dir: Optional[Path] = None,
+    max_age_days: int = 7,
+    max_size_mb: int = 100,
+) -> int:
+    """
+    Clean up old trace files.
+    
+    DRY: Uses cleanup_old_sessions from praisonaiagents.storage.base.
+    
+    Args:
+        traces_dir: Directory to clean (default: ~/.praison/traces)
+        max_age_days: Delete traces older than this
+        max_size_mb: Delete oldest traces if total size exceeds this
+        
+    Returns:
+        Number of files deleted
+    """
+    if traces_dir is None:
+        traces_dir = get_traces_dir()
+    
+    # DRY: Use common cleanup function with .jsonl suffix
+    return _cleanup_old_sessions(
+        storage_dir=Path(traces_dir),
+        suffix=".jsonl",
+        max_age_days=max_age_days,
+        max_size_mb=max_size_mb,
+    )

src/praisonai/praisonai/replay/writer.py

@@ -0,0 +1,127 @@
+"""
+Context Trace Writer for PraisonAI.
+
+Writes context events to JSONL files for later replay.
+Thread-safe with buffered writes for performance.
+"""
+
+import json
+import threading
+from pathlib import Path
+from typing import List, Optional
+
+from .storage import get_trace_path
+
+
+class ContextTraceWriter:
+    """
+    JSONL writer for context trace events.
+    
+    Implements the ContextTraceSink protocol from praisonaiagents.
+    Writes events to a JSONL file with buffering for performance.
+    
+    Usage:
+        writer = ContextTraceWriter(session_id="my-session")
+        
+        # Emit events
+        writer.emit(event)
+        
+        # Flush and close when done
+        writer.close()
+    """
+    
+    def __init__(
+        self,
+        session_id: str,
+        traces_dir: Optional[Path] = None,
+        buffer_size: int = 100,
+        auto_flush: bool = True,
+    ):
+        """
+        Initialize the writer.
+        
+        Args:
+            session_id: Session identifier (used as filename)
+            traces_dir: Directory for trace files (default: ~/.praison/traces)
+            buffer_size: Number of events to buffer before flushing
+            auto_flush: Whether to auto-flush on buffer full
+        """
+        self._session_id = session_id
+        self._path = get_trace_path(session_id, traces_dir)
+        self._buffer_size = buffer_size
+        self._auto_flush = auto_flush
+        
+        self._buffer: List[str] = []
+        self._lock = threading.Lock()
+        self._closed = False
+        
+        # Ensure parent directory exists
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+    
+    def emit(self, event) -> None:
+        """
+        Emit a context event.
+        
+        Args:
+            event: ContextEvent to write (must have to_dict() method)
+        """
+        if self._closed:
+            return
+        
+        # Serialize event
+        try:
+            event_dict = event.to_dict()
+            json_line = json.dumps(event_dict, default=str)
+        except Exception:
+            return
+        
+        with self._lock:
+            self._buffer.append(json_line)
+            
+            if self._auto_flush and len(self._buffer) >= self._buffer_size:
+                self._flush_buffer()
+    
+    def flush(self) -> None:
+        """Flush buffered events to disk."""
+        with self._lock:
+            self._flush_buffer()
+    
+    def _flush_buffer(self) -> None:
+        """Internal flush (must hold lock)."""
+        if not self._buffer:
+            return
+        
+        try:
+            with open(self._path, "a", encoding="utf-8") as f:
+                for line in self._buffer:
+                    f.write(line + "\n")
+            self._buffer.clear()
+        except Exception:
+            pass
+    
+    def close(self) -> None:
+        """Close the writer and flush remaining events."""
+        if self._closed:
+            return
+        
+        self.flush()
+        self._closed = True
+    
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self._session_id
+    
+    @property
+    def path(self) -> Path:
+        """Get the trace file path."""
+        return self._path
+    
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
+        return False

src/praisonai/praisonai/templates/security.py

@@ -250,8 +250,15 @@ class TemplateSecurity:
         errors = []
         total_size = 0
         
+        # Directories to skip during validation
+        SKIP_DIRS = {"__pycache__", ".git", ".svn", ".hg", "node_modules", ".venv", "venv"}
+        
         for file_path in directory.rglob("*"):
             if file_path.is_file():
+                # Skip files in excluded directories
+                if any(skip_dir in file_path.parts for skip_dir in SKIP_DIRS):
+                    continue
+                
                 # Check path safety
                 if not self.validate_path(str(file_path.relative_to(directory))):
                     errors.append(f"Unsafe path: {file_path.relative_to(directory)}")

src/praisonai/praisonai/train/__init__.py

@@ -0,0 +1,77 @@
+"""
+PraisonAI Training Module.
+
+This module provides two types of training:
+
+1. **LLM Fine-tuning** (`praisonai train dataset.json`)
+   - Fine-tune language models using Unsloth
+   - Requires GPU and heavy dependencies
+
+2. **Agent Training** (`praisonai train agents`)
+   - Iteratively improve agent behavior through feedback
+   - Two modes:
+     - LLM-as-judge (default): Automated review and optimization
+     - Human-in-the-loop (--human): Interactive feedback
+
+Usage:
+    # LLM Fine-tuning (existing)
+    praisonai train dataset.json
+    
+    # Agent Training (new)
+    praisonai train agents my_agent.yaml
+    praisonai train agents my_agent.yaml --human
+    praisonai train agents my_agent.yaml --iterations 5
+
+Programmatic Usage:
+    from praisonai.train.agents import AgentTrainer, TrainingScenario
+    
+    trainer = AgentTrainer(agent=my_agent)
+    report = trainer.run()
+"""
+
+__all__ = [
+    # Agent training
+    "AgentTrainer",
+    "TrainingScenario",
+    "TrainingIteration",
+    "TrainingReport",
+    "TrainingGrader",
+    "TrainingStorage",
+    # LLM fine-tuning (lazy)
+    "TrainModel",
+]
+
+
+def __getattr__(name: str):
+    """Lazy load training components to avoid heavy imports."""
+    # Agent training components (lightweight)
+    if name == "AgentTrainer":
+        from .agents.orchestrator import AgentTrainer
+        return AgentTrainer
+    
+    if name == "TrainingScenario":
+        from .agents.models import TrainingScenario
+        return TrainingScenario
+    
+    if name == "TrainingIteration":
+        from .agents.models import TrainingIteration
+        return TrainingIteration
+    
+    if name == "TrainingReport":
+        from .agents.models import TrainingReport
+        return TrainingReport
+    
+    if name == "TrainingGrader":
+        from .agents.grader import TrainingGrader
+        return TrainingGrader
+    
+    if name == "TrainingStorage":
+        from .agents.storage import TrainingStorage
+        return TrainingStorage
+    
+    # LLM fine-tuning (heavy - only load when needed)
+    if name == "TrainModel":
+        from .llm.trainer import TrainModel
+        return TrainModel
+    
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai/praisonai/train/agents/__init__.py

@@ -0,0 +1,69 @@
+"""
+Agent Training Module.
+
+Iteratively improve agent behavior through feedback loops.
+
+Two Modes:
+1. **LLM-as-Judge (default)**: Automated review and optimization
+   - Agent runs with input
+   - LLM grades the output (1-10 score)
+   - LLM suggests improvements
+   - Agent re-runs with improvements
+   - Repeat for N iterations
+
+2. **Human-in-the-Loop (--human flag)**: Interactive feedback
+   - Agent runs with input
+   - Human reviews and provides feedback
+   - Agent re-runs with human feedback
+   - Repeat for N iterations
+
+Usage:
+    # CLI
+    praisonai train agents my_agent.yaml
+    praisonai train agents my_agent.yaml --human
+    praisonai train agents my_agent.yaml --iterations 5
+    
+    # Python
+    from praisonai.train.agents import AgentTrainer, TrainingScenario
+    
+    trainer = AgentTrainer(agent=my_agent)
+    report = trainer.run()
+"""
+
+__all__ = [
+    "AgentTrainer",
+    "TrainingScenario",
+    "TrainingIteration",
+    "TrainingReport",
+    "TrainingGrader",
+    "TrainingStorage",
+]
+
+
+def __getattr__(name: str):
+    """Lazy load components."""
+    if name == "AgentTrainer":
+        from .orchestrator import AgentTrainer
+        return AgentTrainer
+    
+    if name == "TrainingScenario":
+        from .models import TrainingScenario
+        return TrainingScenario
+    
+    if name == "TrainingIteration":
+        from .models import TrainingIteration
+        return TrainingIteration
+    
+    if name == "TrainingReport":
+        from .models import TrainingReport
+        return TrainingReport
+    
+    if name == "TrainingGrader":
+        from .grader import TrainingGrader
+        return TrainingGrader
+    
+    if name == "TrainingStorage":
+        from .storage import TrainingStorage
+        return TrainingStorage
+    
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai/praisonai/train/agents/grader.py

@@ -0,0 +1,63 @@
+"""
+Training Grader for Agent Training.
+
+Uses LLM-as-judge to grade agent outputs and provide improvement suggestions.
+
+DRY: This module now inherits from praisonaiagents.eval.grader.BaseLLMGrader
+to share common prompt building and response parsing logic.
+"""
+
+import logging
+from typing import Optional
+
+# DRY: Import base grader and result from praisonaiagents
+from praisonaiagents.eval.grader import BaseLLMGrader, GradeResult
+
+logger = logging.getLogger(__name__)
+
+# Re-export GradeResult for backward compatibility
+__all__ = ['GradeResult', 'TrainingGrader']
+
+
+class TrainingGrader(BaseLLMGrader):
+    """
+    Grades agent outputs using LLM-as-judge.
+    
+    DRY: Inherits from BaseLLMGrader which provides common:
+    - Prompt building (_build_prompt)
+    - Response parsing (_parse_response)
+    - Lazy litellm import (_get_litellm)
+    - Sync and async grading (grade, grade_async)
+    
+    Uses an LLM to evaluate the quality of agent outputs and provide
+    improvement suggestions. This enables automated training without
+    human intervention.
+    
+    Usage:
+        grader = TrainingGrader()
+        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")
+        print(f"Suggestions: {result.suggestions}")
+    """
+    
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        temperature: float = 0.1,
+    ):
+        """
+        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)
+        """
+        # DRY: Use parent class initialization
+        super().__init__(model=model, temperature=temperature, max_tokens=500)
+    
+    # All methods (grade, grade_async, _build_prompt, _parse_response)
+    # are inherited from BaseLLMGrader - no duplication needed!

src/praisonai/praisonai/train/agents/models.py

@@ -0,0 +1,222 @@
+"""
+Data models for Agent Training.
+
+Provides dataclasses for training scenarios, iterations, and reports.
+Designed to be simple and JSON-serializable.
+"""
+
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import statistics
+
+
+@dataclass
+class TrainingScenario:
+    """
+    A scenario for agent training.
+    
+    Represents a single input/output pair that the agent should learn from.
+    
+    Attributes:
+        id: Unique identifier for this scenario
+        input_text: The input prompt to give the agent
+        expected_output: Optional expected output for comparison
+        context: Additional context or metadata
+    
+    Example:
+        scenario = TrainingScenario(
+            id="greeting-1",
+            input_text="Hello, how are you?",
+            expected_output="I'm doing well, thank you!"
+        )
+    """
+    id: str
+    input_text: str
+    expected_output: Optional[str] = None
+    context: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TrainingScenario":
+        """Create from dictionary."""
+        return cls(
+            id=data["id"],
+            input_text=data["input_text"],
+            expected_output=data.get("expected_output"),
+            context=data.get("context", {})
+        )
+
+
+@dataclass
+class TrainingIteration:
+    """
+    Data from one training iteration.
+    
+    Captures the input, output, score, and feedback for a single iteration.
+    
+    Attributes:
+        iteration_num: Which iteration this is (1-based)
+        scenario_id: ID of the scenario being trained
+        input_text: The input given to the agent
+        output: The agent's output
+        score: Quality score (1-10)
+        feedback: Feedback text (from LLM or human)
+        suggestions: List of improvement suggestions
+        timestamp: When this iteration occurred
+        metadata: Additional metadata
+    """
+    iteration_num: int
+    scenario_id: str
+    input_text: str
+    output: str
+    score: float
+    feedback: str
+    suggestions: List[str] = field(default_factory=list)
+    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat())
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TrainingIteration":
+        """Create from dictionary."""
+        return cls(
+            iteration_num=data["iteration_num"],
+            scenario_id=data["scenario_id"],
+            input_text=data["input_text"],
+            output=data["output"],
+            score=data["score"],
+            feedback=data["feedback"],
+            suggestions=data.get("suggestions", []),
+            timestamp=data.get("timestamp", datetime.utcnow().isoformat()),
+            metadata=data.get("metadata", {})
+        )
+
+
+@dataclass
+class TrainingReport:
+    """
+    Summary report of a training session.
+    
+    Contains all iterations and computed statistics.
+    
+    Attributes:
+        session_id: Unique session identifier
+        iterations: List of all training iterations
+        total_iterations: Total number of iterations run
+        started_at: When training started
+        completed_at: When training completed
+        metadata: Additional metadata
+    """
+    session_id: str
+    iterations: List[TrainingIteration]
+    total_iterations: int
+    started_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
+    completed_at: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    @property
+    def scores(self) -> List[float]:
+        """Get list of all scores."""
+        return [it.score for it in self.iterations]
+    
+    @property
+    def avg_score(self) -> float:
+        """Calculate average score across all iterations."""
+        if not self.scores:
+            return 0.0
+        return statistics.mean(self.scores)
+    
+    @property
+    def min_score(self) -> float:
+        """Get minimum score."""
+        if not self.scores:
+            return 0.0
+        return min(self.scores)
+    
+    @property
+    def max_score(self) -> float:
+        """Get maximum score."""
+        if not self.scores:
+            return 0.0
+        return max(self.scores)
+    
+    @property
+    def improvement(self) -> float:
+        """Calculate improvement from first to last iteration."""
+        if len(self.scores) < 2:
+            return 0.0
+        return self.scores[-1] - self.scores[0]
+    
+    @property
+    def passed(self) -> bool:
+        """Check if training passed (avg score >= 7)."""
+        return self.avg_score >= 7.0
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "session_id": self.session_id,
+            "iterations": [it.to_dict() for it in self.iterations],
+            "total_iterations": self.total_iterations,
+            "started_at": self.started_at,
+            "completed_at": self.completed_at,
+            "metadata": self.metadata,
+            "avg_score": self.avg_score,
+            "min_score": self.min_score,
+            "max_score": self.max_score,
+            "improvement": self.improvement,
+            "passed": self.passed,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TrainingReport":
+        """Create from dictionary."""
+        iterations = [
+            TrainingIteration.from_dict(it) for it in data.get("iterations", [])
+        ]
+        return cls(
+            session_id=data["session_id"],
+            iterations=iterations,
+            total_iterations=data["total_iterations"],
+            started_at=data.get("started_at", datetime.utcnow().isoformat()),
+            completed_at=data.get("completed_at"),
+            metadata=data.get("metadata", {})
+        )
+    
+    def print_summary(self) -> None:
+        """Print a summary of the training results."""
+        try:
+            from rich.console import Console
+            from rich.table import Table
+            
+            console = Console()
+            
+            table = Table(title="Agent Training Summary")
+            table.add_column("Metric", style="cyan")
+            table.add_column("Value", style="green")
+            
+            table.add_row("Session ID", self.session_id)
+            table.add_row("Total Iterations", str(self.total_iterations))
+            table.add_row("Average Score", f"{self.avg_score:.2f}")
+            table.add_row("Min Score", f"{self.min_score:.2f}")
+            table.add_row("Max Score", f"{self.max_score:.2f}")
+            table.add_row("Improvement", f"{self.improvement:+.2f}")
+            table.add_row("Status", "✅ PASSED" if self.passed else "❌ NEEDS WORK")
+            
+            console.print(table)
+        except ImportError:
+            print("Agent Training Summary")
+            print(f"  Session ID: {self.session_id}")
+            print(f"  Total Iterations: {self.total_iterations}")
+            print(f"  Average Score: {self.avg_score:.2f}")
+            print(f"  Min Score: {self.min_score:.2f}")
+            print(f"  Max Score: {self.max_score:.2f}")
+            print(f"  Improvement: {self.improvement:+.2f}")
+            print(f"  Status: {'PASSED' if self.passed else 'NEEDS WORK'}")

src/praisonai/praisonai/train/agents/orchestrator.py

@@ -0,0 +1,394 @@
+"""
+Agent Training Orchestrator.
+
+Orchestrates the training loop for iteratively improving agent behavior.
+Supports two modes:
+1. LLM-as-Judge (default): Automated review and optimization
+2. Human-in-the-Loop: Interactive feedback from humans
+
+Usage:
+    # LLM mode (default)
+    trainer = AgentTrainer(agent=my_agent)
+    trainer.add_scenario(TrainingScenario(id="s1", input_text="Hello"))
+    report = trainer.run()
+    
+    # Human mode
+    trainer = AgentTrainer(agent=my_agent, human_mode=True)
+    trainer.add_scenario(TrainingScenario(id="s1", input_text="Hello"))
+    report = trainer.run()
+"""
+
+import logging
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union, TYPE_CHECKING
+
+from .grader import TrainingGrader, GradeResult
+from .models import TrainingIteration, TrainingReport, TrainingScenario
+from .storage import TrainingStorage
+
+if TYPE_CHECKING:
+    from praisonaiagents import Agent, Agents
+
+logger = logging.getLogger(__name__)
+
+
+class AgentTrainer:
+    """
+    Orchestrates agent training through iterative feedback loops.
+    
+    Two modes are supported:
+    
+    1. **LLM-as-Judge (default)**: The agent runs, an LLM grades the output,
+       provides suggestions, and the agent re-runs with the feedback.
+       This is fully automated and requires no human intervention.
+    
+    2. **Human-in-the-Loop**: The agent runs, a human reviews the output
+       and provides feedback, then the agent re-runs with the feedback.
+       This requires interactive input from the user.
+    
+    Example:
+        from praisonai.train.agents import AgentTrainer, TrainingScenario
+        from praisonaiagents import Agent
+        
+        agent = Agent(instructions="You are a helpful assistant")
+        
+        trainer = AgentTrainer(agent=agent, iterations=3)
+        trainer.add_scenario(TrainingScenario(
+            id="greeting",
+            input_text="Hello, how are you?",
+            expected_output="I'm doing well, thank you for asking!"
+        ))
+        
+        report = trainer.run()
+        report.print_summary()
+    """
+    
+    def __init__(
+        self,
+        agent: Union["Agent", "Agents", Callable[[str], str]],
+        iterations: int = 3,
+        human_mode: bool = False,
+        grader: Optional[TrainingGrader] = None,
+        storage_dir: Optional[Path] = None,
+        storage_backend: Optional[Any] = None,
+        verbose: bool = True,
+    ):
+        """
+        Initialize the trainer.
+        
+        Args:
+            agent: Agent, Agents, or callable to train
+            iterations: Number of training iterations (default: 3)
+            human_mode: Use human feedback instead of LLM grading (default: False)
+            grader: Custom grader instance (default: TrainingGrader())
+            storage_dir: Directory for storing training data
+            storage_backend: Storage backend (FileBackend, SQLiteBackend, RedisBackend)
+            verbose: Print progress during training (default: True)
+        """
+        self.agent = agent
+        self.iterations = iterations
+        self.human_mode = human_mode
+        self.grader = grader
+        self.storage_dir = storage_dir
+        self.storage_backend = storage_backend
+        self.verbose = verbose
+        
+        self.session_id = f"train-{uuid.uuid4().hex[:8]}"
+        self.scenarios: List[TrainingScenario] = []
+        
+        self._storage: Optional[TrainingStorage] = None
+    
+    def add_scenario(
+        self,
+        scenario: Union[TrainingScenario, Dict[str, Any]],
+    ) -> None:
+        """
+        Add a training scenario.
+        
+        Args:
+            scenario: TrainingScenario or dict with scenario data
+        """
+        if isinstance(scenario, dict):
+            scenario = TrainingScenario.from_dict(scenario)
+        self.scenarios.append(scenario)
+    
+    def add_scenarios(
+        self,
+        scenarios: List[Union[TrainingScenario, Dict[str, Any]]],
+    ) -> None:
+        """
+        Add multiple training scenarios.
+        
+        Args:
+            scenarios: List of TrainingScenario or dicts
+        """
+        for scenario in scenarios:
+            self.add_scenario(scenario)
+    
+    def _get_agent_output(self, prompt: str) -> str:
+        """
+        Get output from the agent.
+        
+        Handles different agent types (Agent, Agents, callable).
+        
+        Args:
+            prompt: Input prompt
+            
+        Returns:
+            Agent's output string
+        """
+        # Callable function
+        if callable(self.agent) and not hasattr(self.agent, 'chat') and not hasattr(self.agent, 'start'):
+            return str(self.agent(prompt))
+        
+        # Agent with chat method
+        if hasattr(self.agent, 'chat'):
+            result = self.agent.chat(prompt)
+            return str(result)
+        
+        # Agents with start method
+        if hasattr(self.agent, 'start'):
+            result = self.agent.start(prompt)
+            if hasattr(result, 'raw'):
+                return str(result.raw)
+            return str(result)
+        
+        raise ValueError(
+            "Agent must have 'chat' method, 'start' method, or be callable"
+        )
+    
+    def _get_human_feedback(self, output: str, scenario: TrainingScenario) -> GradeResult:
+        """
+        Get feedback from human.
+        
+        Args:
+            output: Agent's output to review
+            scenario: The scenario being trained
+            
+        Returns:
+            GradeResult with human feedback
+        """
+        print("\n" + "=" * 60)
+        print("AGENT OUTPUT:")
+        print("-" * 60)
+        print(output)
+        print("-" * 60)
+        
+        if scenario.expected_output:
+            print("\nEXPECTED OUTPUT:")
+            print(scenario.expected_output)
+            print("-" * 60)
+        
+        # Get score
+        while True:
+            try:
+                score_input = input("\nScore (1-10): ").strip()
+                score = float(score_input)
+                if 1 <= score <= 10:
+                    break
+                print("Please enter a number between 1 and 10")
+            except ValueError:
+                print("Please enter a valid number")
+        
+        # Get feedback
+        feedback = input("Feedback (what could be improved?): ").strip()
+        if not feedback:
+            feedback = "No specific feedback provided"
+        
+        # Get suggestions
+        suggestions_input = input("Suggestions (comma-separated, or press Enter to skip): ").strip()
+        suggestions = []
+        if suggestions_input:
+            suggestions = [s.strip() for s in suggestions_input.split(",") if s.strip()]
+        
+        return GradeResult(
+            score=score,
+            reasoning=feedback,
+            suggestions=suggestions,
+            input_text=scenario.input_text,
+            output=output,
+            expected_output=scenario.expected_output,
+        )
+    
+    def _build_improvement_prompt(
+        self,
+        original_prompt: str,
+        previous_output: str,
+        grade_result: GradeResult,
+    ) -> str:
+        """
+        Build prompt with improvement feedback.
+        
+        Args:
+            original_prompt: Original input
+            previous_output: Previous agent output
+            grade_result: Grading result with feedback
+            
+        Returns:
+            Enhanced prompt with feedback
+        """
+        prompt = original_prompt
+        
+        if grade_result.suggestions:
+            suggestions_text = "\n".join(f"- {s}" for s in grade_result.suggestions)
+            prompt += "\n\n[FEEDBACK FROM PREVIOUS ATTEMPT]\n"
+            prompt += f"Score: {grade_result.score}/10\n"
+            prompt += f"Feedback: {grade_result.reasoning}\n"
+            prompt += f"Suggestions for improvement:\n{suggestions_text}\n"
+            prompt += "\nPlease provide an improved response addressing this feedback."
+        
+        return prompt
+    
+    def run(self) -> TrainingReport:
+        """
+        Run the training loop.
+        
+        Returns:
+            TrainingReport with all iterations and statistics
+            
+        Raises:
+            ValueError: If no scenarios have been added
+        """
+        if not self.scenarios:
+            raise ValueError(
+                "No scenarios added. Use add_scenario() to add training scenarios."
+            )
+        
+        # Initialize storage
+        self._storage = TrainingStorage(
+            session_id=self.session_id,
+            storage_dir=self.storage_dir,
+            backend=self.storage_backend,
+        )
+        
+        # Initialize grader for LLM mode
+        if not self.human_mode and self.grader is None:
+            self.grader = TrainingGrader()
+        
+        # Save scenarios
+        for scenario in self.scenarios:
+            self._storage.save_scenario(scenario)
+        
+        all_iterations: List[TrainingIteration] = []
+        started_at = datetime.utcnow().isoformat()
+        
+        if self.verbose:
+            mode = "Human-in-the-Loop" if self.human_mode else "LLM-as-Judge"
+            print("\n" + "="*60)
+            print("Starting Agent Training")
+            print(f"Mode: {mode}")
+            print(f"Scenarios: {len(self.scenarios)}")
+            print(f"Iterations: {self.iterations}")
+            print(f"Session: {self.session_id}")
+            print("="*60 + "\n")
+        
+        for scenario in self.scenarios:
+            if self.verbose:
+                print(f"\n--- Scenario: {scenario.id} ---")
+            
+            current_prompt = scenario.input_text
+            previous_output = None
+            previous_grade = None
+            
+            for iteration_num in range(1, self.iterations + 1):
+                if self.verbose:
+                    print(f"\nIteration {iteration_num}/{self.iterations}")
+                
+                # Build prompt with feedback if not first iteration
+                if previous_grade is not None and previous_output is not None:
+                    current_prompt = self._build_improvement_prompt(
+                        scenario.input_text,
+                        previous_output,
+                        previous_grade,
+                    )
+                
+                # Get agent output
+                try:
+                    output = self._get_agent_output(current_prompt)
+                except Exception as e:
+                    logger.error(f"Agent error: {e}")
+                    output = f"Error: {str(e)}"
+                
+                # Grade the output
+                if self.human_mode:
+                    grade_result = self._get_human_feedback(output, scenario)
+                else:
+                    grade_result = self.grader.grade(
+                        input_text=scenario.input_text,
+                        output=output,
+                        expected_output=scenario.expected_output,
+                    )
+                
+                if self.verbose:
+                    print(f"  Score: {grade_result.score}/10")
+                    if grade_result.suggestions:
+                        print(f"  Suggestions: {', '.join(grade_result.suggestions[:2])}")
+                
+                # Create iteration record
+                iteration = TrainingIteration(
+                    iteration_num=iteration_num,
+                    scenario_id=scenario.id,
+                    input_text=scenario.input_text,
+                    output=output,
+                    score=grade_result.score,
+                    feedback=grade_result.reasoning,
+                    suggestions=grade_result.suggestions,
+                    metadata={
+                        "mode": "human" if self.human_mode else "llm",
+                        "prompt_used": current_prompt,
+                    }
+                )
+                
+                all_iterations.append(iteration)
+                self._storage.save_iteration(iteration)
+                
+                # Update for next iteration
+                previous_output = output
+                previous_grade = grade_result
+                
+                # Early stop if score is high enough
+                if grade_result.score >= 9.5:
+                    if self.verbose:
+                        print("  ✓ Excellent score achieved, stopping early")
+                    break
+        
+        # Create report
+        report = TrainingReport(
+            session_id=self.session_id,
+            iterations=all_iterations,
+            total_iterations=len(all_iterations),
+            started_at=started_at,
+            completed_at=datetime.utcnow().isoformat(),
+            metadata={
+                "mode": "human" if self.human_mode else "llm",
+                "scenarios_count": len(self.scenarios),
+                "target_iterations": self.iterations,
+            }
+        )
+        
+        self._storage.save_report(report)
+        
+        if self.verbose:
+            print("\n" + "="*60)
+            print("Training Complete!")
+            report.print_summary()
+        
+        return report
+    
+    async def run_async(self) -> TrainingReport:
+        """
+        Run the training loop asynchronously.
+        
+        Note: Human mode is not supported in async.
+        
+        Returns:
+            TrainingReport with all iterations and statistics
+        """
+        if self.human_mode:
+            raise ValueError("Human mode is not supported in async training")
+        
+        # For now, delegate to sync version
+        # TODO: Implement true async with acompletion
+        return self.run()

src/praisonai/praisonai/train/agents/storage.py

@@ -0,0 +1,282 @@
+"""
+Storage for Agent Training data.
+
+Provides JSON-based persistence for training iterations, scenarios, and reports.
+DRY: Follows the same pattern as praisonaiagents.memory.learn.stores.BaseStore.
+"""
+
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional, TYPE_CHECKING
+
+# DRY: Import base classes from praisonaiagents.storage
+from praisonaiagents.storage.models import BaseSessionInfo
+from praisonaiagents.storage.base import BaseJSONStore, list_json_sessions, cleanup_old_sessions as _cleanup_old_sessions
+
+if TYPE_CHECKING:
+    from praisonaiagents.storage.protocols import StorageBackendProtocol
+
+from .models import TrainingIteration, TrainingReport, TrainingScenario
+
+
+DEFAULT_STORAGE_DIR = "~/.praison/train"
+
+
+def get_storage_dir() -> Path:
+    """
+    Get the training storage directory.
+    
+    Returns:
+        Path to storage directory (creates if not exists)
+    """
+    storage_dir = Path(os.path.expanduser(DEFAULT_STORAGE_DIR))
+    storage_dir.mkdir(parents=True, exist_ok=True)
+    return storage_dir
+
+
+class TrainingSessionInfo(BaseSessionInfo):
+    """
+    Information about a training session.
+    
+    DRY: Inherits from BaseSessionInfo which provides:
+    - session_id, path, size_bytes, created_at, modified_at, item_count
+    - to_dict(), from_dict(), from_path() methods
+    """
+    
+    @property
+    def iteration_count(self) -> int:
+        """Alias for item_count for backward compatibility."""
+        return self.item_count
+    
+    def to_dict(self):
+        """Override to include iteration_count for backward compatibility."""
+        d = super().to_dict()
+        d["iteration_count"] = self.iteration_count
+        return d
+
+
+class TrainingStorage:
+    """
+    JSON-based storage for training data.
+    
+    Stores training iterations, scenarios, and reports to JSON files.
+    Each training session gets its own file.
+    
+    DRY: Uses BaseJSONStore internally for thread-safe, file-locked storage.
+    Supports pluggable backends (file, sqlite, etc.) via the backend parameter.
+    
+    Usage:
+        storage = TrainingStorage(session_id="train-abc123")
+        storage.save_iteration(iteration)
+        iterations = storage.load_iterations()
+        
+        # With SQLite backend:
+        from praisonaiagents.storage import SQLiteBackend
+        storage = TrainingStorage(
+            session_id="train-abc123",
+            backend=SQLiteBackend("training.db")
+        )
+    """
+    
+    def __init__(
+        self,
+        session_id: str,
+        storage_dir: Optional[Path] = None,
+        backend: Optional["StorageBackendProtocol"] = None,
+    ):
+        """
+        Initialize storage.
+        
+        Args:
+            session_id: Unique session identifier
+            storage_dir: Directory for storage files (default: ~/.praison/train)
+            backend: Optional storage backend (file, sqlite, etc.)
+                     If provided, storage_dir is ignored and backend is used.
+        """
+        self.session_id = session_id
+        self._backend = backend
+        self.storage_dir = storage_dir or get_storage_dir()
+        self.storage_dir = Path(self.storage_dir)
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+        
+        # DRY: Use BaseJSONStore for thread-safe storage
+        self._store = BaseJSONStore(
+            storage_path=self.storage_path,
+            backend=backend,
+        )
+        
+        self._data = self._default_data()
+        
+        # Load existing data if available
+        if self._store.exists():
+            self._load()
+    
+    def _default_data(self):
+        """Return default data structure."""
+        return {
+            "session_id": self.session_id,
+            "created_at": datetime.utcnow().isoformat(),
+            "scenarios": [],
+            "iterations": [],
+            "report": None,
+        }
+    
+    @property
+    def storage_path(self) -> Path:
+        """Get path to storage file."""
+        return self.storage_dir / f"{self.session_id}.json"
+    
+    def _load(self) -> None:
+        """Load data from storage using BaseJSONStore."""
+        loaded = self._store.load()
+        if loaded:
+            self._data = loaded
+    
+    def _save(self) -> None:
+        """Save data to storage using BaseJSONStore."""
+        self._data["updated_at"] = datetime.utcnow().isoformat()
+        self._store.save(self._data)
+    
+    def save_iteration(self, iteration: TrainingIteration) -> None:
+        """
+        Save a training iteration.
+        
+        Args:
+            iteration: The iteration to save
+        """
+        self._data["iterations"].append(iteration.to_dict())
+        self._save()
+    
+    def load_iterations(self) -> List[TrainingIteration]:
+        """
+        Load all iterations from storage.
+        
+        Returns:
+            List of TrainingIteration objects
+        """
+        return [
+            TrainingIteration.from_dict(it)
+            for it in self._data.get("iterations", [])
+        ]
+    
+    def save_scenario(self, scenario: TrainingScenario) -> None:
+        """
+        Save a training scenario.
+        
+        Args:
+            scenario: The scenario to save
+        """
+        self._data["scenarios"].append(scenario.to_dict())
+        self._save()
+    
+    def load_scenarios(self) -> List[TrainingScenario]:
+        """
+        Load all scenarios from storage.
+        
+        Returns:
+            List of TrainingScenario objects
+        """
+        return [
+            TrainingScenario.from_dict(s)
+            for s in self._data.get("scenarios", [])
+        ]
+    
+    def save_report(self, report: TrainingReport) -> None:
+        """
+        Save the training report.
+        
+        Args:
+            report: The report to save
+        """
+        self._data["report"] = report.to_dict()
+        self._save()
+    
+    def load_report(self) -> Optional[TrainingReport]:
+        """
+        Load the training report.
+        
+        Returns:
+            TrainingReport or None if not saved
+        """
+        report_data = self._data.get("report")
+        if report_data:
+            return TrainingReport.from_dict(report_data)
+        return None
+    
+    def clear(self) -> None:
+        """Clear all stored data."""
+        self._data = {
+            "session_id": self.session_id,
+            "created_at": datetime.utcnow().isoformat(),
+            "scenarios": [],
+            "iterations": [],
+            "report": None,
+        }
+        self._save()
+
+
+def list_training_sessions(
+    storage_dir: Optional[Path] = None,
+    limit: int = 50,
+) -> List[TrainingSessionInfo]:
+    """
+    List all training sessions.
+    
+    DRY: Uses list_json_sessions from praisonaiagents.storage.base.
+    
+    Args:
+        storage_dir: Directory to search (default: ~/.praison/train)
+        limit: Maximum number of sessions to return
+        
+    Returns:
+        List of TrainingSessionInfo objects, sorted by modification time (newest first)
+    """
+    if storage_dir is None:
+        storage_dir = get_storage_dir()
+    
+    # DRY: Use common list_json_sessions function
+    base_sessions = list_json_sessions(Path(storage_dir), suffix=".json", limit=limit)
+    
+    # Convert BaseSessionInfo to TrainingSessionInfo for backward compatibility
+    return [
+        TrainingSessionInfo(
+            session_id=s.session_id,
+            path=s.path,
+            size_bytes=s.size_bytes,
+            created_at=s.created_at,
+            modified_at=s.modified_at,
+            item_count=s.item_count,
+        )
+        for s in base_sessions
+    ]
+
+
+def cleanup_old_sessions(
+    storage_dir: Optional[Path] = None,
+    max_age_days: int = 30,
+    max_size_mb: int = 100,
+) -> int:
+    """
+    Clean up old training sessions.
+    
+    DRY: Uses cleanup_old_sessions from praisonaiagents.storage.base.
+    
+    Args:
+        storage_dir: Directory to clean (default: ~/.praison/train)
+        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 storage_dir is None:
+        storage_dir = get_storage_dir()
+    
+    # DRY: Use common cleanup function
+    return _cleanup_old_sessions(
+        storage_dir=Path(storage_dir),
+        suffix=".json",
+        max_age_days=max_age_days,
+        max_size_mb=max_size_mb,
+    )

src/praisonai/praisonai/train/llm/__init__.py

@@ -0,0 +1,23 @@
+"""
+LLM Fine-tuning Module.
+
+This module provides LLM fine-tuning using Unsloth's fast training framework.
+Supports ShareGPT and Alpaca-style datasets.
+
+Usage:
+    praisonai train dataset.json
+    praisonai train --model llama-3.1 dataset.json
+
+Note: This module requires heavy dependencies (torch, unsloth, transformers).
+They are only loaded when this module is explicitly imported.
+"""
+
+__all__ = ["TrainModel"]
+
+
+def __getattr__(name: str):
+    """Lazy load to avoid heavy imports."""
+    if name == "TrainModel":
+        from .trainer import TrainModel
+        return TrainModel
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/praisonai/praisonai/train/llm/trainer.py

@@ -0,0 +1,562 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+This script finetunes a model using Unsloth's fast training framework.
+It supports both ShareGPT and Alpaca‑style datasets by converting raw conversation
+data into plain-text prompts using a chat template, then pre‑tokenizing the prompts.
+Extra debug logging is added to help trace the root cause of errors.
+"""
+
+import os
+import sys
+import yaml
+import torch
+import shutil
+import subprocess
+from transformers import TextStreamer
+from unsloth import FastLanguageModel, is_bfloat16_supported
+from trl import SFTTrainer
+from transformers import TrainingArguments
+from datasets import load_dataset, concatenate_datasets
+from psutil import virtual_memory
+from unsloth.chat_templates import standardize_sharegpt, get_chat_template
+from functools import partial
+
+#####################################
+# Step 1: Formatting Raw Conversations
+#####################################
+def formatting_prompts_func(examples, tokenizer):
+    """
+    Converts each example's conversation into a single plain-text prompt.
+    If the example has a "conversations" field, process it as ShareGPT-style.
+    Otherwise, assume Alpaca-style data with "instruction", "input", and "output" fields.
+    """
+    print("DEBUG: formatting_prompts_func() received batch with keys:", list(examples.keys()))
+    texts = []
+    # Check if the example has a "conversations" field.
+    if "conversations" in examples:
+        for convo in examples["conversations"]:
+            try:
+                formatted = tokenizer.apply_chat_template(
+                    convo,
+                    tokenize=False,  # Return a plain string
+                    add_generation_prompt=False
+                )
+            except Exception as e:
+                print(f"ERROR in apply_chat_template (conversations): {e}")
+                formatted = ""
+            # Flatten list if necessary
+            if isinstance(formatted, list):
+                formatted = formatted[0] if len(formatted) == 1 else "\n".join(formatted)
+            texts.append(formatted)
+    else:
+        # Assume Alpaca format: use "instruction", "input", and "output" keys.
+        instructions = examples.get("instruction", [])
+        inputs_list = examples.get("input", [])
+        outputs_list = examples.get("output", [])
+        # If any field is missing, replace with empty string.
+        for ins, inp, out in zip(instructions, inputs_list, outputs_list):
+            # Create a conversation-like structure.
+            convo = [
+                {"role": "user", "content": ins + (f"\nInput: {inp}" if inp.strip() != "" else "")},
+                {"role": "assistant", "content": out}
+            ]
+            try:
+                formatted = tokenizer.apply_chat_template(
+                    convo,
+                    tokenize=False,
+                    add_generation_prompt=False
+                )
+            except Exception as e:
+                print(f"ERROR in apply_chat_template (alpaca): {e}")
+                formatted = ""
+            if isinstance(formatted, list):
+                formatted = formatted[0] if len(formatted) == 1 else "\n".join(formatted)
+            texts.append(formatted)
+    if texts:
+        print("DEBUG: Raw texts sample (first 200 chars):", texts[0][:200])
+    return {"text": texts}
+
+#####################################
+# Step 2: Tokenizing the Prompts
+#####################################
+def tokenize_function(examples, hf_tokenizer, max_length):
+    """
+    Tokenizes a batch of text prompts with padding and truncation enabled.
+    """
+    flat_texts = []
+    for t in examples["text"]:
+        if isinstance(t, list):
+            t = t[0] if len(t) == 1 else " ".join(t)
+        flat_texts.append(t)
+    print("DEBUG: Tokenizing a batch of size:", len(flat_texts))
+    tokenized = hf_tokenizer(
+        flat_texts,
+        padding="max_length",
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt",
+    )
+    tokenized = {key: value.tolist() for key, value in tokenized.items()}
+    sample_key = list(tokenized.keys())[0]
+    print("DEBUG: Tokenized sample (first 10 tokens of", sample_key, "):", tokenized[sample_key][0][:10])
+    return tokenized
+
+#####################################
+# Main Training Class
+#####################################
+class TrainModel:
+    def __init__(self, config_path="config.yaml"):
+        self.load_config(config_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.hf_tokenizer = None   # The underlying HF tokenizer
+        self.chat_tokenizer = None # Chat wrapper for formatting
+
+    def load_config(self, path):
+        with open(path, "r") as file:
+            self.config = yaml.safe_load(file)
+        print("DEBUG: Loaded config:", self.config)
+
+    def print_system_info(self):
+        print("DEBUG: PyTorch version:", torch.__version__)
+        print("DEBUG: CUDA version:", torch.version.cuda)
+        if torch.cuda.is_available():
+            print("DEBUG: CUDA Device Capability:", torch.cuda.get_device_capability())
+        else:
+            print("DEBUG: CUDA is not available")
+        print("DEBUG: Python Version:", sys.version)
+        print("DEBUG: Python Path:", sys.executable)
+
+    def check_gpu(self):
+        gpu_stats = torch.cuda.get_device_properties(0)
+        print(f"DEBUG: GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory/(1024**3),3)} GB.")
+
+    def check_ram(self):
+        ram_gb = virtual_memory().total / 1e9
+        print(f"DEBUG: Your runtime has {ram_gb:.1f} gigabytes of available RAM")
+        if ram_gb < 20:
+            print("DEBUG: Not using a high-RAM runtime")
+        else:
+            print("DEBUG: You are using a high-RAM runtime!")
+
+    def prepare_model(self):
+        print("DEBUG: Preparing model and tokenizer...")
+        self.model, original_tokenizer = FastLanguageModel.from_pretrained(
+            model_name=self.config["model_name"],
+            max_seq_length=self.config["max_seq_length"],
+            dtype=None,
+            load_in_4bit=self.config["load_in_4bit"],
+        )
+        print("DEBUG: Model and original tokenizer loaded.")
+        if original_tokenizer.pad_token is None:
+            original_tokenizer.pad_token = original_tokenizer.eos_token
+        original_tokenizer.model_max_length = self.config["max_seq_length"]
+        self.chat_tokenizer = get_chat_template(original_tokenizer, chat_template="llama-3.1")
+        self.hf_tokenizer = original_tokenizer
+        print("DEBUG: Chat tokenizer created; HF tokenizer saved.")
+        self.model = FastLanguageModel.get_peft_model(
+            self.model,
+            r=16,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+            lora_alpha=16,
+            lora_dropout=0,
+            bias="none",
+            use_gradient_checkpointing="unsloth",
+            random_state=3407,
+            use_rslora=False,
+            loftq_config=None,
+        )
+        print("DEBUG: LoRA adapters added.")
+
+    def process_dataset(self, dataset_info):
+        dataset_name = dataset_info["name"]
+        split_type = dataset_info.get("split_type", "train")
+        print(f"DEBUG: Loading dataset '{dataset_name}' split '{split_type}'...")
+        dataset = load_dataset(dataset_name, split=split_type)
+        print("DEBUG: Dataset columns:", dataset.column_names)
+        if "conversations" in dataset.column_names:
+            print("DEBUG: Standardizing dataset (ShareGPT style)...")
+            dataset = standardize_sharegpt(dataset)
+        else:
+            print("DEBUG: Dataset does not have 'conversations'; assuming Alpaca format.")
+        print("DEBUG: Applying formatting function to dataset...")
+        format_func = partial(formatting_prompts_func, tokenizer=self.chat_tokenizer)
+        dataset = dataset.map(format_func, batched=True, remove_columns=dataset.column_names)
+        sample = dataset[0]
+        print("DEBUG: Sample processed example keys:", list(sample.keys()))
+        if "text" in sample:
+            print("DEBUG: Sample processed 'text' type:", type(sample["text"]))
+            print("DEBUG: Sample processed 'text' content (first 200 chars):", sample["text"][:200])
+        else:
+            print("DEBUG: Processed sample does not contain 'text'.")
+        return dataset
+
+    def tokenize_dataset(self, dataset):
+        print("DEBUG: Tokenizing the entire dataset...")
+        tokenized_dataset = dataset.map(
+            lambda examples: tokenize_function(examples, self.hf_tokenizer, self.config["max_seq_length"]),
+            batched=True
+        )
+        tokenized_dataset = tokenized_dataset.remove_columns(["text"])
+        print("DEBUG: Tokenized dataset sample keys:", tokenized_dataset[0].keys())
+        return tokenized_dataset
+
+    def load_datasets(self):
+        datasets = []
+        for dataset_info in self.config["dataset"]:
+            print("DEBUG: Processing dataset info:", dataset_info)
+            datasets.append(self.process_dataset(dataset_info))
+        combined = concatenate_datasets(datasets)
+        print("DEBUG: Combined dataset has", len(combined), "examples.")
+        return combined
+
+    def train_model(self):
+        print("DEBUG: Starting training...")
+        raw_dataset = self.load_datasets()
+        tokenized_dataset = self.tokenize_dataset(raw_dataset)
+        print("DEBUG: Dataset tokenization complete.")
+        # Build the training arguments parameters dynamically
+        ta_params = {
+            "per_device_train_batch_size": self.config.get("per_device_train_batch_size", 2),
+            "gradient_accumulation_steps": self.config.get("gradient_accumulation_steps", 2),
+            "warmup_steps": self.config.get("warmup_steps", 50),
+            "max_steps": self.config.get("max_steps", 2800),
+            "learning_rate": self.config.get("learning_rate", 2e-4),
+            "fp16": self.config.get("fp16", not is_bfloat16_supported()),
+            "bf16": self.config.get("bf16", is_bfloat16_supported()),
+            "logging_steps": self.config.get("logging_steps", 15),
+            "optim": self.config.get("optim", "adamw_8bit"),
+            "weight_decay": self.config.get("weight_decay", 0.01),
+            "lr_scheduler_type": self.config.get("lr_scheduler_type", "linear"),
+            "seed": self.config.get("seed", 3407),
+            "output_dir": self.config.get("output_dir", "outputs"),
+            "report_to": "none" if not os.getenv("PRAISON_WANDB") else "wandb",
+            "remove_unused_columns": self.config.get("remove_unused_columns", False)
+        }
+        if os.getenv("PRAISON_WANDB"):
+            ta_params["save_steps"] = self.config.get("save_steps", 100)
+            ta_params["run_name"] = os.getenv("PRAISON_WANDB_RUN_NAME", "praisonai-train")
+
+        training_args = TrainingArguments(**ta_params)
+        # Since the dataset is pre-tokenized, we supply a dummy dataset_text_field.
+        trainer = SFTTrainer(
+            model=self.model,
+            tokenizer=self.hf_tokenizer,
+            train_dataset=tokenized_dataset,
+            dataset_text_field="input_ids",  # Dummy field since data is numeric
+            max_seq_length=self.config["max_seq_length"],
+            dataset_num_proc=1,  # Use a single process to avoid pickling issues
+            packing=False,
+            args=training_args,
+        )
+        from unsloth.chat_templates import train_on_responses_only
+        trainer = train_on_responses_only(
+            trainer,
+            instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
+            response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
+        )
+        print("DEBUG: Beginning trainer.train() ...")
+        trainer.train()
+        print("DEBUG: Training complete. Saving model and tokenizer locally...")
+        self.model.save_pretrained("lora_model")
+        self.hf_tokenizer.save_pretrained("lora_model")
+        print("DEBUG: Saved model and tokenizer to 'lora_model'.")
+
+    def inference(self, instruction, input_text):
+        FastLanguageModel.for_inference(self.model)
+        messages = [{"role": "user", "content": f"{instruction}\n\nInput: {input_text}"}]
+        inputs = self.hf_tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to("cuda")
+        outputs = self.model.generate(
+            input_ids=inputs,
+            max_new_tokens=64,
+            use_cache=True,
+            temperature=1.5,
+            min_p=0.1
+        )
+        print("DEBUG: Inference output:", self.hf_tokenizer.batch_decode(outputs))
+
+    def load_model(self):
+        from unsloth import FastLanguageModel
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=self.config["output_dir"],
+            max_seq_length=2048,
+            dtype=None,
+            load_in_4bit=self.config["load_in_4bit"],
+        )
+        return model, tokenizer
+
+    def save_model_merged(self):
+        if os.path.exists(self.config["hf_model_name"]):
+            shutil.rmtree(self.config["hf_model_name"])
+        self.model.push_to_hub_merged(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            save_method="merged_16bit",
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def push_model_gguf(self):
+        self.model.push_to_hub_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method=self.config["quantization_method"],
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def save_model_gguf(self):
+        self.model.save_pretrained_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method="q4_k_m"
+        )
+
+    def prepare_modelfile_content(self):
+        output_model = self.config["hf_model_name"]
+        model_name = self.config["model_name"].lower()
+        # Mapping from model name keywords to their default TEMPLATE and stop tokens (and optional SYSTEM/num_ctx)
+        mapping = {
+            "llama": {
+                "template": """<|start_header_id|>system<|end_header_id|>
+    Cutting Knowledge Date: December 2023
+    {{ if .System }}{{ .System }}
+    {{- end }}
+    {{- if .Tools }}When you receive a tool call response, use the output to format an answer to the orginal user question.
+    You are a helpful assistant with tool calling capabilities.
+    {{- end }}<|eot_id|>
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 }}
+    {{- if eq .Role "user" }}<|start_header_id|>user<|end_header_id|>
+    {{- if and $.Tools $last }}
+    Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+    Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. Do not use variables.
+    {{ range $.Tools }}
+    {{- . }}
+    {{ end }}
+    {{ .Content }}<|eot_id|>
+    {{- else }}
+    {{ .Content }}<|eot_id|>
+    {{- end }}{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- else if eq .Role "assistant" }}<|start_header_id|>assistant<|end_header_id|>
+    {{- if .ToolCalls }}
+    {{ range .ToolCalls }}
+    {"name": "{{ .Function.Name }}", "parameters": {{ .Function.Arguments }}}{{ end }}
+    {{- else }}
+    {{ .Content }}
+    {{- end }}{{ if not $last }}<|eot_id|>{{ end }}
+    {{- else if eq .Role "tool" }}<|start_header_id|>ipython<|end_header_id|>
+    {{ .Content }}<|eot_id|>{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- end }}
+    {{- end }}""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            },
+            "qwen": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "system": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+                "num_ctx": 32768,
+                "stop_tokens": ["<|endoftext|>"]
+            },
+            "mistral": {
+                "template": "[INST] {{ if .System }}{{ .System }} {{ end }}{{ .Prompt }} [/INST]",
+                "stop_tokens": ["[INST]", "[/INST]"]
+            },
+            "phi": {
+                "template": """{{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    <|im_start|>{{ .Role }}<|im_sep|>
+    {{ .Content }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_end|>
+    <|im_start|>assistant<|im_sep|>
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["<|im_start|>", "<|im_end|>", "<|im_sep|>"]
+            },
+            "deepseek": {
+                "template": """{{- if .System }}{{ .System }}{{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1}}
+    {{- if eq .Role "user" }}
+    {{ .Content }}
+    {{- else if eq .Role "assistant" }}
+    {{ .Content }}{{- if not $last }}
+    {{- end }}
+    {{- end }}
+    {{- if and $last (ne .Role "assistant") }}
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["", "", "", ""]
+            },
+            "llava": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "stop_tokens": ["</s>", "USER:", "ASSSISTANT:"]
+            }
+        }
+        # Select mapping by checking if any key is in the model_name.
+        chosen = None
+        for key, settings in mapping.items():
+            if key in model_name:
+                chosen = settings
+                break
+        if chosen is None:
+            # Fallback default
+            chosen = {
+                "template": """{{ if .System }}<|start_header_id|>system<|end_header_id|>
+    {{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>
+    {{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>
+    {{ .Response }}<|eot_id|>""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            }
+        # Build the stop parameter lines.
+        stop_params = "\n".join([f"PARAMETER stop {token}" for token in chosen["stop_tokens"]])
+        # Optionally include a SYSTEM line and num_ctx if defined in the mapping.
+        system_line = ""
+        if "system" in chosen:
+            system_line = f"SYSTEM {chosen['system']}\n"
+        num_ctx_line = ""
+        if "num_ctx" in chosen:
+            num_ctx_line = f"PARAMETER num_ctx {chosen['num_ctx']}\n"
+        # Assemble and return the modelfile content.
+        return f"""FROM {output_model}
+    TEMPLATE \"\"\"{chosen['template']}\"\"\"
+    {system_line}{num_ctx_line}{stop_params}
+    """
+
+    def create_and_push_ollama_model(self):
+        modelfile_content = self.prepare_modelfile_content()
+        with open("Modelfile", "w") as file:
+            file.write(modelfile_content)
+        subprocess.run(["ollama", "serve"])
+        subprocess.run(["ollama", "create", f"{self.config['ollama_model']}:{self.config['model_parameters']}", "-f", "Modelfile"])
+        subprocess.run(["ollama", "push", f"{self.config['ollama_model']}:{self.config['model_parameters']}"])
+
+    def run(self):
+        self.print_system_info()
+        self.check_gpu()
+        self.check_ram()
+        if self.config.get("train", "true").lower() == "true":
+            self.prepare_model()
+            self.train_model()
+        if self.config.get("huggingface_save", "true").lower() == "true":
+            self.save_model_merged()
+        if self.config.get("huggingface_save_gguf", "true").lower() == "true":
+            self.push_model_gguf()
+        if self.config.get("ollama_save", "true").lower() == "true":
+            self.create_and_push_ollama_model()
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="PraisonAI Training Script")
+    parser.add_argument("command", choices=["train"], help="Command to execute")
+    parser.add_argument("--config", default="config.yaml", help="Path to configuration file")
+    parser.add_argument("--model", type=str, help="Model name")
+    parser.add_argument("--hf", type=str, help="Hugging Face model name")
+    parser.add_argument("--ollama", type=str, help="Ollama model name")
+    parser.add_argument("--dataset", type=str, help="Dataset name for training")
+    args = parser.parse_args()
+
+    if args.command == "train":
+        trainer_obj = TrainModel(config_path=args.config)
+        trainer_obj.run()
+
+if __name__ == "__main__":
+    main()

src/praisonai/praisonai/version.py

@@ -1 +1 @@
-__version__ = "3.9.31"
+__version__ = "3.10.4"

src/praisonai/pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "rich>=13.7",
     "markdown>=3.5",
     "pyparsing>=3.0.0",
-    "praisonaiagents>=0.12.16",
+    "praisonaiagents>=0.13.4",
     "python-dotenv>=0.19.0",
     "instructor>=1.3.3",
     "PyYAML>=6.0",

src/praisonai/tests/unit/train/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for praisonai.train module."""

src/praisonai/tests/unit/train/test_grader.py

@@ -0,0 +1,221 @@
+"""
+TDD Tests for Agent Training Grader.
+
+Tests for TrainingGrader class that grades agent outputs using LLM-as-judge.
+DRY: Reuses AccuracyEvaluator pattern from praisonaiagents.eval.
+"""
+
+from unittest.mock import Mock, patch
+import sys
+
+
+class TestTrainingGrader:
+    """Tests for TrainingGrader class."""
+    
+    def test_create_grader_default(self):
+        """Test creating grader with defaults."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader()
+        
+        assert grader.model == "gpt-4o-mini"
+    
+    def test_create_grader_custom_model(self):
+        """Test creating grader with custom model."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader(model="gpt-4o")
+        
+        assert grader.model == "gpt-4o"
+    
+    def test_grade_output_returns_score(self):
+        """Test grading returns a score between 1-10."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        # Mock litellm module
+        mock_litellm = Mock()
+        mock_response = Mock()
+        mock_response.choices = [Mock()]
+        mock_response.choices[0].message.content = "SCORE: 8\nREASONING: Good response"
+        mock_litellm.completion.return_value = mock_response
+        
+        with patch.dict(sys.modules, {'litellm': mock_litellm}):
+            grader = TrainingGrader()
+            result = grader.grade(
+                input_text="What is Python?",
+                output="Python is a programming language",
+                expected_output="Python is a high-level programming language"
+            )
+        
+        assert result.score == 8.0
+        assert "Good response" in result.reasoning
+    
+    def test_grade_output_clamps_score(self):
+        """Test that scores are clamped to 1-10 range."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        # Mock LLM returning out-of-range score
+        mock_litellm = Mock()
+        mock_response = Mock()
+        mock_response.choices = [Mock()]
+        mock_response.choices[0].message.content = "SCORE: 15\nREASONING: Amazing"
+        mock_litellm.completion.return_value = mock_response
+        
+        with patch.dict(sys.modules, {'litellm': mock_litellm}):
+            grader = TrainingGrader()
+            result = grader.grade(
+                input_text="Test",
+                output="Output"
+            )
+        
+        assert result.score == 10.0  # Clamped to max
+    
+    def test_grade_without_expected_output(self):
+        """Test grading without expected output (quality assessment only)."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        mock_litellm = Mock()
+        mock_response = Mock()
+        mock_response.choices = [Mock()]
+        mock_response.choices[0].message.content = "SCORE: 7\nREASONING: Decent quality"
+        mock_litellm.completion.return_value = mock_response
+        
+        with patch.dict(sys.modules, {'litellm': mock_litellm}):
+            grader = TrainingGrader()
+            result = grader.grade(
+                input_text="Explain AI",
+                output="AI is artificial intelligence that mimics human thinking"
+            )
+        
+        assert result.score == 7.0
+    
+    def test_grade_returns_suggestions(self):
+        """Test that grading returns improvement suggestions."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        mock_litellm = Mock()
+        mock_response = Mock()
+        mock_response.choices = [Mock()]
+        mock_response.choices[0].message.content = """SCORE: 6
+REASONING: Could be better
+SUGGESTIONS:
+- Add more detail
+- Include examples
+- Be more concise"""
+        mock_litellm.completion.return_value = mock_response
+        
+        with patch.dict(sys.modules, {'litellm': mock_litellm}):
+            grader = TrainingGrader()
+            result = grader.grade(
+                input_text="Test",
+                output="Output"
+            )
+        
+        assert len(result.suggestions) >= 1
+    
+    def test_grade_handles_malformed_response(self):
+        """Test graceful handling of malformed LLM response."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        mock_litellm = Mock()
+        mock_response = Mock()
+        mock_response.choices = [Mock()]
+        mock_response.choices[0].message.content = "This is not formatted correctly"
+        mock_litellm.completion.return_value = mock_response
+        
+        with patch.dict(sys.modules, {'litellm': mock_litellm}):
+            grader = TrainingGrader()
+            result = grader.grade(
+                input_text="Test",
+                output="Output"
+            )
+        
+        # Should return default score, not crash
+        assert 1.0 <= result.score <= 10.0
+    
+    def test_grade_result_dataclass(self):
+        """Test GradeResult dataclass structure."""
+        from praisonai.train.agents.grader import GradeResult
+        
+        result = GradeResult(
+            score=8.5,
+            reasoning="Good work",
+            suggestions=["Add examples"],
+            input_text="Input",
+            output="Output"
+        )
+        
+        assert result.score == 8.5
+        assert result.reasoning == "Good work"
+        assert result.suggestions == ["Add examples"]
+    
+    def test_grade_result_to_dict(self):
+        """Test converting GradeResult to dictionary."""
+        from praisonai.train.agents.grader import GradeResult
+        
+        result = GradeResult(
+            score=7.0,
+            reasoning="OK",
+            suggestions=[],
+            input_text="In",
+            output="Out"
+        )
+        
+        d = result.to_dict()
+        assert d["score"] == 7.0
+        assert "reasoning" in d
+        assert "suggestions" in d
+
+
+class TestTrainingGraderPrompts:
+    """Tests for grader prompt generation."""
+    
+    def test_prompt_includes_input(self):
+        """Test that grading prompt includes input text."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader()
+        prompt = grader._build_prompt(
+            input_text="What is machine learning?",
+            output="ML is a subset of AI"
+        )
+        
+        assert "What is machine learning?" in prompt
+    
+    def test_prompt_includes_output(self):
+        """Test that grading prompt includes agent output."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader()
+        prompt = grader._build_prompt(
+            input_text="Test",
+            output="This is the agent's response"
+        )
+        
+        assert "This is the agent's response" in prompt
+    
+    def test_prompt_includes_expected_when_provided(self):
+        """Test that prompt includes expected output when provided."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader()
+        prompt = grader._build_prompt(
+            input_text="Test",
+            output="Output",
+            expected_output="Expected response here"
+        )
+        
+        assert "Expected response here" in prompt
+    
+    def test_prompt_requests_score_format(self):
+        """Test that prompt requests specific score format."""
+        from praisonai.train.agents.grader import TrainingGrader
+        
+        grader = TrainingGrader()
+        prompt = grader._build_prompt(
+            input_text="Test",
+            output="Output"
+        )
+        
+        assert "SCORE:" in prompt
+        assert "REASONING:" in prompt

src/praisonai/tests/unit/train/test_models.py

@@ -0,0 +1,219 @@
+"""
+TDD Tests for Agent Training Models.
+
+Tests for TrainingScenario, TrainingIteration, and TrainingReport dataclasses.
+"""
+
+from datetime import datetime
+
+
+class TestTrainingScenario:
+    """Tests for TrainingScenario dataclass."""
+    
+    def test_create_scenario_minimal(self):
+        """Test creating scenario with minimal required fields."""
+        from praisonai.train.agents.models import TrainingScenario
+        
+        scenario = TrainingScenario(
+            id="test-1",
+            input_text="What is 2+2?"
+        )
+        
+        assert scenario.id == "test-1"
+        assert scenario.input_text == "What is 2+2?"
+        assert scenario.expected_output is None
+        assert scenario.context == {}
+    
+    def test_create_scenario_full(self):
+        """Test creating scenario with all fields."""
+        from praisonai.train.agents.models import TrainingScenario
+        
+        scenario = TrainingScenario(
+            id="test-2",
+            input_text="Explain quantum computing",
+            expected_output="Quantum computing uses qubits...",
+            context={"domain": "science", "difficulty": "advanced"}
+        )
+        
+        assert scenario.id == "test-2"
+        assert scenario.expected_output == "Quantum computing uses qubits..."
+        assert scenario.context["domain"] == "science"
+    
+    def test_scenario_to_dict(self):
+        """Test converting scenario to dictionary."""
+        from praisonai.train.agents.models import TrainingScenario
+        
+        scenario = TrainingScenario(
+            id="test-3",
+            input_text="Hello"
+        )
+        
+        d = scenario.to_dict()
+        assert d["id"] == "test-3"
+        assert d["input_text"] == "Hello"
+        assert "expected_output" in d
+    
+    def test_scenario_from_dict(self):
+        """Test creating scenario from dictionary."""
+        from praisonai.train.agents.models import TrainingScenario
+        
+        data = {
+            "id": "test-4",
+            "input_text": "Test input",
+            "expected_output": "Test output"
+        }
+        
+        scenario = TrainingScenario.from_dict(data)
+        assert scenario.id == "test-4"
+        assert scenario.input_text == "Test input"
+
+
+class TestTrainingIteration:
+    """Tests for TrainingIteration dataclass."""
+    
+    def test_create_iteration(self):
+        """Test creating a training iteration."""
+        from praisonai.train.agents.models import TrainingIteration
+        
+        iteration = TrainingIteration(
+            iteration_num=1,
+            scenario_id="test-1",
+            input_text="What is AI?",
+            output="AI is artificial intelligence...",
+            score=8.5,
+            feedback="Good explanation but could be more concise",
+            suggestions=["Be more concise", "Add examples"]
+        )
+        
+        assert iteration.iteration_num == 1
+        assert iteration.scenario_id == "test-1"
+        assert iteration.score == 8.5
+        assert len(iteration.suggestions) == 2
+    
+    def test_iteration_has_timestamp(self):
+        """Test that iteration has timestamp."""
+        from praisonai.train.agents.models import TrainingIteration
+        
+        iteration = TrainingIteration(
+            iteration_num=1,
+            scenario_id="test-1",
+            input_text="Test",
+            output="Output",
+            score=7.0,
+            feedback="OK"
+        )
+        
+        assert iteration.timestamp is not None
+        # Should be a valid ISO format string
+        datetime.fromisoformat(iteration.timestamp)
+    
+    def test_iteration_to_dict(self):
+        """Test converting iteration to dictionary."""
+        from praisonai.train.agents.models import TrainingIteration
+        
+        iteration = TrainingIteration(
+            iteration_num=2,
+            scenario_id="test-2",
+            input_text="Input",
+            output="Output",
+            score=9.0,
+            feedback="Excellent"
+        )
+        
+        d = iteration.to_dict()
+        assert d["iteration_num"] == 2
+        assert d["score"] == 9.0
+        assert "timestamp" in d
+
+
+class TestTrainingReport:
+    """Tests for TrainingReport dataclass."""
+    
+    def test_create_report(self):
+        """Test creating a training report."""
+        from praisonai.train.agents.models import TrainingReport, TrainingIteration
+        
+        iterations = [
+            TrainingIteration(
+                iteration_num=1,
+                scenario_id="s1",
+                input_text="Test",
+                output="Output 1",
+                score=6.0,
+                feedback="Needs improvement"
+            ),
+            TrainingIteration(
+                iteration_num=2,
+                scenario_id="s1",
+                input_text="Test",
+                output="Output 2",
+                score=8.0,
+                feedback="Better"
+            ),
+        ]
+        
+        report = TrainingReport(
+            session_id="train-abc123",
+            iterations=iterations,
+            total_iterations=2
+        )
+        
+        assert report.session_id == "train-abc123"
+        assert len(report.iterations) == 2
+        assert report.total_iterations == 2
+    
+    def test_report_avg_score(self):
+        """Test calculating average score."""
+        from praisonai.train.agents.models import TrainingReport, TrainingIteration
+        
+        iterations = [
+            TrainingIteration(iteration_num=1, scenario_id="s1", input_text="T", output="O", score=6.0, feedback="F"),
+            TrainingIteration(iteration_num=2, scenario_id="s1", input_text="T", output="O", score=8.0, feedback="F"),
+            TrainingIteration(iteration_num=3, scenario_id="s1", input_text="T", output="O", score=10.0, feedback="F"),
+        ]
+        
+        report = TrainingReport(session_id="test", iterations=iterations, total_iterations=3)
+        
+        assert report.avg_score == 8.0  # (6+8+10)/3
+    
+    def test_report_improvement(self):
+        """Test calculating improvement from first to last."""
+        from praisonai.train.agents.models import TrainingReport, TrainingIteration
+        
+        iterations = [
+            TrainingIteration(iteration_num=1, scenario_id="s1", input_text="T", output="O", score=5.0, feedback="F"),
+            TrainingIteration(iteration_num=2, scenario_id="s1", input_text="T", output="O", score=7.0, feedback="F"),
+            TrainingIteration(iteration_num=3, scenario_id="s1", input_text="T", output="O", score=9.0, feedback="F"),
+        ]
+        
+        report = TrainingReport(session_id="test", iterations=iterations, total_iterations=3)
+        
+        assert report.improvement == 4.0  # 9.0 - 5.0
+    
+    def test_report_to_dict(self):
+        """Test converting report to dictionary."""
+        from praisonai.train.agents.models import TrainingReport
+        
+        report = TrainingReport(
+            session_id="test-session",
+            iterations=[],
+            total_iterations=0
+        )
+        
+        d = report.to_dict()
+        assert d["session_id"] == "test-session"
+        assert "avg_score" in d
+        assert "improvement" in d
+    
+    def test_report_empty_iterations(self):
+        """Test report with no iterations."""
+        from praisonai.train.agents.models import TrainingReport
+        
+        report = TrainingReport(
+            session_id="empty",
+            iterations=[],
+            total_iterations=0
+        )
+        
+        assert report.avg_score == 0.0
+        assert report.improvement == 0.0

src/praisonai/tests/unit/train/test_orchestrator.py

@@ -0,0 +1,326 @@
+"""
+TDD Tests for Agent Training Orchestrator.
+
+Tests for AgentTrainer class that orchestrates the training loop.
+"""
+
+from unittest.mock import Mock, patch
+import tempfile
+from pathlib import Path
+
+
+class TestAgentTrainer:
+    """Tests for AgentTrainer class."""
+    
+    def test_create_trainer_with_agent(self):
+        """Test creating trainer with an agent."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        
+        mock_agent = Mock()
+        mock_agent.chat = Mock(return_value="Agent response")
+        
+        trainer = AgentTrainer(agent=mock_agent)
+        
+        assert trainer.agent == mock_agent
+        assert trainer.iterations == 3  # Default
+        assert trainer.human_mode is False  # Default is LLM mode
+    
+    def test_create_trainer_with_custom_iterations(self):
+        """Test creating trainer with custom iteration count."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        
+        mock_agent = Mock()
+        trainer = AgentTrainer(agent=mock_agent, iterations=5)
+        
+        assert trainer.iterations == 5
+    
+    def test_create_trainer_human_mode(self):
+        """Test creating trainer in human-in-the-loop mode."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        
+        mock_agent = Mock()
+        trainer = AgentTrainer(agent=mock_agent, human_mode=True)
+        
+        assert trainer.human_mode is True
+    
+    def test_trainer_generates_session_id(self):
+        """Test that trainer generates unique session ID."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        
+        mock_agent = Mock()
+        trainer1 = AgentTrainer(agent=mock_agent)
+        trainer2 = AgentTrainer(agent=mock_agent)
+        
+        assert trainer1.session_id.startswith("train-")
+        assert trainer1.session_id != trainer2.session_id
+    
+    def test_add_scenario(self):
+        """Test adding training scenarios."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        
+        mock_agent = Mock()
+        trainer = AgentTrainer(agent=mock_agent)
+        
+        scenario = TrainingScenario(
+            id="s1",
+            input_text="What is AI?"
+        )
+        trainer.add_scenario(scenario)
+        
+        assert len(trainer.scenarios) == 1
+        assert trainer.scenarios[0].id == "s1"
+    
+    def test_add_scenario_from_dict(self):
+        """Test adding scenario from dictionary."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        
+        mock_agent = Mock()
+        trainer = AgentTrainer(agent=mock_agent)
+        
+        trainer.add_scenario({
+            "id": "s2",
+            "input_text": "Explain Python",
+            "expected_output": "Python is a language"
+        })
+        
+        assert len(trainer.scenarios) == 1
+        assert trainer.scenarios[0].input_text == "Explain Python"
+    
+    @patch("praisonai.train.agents.orchestrator.TrainingGrader")
+    def test_run_llm_mode_single_iteration(self, mock_grader_class):
+        """Test running single iteration in LLM mode."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        from praisonai.train.agents.grader import GradeResult
+        
+        # Setup mocks
+        mock_agent = Mock()
+        mock_agent.chat = Mock(return_value="Agent response")
+        
+        mock_grader = Mock()
+        mock_grader.grade = Mock(return_value=GradeResult(
+            score=8.0,
+            reasoning="Good",
+            suggestions=["Be more specific"],
+            input_text="Test",
+            output="Agent response"
+        ))
+        mock_grader_class.return_value = mock_grader
+        
+        # Create trainer
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=mock_agent,
+                iterations=1,
+                storage_dir=Path(tmpdir)
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Test"))
+            
+            report = trainer.run()
+        
+        assert report is not None
+        assert len(report.iterations) == 1
+        assert report.iterations[0].score == 8.0
+    
+    @patch("praisonai.train.agents.orchestrator.TrainingGrader")
+    def test_run_llm_mode_multiple_iterations(self, mock_grader_class):
+        """Test running multiple iterations shows improvement."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        from praisonai.train.agents.grader import GradeResult
+        
+        mock_agent = Mock()
+        call_count = [0]
+        
+        def agent_response(prompt):
+            call_count[0] += 1
+            return f"Response {call_count[0]}"
+        
+        mock_agent.chat = Mock(side_effect=agent_response)
+        
+        # Grader returns improving scores
+        scores = [6.0, 7.5, 9.0]
+        score_idx = [0]
+        
+        def grade_response(*args, **kwargs):
+            score = scores[min(score_idx[0], len(scores)-1)]
+            score_idx[0] += 1
+            return GradeResult(
+                score=score,
+                reasoning=f"Score {score}",
+                suggestions=["Improve"],
+                input_text="Test",
+                output="Output"
+            )
+        
+        mock_grader = Mock()
+        mock_grader.grade = Mock(side_effect=grade_response)
+        mock_grader_class.return_value = mock_grader
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=mock_agent,
+                iterations=3,
+                storage_dir=Path(tmpdir)
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Test"))
+            
+            report = trainer.run()
+        
+        assert len(report.iterations) == 3
+        assert report.improvement > 0  # Should show improvement
+    
+    @patch("builtins.input")
+    def test_run_human_mode_prompts_for_feedback(self, mock_input):
+        """Test that human mode prompts for feedback."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        
+        mock_agent = Mock()
+        mock_agent.chat = Mock(return_value="Agent response")
+        
+        # Mock input sequence: score, feedback, suggestions
+        mock_input.side_effect = ["8", "Good feedback", "suggestion1, suggestion2"]
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=mock_agent,
+                iterations=1,
+                human_mode=True,
+                storage_dir=Path(tmpdir),
+                verbose=False,  # Suppress output in tests
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Test"))
+            
+            report = trainer.run()
+        
+        # Should have called input() for feedback
+        mock_input.assert_called()
+        assert len(report.iterations) == 1
+    
+    def test_run_without_scenarios_raises(self):
+        """Test that running without scenarios raises error."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        import pytest
+        
+        mock_agent = Mock()
+        trainer = AgentTrainer(agent=mock_agent)
+        
+        with pytest.raises(ValueError, match="No scenarios"):
+            trainer.run()
+    
+    @patch("praisonai.train.agents.orchestrator.TrainingGrader")
+    def test_run_saves_to_storage(self, mock_grader_class):
+        """Test that run saves iterations to storage."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        from praisonai.train.agents.grader import GradeResult
+        import json
+        
+        mock_agent = Mock()
+        mock_agent.chat = Mock(return_value="Response")
+        
+        mock_grader = Mock()
+        mock_grader.grade = Mock(return_value=GradeResult(
+            score=7.0,
+            reasoning="OK",
+            suggestions=[],
+            input_text="Test",
+            output="Response"
+        ))
+        mock_grader_class.return_value = mock_grader
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=mock_agent,
+                iterations=1,
+                storage_dir=Path(tmpdir)
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Test"))
+            
+            trainer.run()
+            
+            # Check storage file exists
+            storage_file = Path(tmpdir) / f"{trainer.session_id}.json"
+            assert storage_file.exists()
+            
+            with open(storage_file) as f:
+                data = json.load(f)
+            
+            assert "iterations" in data
+
+
+class TestAgentTrainerWithAgents:
+    """Tests for AgentTrainer with multi-agent (Agents class)."""
+    
+    @patch("praisonai.train.agents.orchestrator.TrainingGrader")
+    def test_trainer_works_with_agents_class(self, mock_grader_class):
+        """Test trainer works with Agents (multi-agent) class."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        from praisonai.train.agents.grader import GradeResult
+        
+        # Mock Agents class (has start() method but NOT chat())
+        mock_agents = Mock(spec=['start'])  # Only has start, not chat
+        mock_agents.start = Mock(return_value="Multi-agent response")
+        
+        mock_grader = Mock()
+        mock_grader.grade = Mock(return_value=GradeResult(
+            score=8.0,
+            reasoning="Good",
+            suggestions=[],
+            input_text="Test",
+            output="Response"
+        ))
+        mock_grader_class.return_value = mock_grader
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=mock_agents,
+                iterations=1,
+                storage_dir=Path(tmpdir),
+                verbose=False,
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Test"))
+            
+            report = trainer.run()
+        
+        assert report is not None
+        mock_agents.start.assert_called()
+
+
+class TestAgentTrainerCallable:
+    """Tests for AgentTrainer with callable functions."""
+    
+    @patch("praisonai.train.agents.orchestrator.TrainingGrader")
+    def test_trainer_works_with_callable(self, mock_grader_class):
+        """Test trainer works with a callable function."""
+        from praisonai.train.agents.orchestrator import AgentTrainer
+        from praisonai.train.agents.models import TrainingScenario
+        from praisonai.train.agents.grader import GradeResult
+        
+        def my_func(prompt):
+            return f"Response to: {prompt}"
+        
+        mock_grader = Mock()
+        mock_grader.grade = Mock(return_value=GradeResult(
+            score=7.5,
+            reasoning="OK",
+            suggestions=[],
+            input_text="Test",
+            output="Response"
+        ))
+        mock_grader_class.return_value = mock_grader
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = AgentTrainer(
+                agent=my_func,
+                iterations=1,
+                storage_dir=Path(tmpdir)
+            )
+            trainer.add_scenario(TrainingScenario(id="s1", input_text="Hello"))
+            
+            report = trainer.run()
+        
+        assert report is not None

src/praisonai/tests/unit/train/test_storage.py

@@ -0,0 +1,249 @@
+"""
+TDD Tests for Agent Training Storage.
+
+Tests for TrainingStorage class that persists training data to JSON.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+
+class TestTrainingStorage:
+    """Tests for TrainingStorage class."""
+    
+    def test_create_storage_default_path(self):
+        """Test creating storage with default path."""
+        from praisonai.train.agents.storage import TrainingStorage
+        
+        storage = TrainingStorage(session_id="test-session")
+        
+        assert storage.session_id == "test-session"
+        assert "train" in str(storage.storage_path)
+    
+    def test_create_storage_custom_path(self):
+        """Test creating storage with custom path."""
+        from praisonai.train.agents.storage import TrainingStorage
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="custom-session",
+                storage_dir=Path(tmpdir)
+            )
+            
+            assert storage.storage_dir == Path(tmpdir)
+    
+    def test_save_iteration(self):
+        """Test saving a training iteration."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingIteration
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="save-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            iteration = TrainingIteration(
+                iteration_num=1,
+                scenario_id="s1",
+                input_text="Test input",
+                output="Test output",
+                score=7.5,
+                feedback="Good job"
+            )
+            
+            storage.save_iteration(iteration)
+            
+            # Verify file exists
+            assert storage.storage_path.exists()
+            
+            # Verify content
+            with open(storage.storage_path) as f:
+                data = json.load(f)
+            
+            assert "iterations" in data
+            assert len(data["iterations"]) == 1
+            assert data["iterations"][0]["score"] == 7.5
+    
+    def test_save_multiple_iterations(self):
+        """Test saving multiple iterations."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingIteration
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="multi-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            for i in range(3):
+                iteration = TrainingIteration(
+                    iteration_num=i + 1,
+                    scenario_id="s1",
+                    input_text=f"Input {i}",
+                    output=f"Output {i}",
+                    score=5.0 + i,
+                    feedback=f"Feedback {i}"
+                )
+                storage.save_iteration(iteration)
+            
+            # Verify all saved
+            with open(storage.storage_path) as f:
+                data = json.load(f)
+            
+            assert len(data["iterations"]) == 3
+    
+    def test_load_iterations(self):
+        """Test loading iterations from storage."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingIteration
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="load-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            # Save some iterations
+            for i in range(2):
+                iteration = TrainingIteration(
+                    iteration_num=i + 1,
+                    scenario_id="s1",
+                    input_text=f"Input {i}",
+                    output=f"Output {i}",
+                    score=6.0 + i,
+                    feedback=f"Feedback {i}"
+                )
+                storage.save_iteration(iteration)
+            
+            # Load them back
+            iterations = storage.load_iterations()
+            
+            assert len(iterations) == 2
+            assert iterations[0].iteration_num == 1
+            assert iterations[1].score == 7.0
+    
+    def test_load_empty_storage(self):
+        """Test loading from empty/nonexistent storage."""
+        from praisonai.train.agents.storage import TrainingStorage
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="empty-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            iterations = storage.load_iterations()
+            
+            assert iterations == []
+    
+    def test_save_scenario(self):
+        """Test saving a training scenario."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingScenario
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="scenario-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            scenario = TrainingScenario(
+                id="s1",
+                input_text="What is Python?",
+                expected_output="Python is a programming language"
+            )
+            
+            storage.save_scenario(scenario)
+            
+            # Verify
+            with open(storage.storage_path) as f:
+                data = json.load(f)
+            
+            assert "scenarios" in data
+            assert len(data["scenarios"]) == 1
+    
+    def test_save_report(self):
+        """Test saving a training report."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingReport
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="report-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            report = TrainingReport(
+                session_id="report-test",
+                iterations=[],
+                total_iterations=0
+            )
+            
+            storage.save_report(report)
+            
+            # Verify
+            with open(storage.storage_path) as f:
+                data = json.load(f)
+            
+            assert "report" in data
+            assert data["report"]["session_id"] == "report-test"
+    
+    def test_storage_uses_json_not_pickle(self):
+        """Verify storage uses JSON format, not pickle."""
+        from praisonai.train.agents.storage import TrainingStorage
+        from praisonai.train.agents.models import TrainingIteration
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = TrainingStorage(
+                session_id="json-test",
+                storage_dir=Path(tmpdir)
+            )
+            
+            iteration = TrainingIteration(
+                iteration_num=1,
+                scenario_id="s1",
+                input_text="Test",
+                output="Output",
+                score=8.0,
+                feedback="Good"
+            )
+            storage.save_iteration(iteration)
+            
+            # File should be JSON (readable as text)
+            content = storage.storage_path.read_text()
+            assert content.startswith("{")  # JSON starts with {
+            
+            # Should be valid JSON
+            json.loads(content)
+    
+    def test_list_sessions(self):
+        """Test listing all training sessions."""
+        from praisonai.train.agents.storage import TrainingStorage, list_training_sessions
+        from praisonai.train.agents.models import TrainingIteration
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+            
+            # Create multiple sessions
+            for i in range(3):
+                storage = TrainingStorage(
+                    session_id=f"session-{i}",
+                    storage_dir=tmppath
+                )
+                iteration = TrainingIteration(
+                    iteration_num=1,
+                    scenario_id="s1",
+                    input_text="Test",
+                    output="Output",
+                    score=7.0,
+                    feedback="OK"
+                )
+                storage.save_iteration(iteration)
+            
+            # List sessions
+            sessions = list_training_sessions(storage_dir=tmppath)
+            
+            assert len(sessions) == 3
+            assert any(s.session_id == "session-0" for s in sessions)

src/praisonai/uv.lock

@@ -4276,7 +4276,7 @@ requires-dist = [
     { name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.1.0" },
     { name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.1.0" },
     { name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.1.0" },
-    { name = "praisonaiagents", specifier = ">=0.12.16" },
+    { name = "praisonaiagents", specifier = ">=0.13.4" },
     { name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
     { name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
     { name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
@@ -4324,7 +4324,7 @@ wheels = [
 
 [[package]]
 name = "praisonaiagents"
-version = "0.12.16"
+version = "0.13.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "aiohttp" },
@@ -4333,9 +4333,9 @@ dependencies = [
     { name = "pydantic" },
     { name = "rich" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c9/97/37b1a5231cdea79503cff3cdffb74620ed8282a676dd60746e09fd9fa599/praisonaiagents-0.12.16.tar.gz", hash = "sha256:ce56bf0fb97bf2d35eb9be6fa19715ace57edbbfde81cf9d41d74a17860514d5", size = 1037591, upload-time = "2026-01-20T12:16:38.047Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/fe/2d/c9e721880cf13f3c0cfb17b36d9a1ed226ed94773dbd32c41f45c2cea773/praisonaiagents-0.13.4.tar.gz", hash = "sha256:e31e37fdb5aa71ac58ac9dc856f4b11f21f61c0b3f8d84cf36f6d55bea6008db", size = 1061250, upload-time = "2026-01-22T07:03:35.194Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/53/01/b29cf7fce20c3c9897488e63fe86470842f1252329dd05a11a9ac581da05/praisonaiagents-0.12.16-py3-none-any.whl", hash = "sha256:ef2f51e2d1ac2b63f99751e99fed612a745d82cf4ecc32bf641bb18de74dfe42", size = 1007888, upload-time = "2026-01-20T12:16:36.195Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/b6/7cb383aafca3d1532f666b529aafb2c25d694d30bda96781d313deaa80bd/praisonaiagents-0.13.4-py3-none-any.whl", hash = "sha256:51825df03125d24d095a3542f27c497085a7baaff3c4fdcce5258538052ca25d", size = 1033775, upload-time = "2026-01-22T07:03:32.829Z" },
 ]
 
 [[package]]