aip-agents-binary 0.0.0b4__cp312-cp312-win_amd64.whl → 0.5.4__cp312-cp312-win_amd64.whl

aip_agents/agent/langgraph_memory_enhancer_agent.pyi

@@ -0,0 +1,49 @@
+from _typeshed import Incomplete
+from aip_agents.agent.langgraph_react_agent import LangGraphReactAgent as LangGraphReactAgent
+from aip_agents.agent.system_instruction_context import get_current_date_context as get_current_date_context
+from aip_agents.memory.guidance import MEM0_MEMORY_RECALL_GUIDANCE as MEM0_MEMORY_RECALL_GUIDANCE
+from aip_agents.tools.memory_search_tool import LongTermMemorySearchTool as LongTermMemorySearchTool, MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0SearchTool as Mem0SearchTool
+from aip_agents.utils.logger import get_logger as get_logger
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+logger: Incomplete
+
+class LangGraphMemoryEnhancerAgent(LangGraphReactAgent):
+    """Simplified mini-agent for automatic memory retrieval and query enhancement.
+
+    This agent has a simple 2-node LangGraph (agent + tools) and uses existing memory
+    infrastructure to enhance user queries with relevant context. It acts as a
+    preprocessing layer that automatically attempts memory retrieval for every query.
+
+    Key features:
+    - Uses runtime `memory_user_id` provided via call arguments (no static storage)
+    - Uses simplified instruction reusing existing guidance
+    - Standard 2-node LangGraph pattern (agent -> tools -> agent)
+    - Automatically enhances queries with memory context when available
+    - Returns original query unchanged if no relevant memories found
+    """
+    def __init__(self, memory, **kwargs) -> None:
+        """Initialize the LangGraphMemoryEnhancerAgent with memory backend and configuration.
+
+        Args:
+            memory: Memory backend instance (Mem0Memory or compatible)
+            **kwargs: Additional arguments passed to BaseLangGraphAgent, including:
+                - memory_agent_id: Fallback user ID for memory operations
+                - model: LLM model to use for memory decisions
+                - Other BaseLangGraphAgent parameters
+        """
+    def define_graph(self, graph_builder: StateGraph) -> CompiledStateGraph:
+        """Define the 3-node memory recall LangGraph for this agent.
+
+        This creates a streamlined ReAct-inspired structure that reuses
+        `LangGraphReactAgent` helpers for robust LM invocation, token usage tracking,
+        error handling, and tool execution.
+
+        Args:
+            graph_builder: LangGraph `StateGraph` builder instance used to register nodes and
+                edges for compilation.
+
+        Returns:
+            CompiledStateGraph: The compiled memory recall graph ready for execution.
+        """

aip_agents/agent/langgraph_react_agent.pyi

@@ -1,81 +1,60 @@
 from _typeshed import Incomplete
 from aip_agents.agent.base_langgraph_agent import BaseLangGraphAgent as BaseLangGraphAgent
+from aip_agents.agent.hitl.langgraph_hitl_mixin import LangGraphHitLMixin as LangGraphHitLMixin
+from aip_agents.agent.hitl.manager import TOOL_EXECUTION_BLOCKING_DECISIONS as TOOL_EXECUTION_BLOCKING_DECISIONS
+from aip_agents.middleware.base import AgentMiddleware as AgentMiddleware, ModelRequest as ModelRequest
+from aip_agents.middleware.manager import MiddlewareManager as MiddlewareManager
+from aip_agents.middleware.todolist import TodoList as TodoList, TodoListMiddleware as TodoListMiddleware
+from aip_agents.schema.a2a import A2AStreamEventType as A2AStreamEventType
+from aip_agents.schema.hitl import ApprovalDecision as ApprovalDecision, HitlMetadata as HitlMetadata
+from aip_agents.schema.langgraph import ToolCallResult as ToolCallResult, ToolStorageParams as ToolStorageParams
+from aip_agents.schema.step_limit import MaxStepsExceededError as MaxStepsExceededError, StepLimitConfig as StepLimitConfig
+from aip_agents.tools.memory_search_tool import MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME
 from aip_agents.tools.tool_config_injector import TOOL_CONFIGS_KEY as TOOL_CONFIGS_KEY
-from aip_agents.types.a2a_events import A2AStreamEventType as A2AStreamEventType
 from aip_agents.utils import add_references_chunks as add_references_chunks
 from aip_agents.utils.langgraph import convert_langchain_messages_to_gllm_messages as convert_langchain_messages_to_gllm_messages, convert_lm_output_to_langchain_message as convert_lm_output_to_langchain_message
 from aip_agents.utils.langgraph.tool_output_management import StoreOutputParams as StoreOutputParams, ToolOutputManager as ToolOutputManager, ToolReferenceError as ToolReferenceError, ToolReferenceResolver as ToolReferenceResolver
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.utils.logger import get_logger as get_logger
+from aip_agents.utils.metadata.activity_metadata_helper import create_tool_activity_info as create_tool_activity_info
+from aip_agents.utils.metadata_helper import Kind as Kind, MetadataFieldKeys as MetadataFieldKeys, Status as Status
+from aip_agents.utils.pii import ToolPIIHandler as ToolPIIHandler, add_pii_mappings as add_pii_mappings
+from aip_agents.utils.reference_helper import extract_references_from_tool as extract_references_from_tool
+from aip_agents.utils.step_limit_manager import StepLimitManager as StepLimitManager
 from aip_agents.utils.token_usage_helper import TOTAL_USAGE_KEY as TOTAL_USAGE_KEY, USAGE_METADATA_KEY as USAGE_METADATA_KEY, add_usage_metadata as add_usage_metadata, extract_and_update_token_usage_from_ai_message as extract_and_update_token_usage_from_ai_message, extract_token_usage_from_tool_output as extract_token_usage_from_tool_output
-from collections.abc import Callable as Callable
+from collections.abc import Awaitable as Awaitable, Sequence
 from dataclasses import dataclass
-from gllm_core.event import EventEmitter as EventEmitter
-from gllm_core.schema import Chunk
+from gllm_core.event import EventEmitter
+from gllm_core.schema import Chunk as Chunk
 from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import BaseMessage as BaseMessage, ToolMessage
-from langchain_core.messages.ai import UsageMetadata as UsageMetadata
-from langchain_core.tools import BaseTool as BaseTool
-from langgraph.graph import StateGraph as StateGraph
+from langchain_core.messages import BaseMessage as BaseMessage
+from langchain_core.messages.ai import UsageMetadata
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
 from langgraph.graph.message import add_messages as add_messages
-from langgraph.graph.state import CompiledStateGraph as CompiledStateGraph
-from langgraph.managed import IsLastStep as IsLastStep, RemainingSteps as RemainingSteps
-from langgraph.types import StreamWriter as StreamWriter
-from typing import Any, Sequence
-from typing_extensions import Annotated, TypedDict
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.managed import IsLastStep, RemainingSteps
+from typing import Annotated, Any
+from typing_extensions import TypedDict
 
 logger: Incomplete
 DEFAULT_INSTRUCTION: str
 TOOL_RUN_STREAMING_METHOD: str
-SAVE_OUTPUT_HISTORY: str
-FORMAT_AGENT_REFERENCE: str
 TOOL_OUTPUT_MANAGER_KEY: str
 CALL_ID_KEY: str
 
 @dataclass
-class ToolCallResult:
-    """Container for the results of a single tool call execution.
-
-    Holds all the data returned from a tool execution in a structured format.
-
-    Attributes:
-        messages: List of tool messages generated from the tool call.
-        artifacts: List of artifacts produced by the tool call.
-        metadata_delta: Dictionary of metadata updates from the tool call.
-        references: List of reference chunks from the tool call.
-        step_usage: Token usage metadata from the tool call.
-    """
-    messages: list[ToolMessage]
-    artifacts: list[dict[str, Any]]
-    metadata_delta: dict[str, Any]
-    references: list[Chunk]
-    step_usage: UsageMetadata | None
-
-@dataclass
-class ToolStorageParams:
-    """Parameters for automatic tool storage.
-
-    Reduces the number of arguments passed to _handle_automatic_tool_storage method.
-
-    Attributes:
-        tool: The tool instance that was executed.
-        tool_output: The output returned by the tool.
-        tool_call: The tool call information from the AI message.
-        tool_call_id: Unique identifier for this tool call.
-        resolved_args: The resolved arguments passed to the tool.
-        state: Current agent state containing tool output manager.
-    """
-    tool: BaseTool
-    tool_output: Any
-    tool_call: dict[str, Any]
-    tool_call_id: str
-    resolved_args: dict[str, Any]
+class ToolCallContext:
+    """Context information for executing a single tool call."""
+    config: dict[str, Any] | None
     state: dict[str, Any]
+    pending_artifacts: list[dict[str, Any]]
+    hitl_decision: ApprovalDecision | None = ...
 
 class ReactAgentState(TypedDict):
     """State schema for the ReAct agent.
 
     Includes messages, step tracking, optional event emission support, artifacts, references,
-    metadata, and tool output management for efficient tool result sharing and reference resolution.
+    metadata, tool output management, and deep agents middleware state (todos, filesystem).
     """
     messages: Annotated[Sequence[BaseMessage], add_messages]
     is_last_step: IsLastStep
@@ -86,8 +65,15 @@ class ReactAgentState(TypedDict):
     metadata: dict[str, Any] | None
     tool_output_manager: ToolOutputManager | None
     total_usage: Annotated[UsageMetadata | None, add_usage_metadata]
-
-class LangGraphReactAgent(BaseLangGraphAgent):
+    pii_mapping: Annotated[dict[str, str] | None, add_pii_mappings]
+    thread_id: str
+    todos: TodoList | None
+    current_step: int
+    delegation_depth: int
+    delegation_chain: list[str]
+    step_limit_config: StepLimitConfig | None
+
+class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
     """A ReAct agent template built on LangGraph.
 
     This agent can use either:
@@ -98,7 +84,8 @@ class LangGraphReactAgent(BaseLangGraphAgent):
     agent -> tools -> agent (loop) -> END
     """
     tool_output_manager: Incomplete
-    def __init__(self, name: str, instruction: str = ..., model: BaseChatModel | str | Any | None = None, tools: Sequence[BaseTool] | None = None, agents: Sequence[Any] | None = None, description: str | None = None, thread_id_key: str = 'thread_id', event_emitter: EventEmitter | None = None, tool_output_manager: ToolOutputManager | None = None, **kwargs: Any) -> None:
+    step_limit_config: Incomplete
+    def __init__(self, name: str, instruction: str = ..., model: BaseChatModel | str | Any | None = None, tools: Sequence[BaseTool] | None = None, agents: Sequence[Any] | None = None, description: str | None = None, thread_id_key: str = 'thread_id', event_emitter: EventEmitter | None = None, tool_output_manager: ToolOutputManager | None = None, planning: bool = False, middlewares: Sequence[AgentMiddleware] | None = None, step_limit_config: StepLimitConfig | None = None, **kwargs: Any) -> None:
         """Initialize the LangGraph ReAct Agent.
 
         Args:
@@ -114,6 +101,12 @@ class LangGraphReactAgent(BaseLangGraphAgent):
                 When provided, enables tool output storage, reference resolution, and sharing capabilities.
                 This enables multi-agent workflows where agents can access each other's tool outputs.
                 If None, tool output management is disabled for this agent.
+            planning: Enable planning capabilities with TodoListMiddleware. Defaults to False.
+            middlewares: Optional sequence of custom middleware to COMPOSE (not override) with built-in middleware.
+                        Execution order: [TodoListMiddleware (if planning=True),
+                                         ...custom middlewares in order provided]
+                        All middleware hooks execute - this extends capabilities, never replaces them.
+            step_limit_config: Optional configuration for step limits and delegation depth.
             **kwargs: Additional keyword arguments passed to BaseLangGraphAgent.
         """
     def define_graph(self, graph_builder: StateGraph) -> CompiledStateGraph:

aip_agents/agent/system_instruction_context.pyi

@@ -0,0 +1,13 @@
+from aip_agents.utils.constants import DefaultTimezone as DefaultTimezone
+
+__all__ = ['get_current_date_context', 'DefaultTimezone']
+
+def get_current_date_context(timezone: str = ...) -> str:
+    """Generate current date context for system prompts.
+
+    Args:
+        timezone: IANA timezone name for date formatting.
+
+    Returns:
+        Formatted date context string for inclusion in system prompts.
+    """

aip_agents/clients/__init__.pyi

@@ -0,0 +1,4 @@
+__all__ = ['langflow']
+
+# Names in __all__ with no definition:
+#   langflow

aip_agents/clients/langflow/__init__.pyi

@@ -0,0 +1,4 @@
+from aip_agents.clients.langflow.client import LangflowApiClient as LangflowApiClient
+from aip_agents.clients.langflow.types import LangflowEventType as LangflowEventType
+
+__all__ = ['LangflowApiClient', 'LangflowEventType']

aip_agents/clients/langflow/client.pyi

@@ -0,0 +1,140 @@
+from _typeshed import Incomplete
+from aip_agents.clients.langflow.types import LangflowEventType as LangflowEventType
+from aip_agents.schema.agent import HttpxClientOptions as HttpxClientOptions
+from aip_agents.utils.logger import get_logger as get_logger
+from collections.abc import AsyncGenerator
+from typing import Any
+
+logger: Incomplete
+DEFAULT_LANGFLOW_BASE_URL: Incomplete
+MAX_PAGE_SIZE: int
+
+class LangflowApiClient:
+    """HTTP client for Langflow API with streaming and non-streaming support.
+
+    This client handles all communication with Langflow APIs, including:
+    - Non-streaming execution
+    - Server-Sent Events (SSE) streaming
+    - Session management for conversation continuity
+    - Error handling and retries
+    - Credential management
+    """
+    flow_id: Incomplete
+    base_url: Incomplete
+    api_key: Incomplete
+    sessions: dict[str, str]
+    client_kwargs: Incomplete
+    def __init__(self, flow_id: str | None = None, base_url: str | None = None, api_key: str | None = None, httpx_client_options: HttpxClientOptions = None) -> None:
+        """Initialize the Langflow API client.
+
+        Args:
+            flow_id: The unique identifier of the Langflow flow to execute.
+            base_url: The base URL of the Langflow API server.
+            api_key: The API key for Langflow authentication.
+            httpx_client_options: HTTP client configuration options for httpx, including timeout.
+        """
+    def get_or_create_session(self, thread_id: str | None = None) -> str:
+        """Get existing session ID or create a new one.
+
+        Args:
+            thread_id: Optional thread ID for session mapping.
+
+        Returns:
+            Session ID for the conversation.
+        """
+    async def call_flow(self, input_value: str, session_id: str | None = None, flow_id: str | None = None, **_: Any) -> dict[str, Any]:
+        """Execute Langflow flow without streaming.
+
+        Args:
+            input_value: The user input to send to the flow.
+            session_id: Optional session ID for conversation continuity.
+            flow_id: Optional flow ID to execute. If None, uses the instance flow_id.
+            **_: Additional keyword arguments.
+
+        Returns:
+            The response from the flow execution.
+
+        Raises:
+            httpx.HTTPError: If the HTTP request fails.
+            ValueError: If the response cannot be parsed.
+        """
+    async def stream_flow(self, input_value: str, session_id: str | None = None, flow_id: str | None = None, **_: Any) -> AsyncGenerator[dict[str, Any], None]:
+        """Execute Langflow flow with streaming.
+
+        Args:
+            input_value: The user input to send to the flow.
+            session_id: Optional session ID for conversation continuity.
+            flow_id: Optional flow ID to execute. If None, uses the instance flow_id.
+            **_: Additional keyword arguments.
+
+        Yields:
+            Parsed streaming events from the Langflow API.
+
+        Raises:
+            httpx.HTTPError: If the HTTP request fails.
+            ValueError: If streaming events cannot be parsed.
+        """
+    def parse_stream_event(self, event_data: dict[str, Any]) -> dict[str, Any] | None:
+        """Parse a single streaming event from Langflow.
+
+        Args:
+            event_data: Raw event data from Langflow streaming response.
+
+        Returns:
+            Parsed event dictionary or None if event should be skipped.
+        """
+    def clear_session(self, thread_id: str) -> None:
+        """Clear session for a specific thread.
+
+        Args:
+            thread_id: Thread ID to clear session for.
+        """
+    def clear_all_sessions(self) -> None:
+        """Clear all stored sessions."""
+    async def get_flows(self, project_id: str | None = None, remove_example_flows: bool = False, components_only: bool = False, header_flows: bool = False, get_all: bool = True, page: int = 1, size: int = 50) -> list[dict[str, Any]]:
+        """Retrieve flows from Langflow API with full control over parameters.
+
+        Based on the official API docs: https://docs.langflow.org/api-flows
+        Uses the exact parameter format from the documentation.
+
+        Args:
+            project_id: Optional project ID to filter flows.
+            remove_example_flows: Whether to exclude example flows. Defaults to False.
+            components_only: Whether to return only components. Defaults to False.
+            header_flows: Whether to return only flow headers. Defaults to False.
+            get_all: Whether to return all flows (ignores pagination). Defaults to True.
+            page: Page number for pagination (ignored if get_all=True). Defaults to 1.
+            size: Number of flows per page (ignored if get_all=True). Defaults to 50.
+
+        Returns:
+            List of flows or flow headers from the Langflow API.
+
+        Raises:
+            httpx.HTTPError: If the HTTP request fails.
+            ValueError: If the response cannot be parsed or invalid parameters provided.
+        """
+    async def get_all_flows(self, project_id: str | None = None, remove_example_flows: bool = False, components_only: bool = False, header_flows: bool = False) -> list[dict[str, Any]]:
+        """Convenience method to get ALL flows using the backend's get_all=true feature.
+
+        This method is a simple wrapper around get_flows() with get_all=True,
+        which uses the Langflow backend's ability to return all flows in one request.
+
+        Args:
+            project_id: Optional project ID to filter flows.
+            remove_example_flows: Whether to exclude example flows. Defaults to False.
+            components_only: Whether to return only components. Defaults to False.
+            header_flows: Whether to return only flow headers. Defaults to False.
+
+        Returns:
+            List of all flows from the Langflow API.
+
+        Raises:
+            httpx.HTTPError: If the HTTP request fails.
+            ValueError: If the response cannot be parsed.
+        """
+    async def health_check(self) -> bool:
+        """Check if the Langflow API is accessible.
+
+        Returns:
+            True if the API is accessible, False otherwise.
+        """

aip_agents/clients/langflow/types.pyi

@@ -0,0 +1,7 @@
+from enum import StrEnum
+
+class LangflowEventType(StrEnum):
+    """Enum for Langflow event types as received from the API."""
+    ADD_MESSAGE: str
+    END: str
+    UNKNOWN: str

aip_agents/constants.pyi

@@ -4,3 +4,4 @@ BASE_DIR: Incomplete
 DEFAULT_AGENT_TIMEOUT: int
 LAST_N_CHATS: int
 MAX_MEMORY_TOKENS: int
+TEXT_PREVIEW_LENGTH: int

aip_agents/examples/compare_streaming_client.pyi

@@ -0,0 +1,48 @@
+from _typeshed import Incomplete
+from aip_agents.agent import LangGraphAgent as LangGraphAgent
+from aip_agents.examples.tools.mock_retrieval_tool import MockRetrievalTool as MockRetrievalTool
+from aip_agents.examples.tools.pii_demo_tools import get_customer_info as get_customer_info, get_employee_data as get_employee_data, get_user_profile as get_user_profile
+from aip_agents.examples.tools.random_chart_tool import RandomChartTool as RandomChartTool
+from aip_agents.examples.tools.table_generator_tool import TableGeneratorTool as TableGeneratorTool
+from aip_agents.schema.a2a import A2AStreamEventType as A2AStreamEventType
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
+from aip_agents.utils.logger import get_logger as get_logger
+from typing import Any
+
+logger: Incomplete
+SERVER_URL: str
+
+def create_local_agent(enable_token_streaming: bool = False) -> LangGraphAgent:
+    """Create a local agent with the same tools as the server.
+
+    Args:
+        enable_token_streaming: Whether to enable token streaming for content_chunk events.
+    """
+def format_chunk_summary(chunk: dict[str, Any], index: int) -> str:
+    """Format a chunk for display."""
+async def run_direct_streaming(agent: LangGraphAgent, query: str, pii_mapping: dict[str, str] | None = None) -> list[dict[str, Any]]:
+    """Run arun_sse_stream and collect all chunks."""
+async def run_connector_streaming(agent: LangGraphAgent, query: str, pii_mapping: dict[str, str] | None = None) -> list[dict[str, Any]]:
+    """Run astream_to_agent via A2A connector and collect all chunks."""
+def get_chunk_keys(chunks: list[dict[str, Any]]) -> set[str]:
+    """Get all unique keys across all chunks."""
+def get_field_types(chunks: list[dict[str, Any]], field: str) -> set[str]:
+    """Get all types seen for a field across chunks.
+
+    Note: StrEnum values are reported as 'str' since they ARE strings
+    and serialize identically over HTTP/SSE.
+    """
+def group_chunks_by_event_type(chunks: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
+    """Group chunks by their event_type."""
+def compare_chunk_structure(chunk1: dict[str, Any], chunk2: dict[str, Any]) -> dict[str, Any]:
+    """Compare structure of two chunks and return differences."""
+def compare_event_type_groups(direct_groups: dict[str, list[dict[str, Any]]], connector_groups: dict[str, list[dict[str, Any]]]) -> dict[str, Any]:
+    """Compare chunks grouped by event type - RAW comparison without filtering."""
+def compare_chunks(direct_chunks: list[dict[str, Any]], connector_chunks: list[dict[str, Any]]) -> dict[str, Any]:
+    """Compare the two chunk lists and produce a comprehensive structural summary."""
+def print_event_type_comparison(event_comparison: dict[str, Any]) -> None:
+    """Print per-event-type comparison details with exact key/value examples."""
+def print_structure_comparison(comparison: dict[str, Any]) -> None:
+    """Print detailed structure comparison."""
+async def main() -> None:
+    """Main comparison workflow."""

aip_agents/examples/compare_streaming_server.pyi

@@ -0,0 +1,18 @@
+from _typeshed import Incomplete
+from aip_agents.agent import LangGraphAgent as LangGraphAgent
+from aip_agents.examples.tools.mock_retrieval_tool import MockRetrievalTool as MockRetrievalTool
+from aip_agents.examples.tools.pii_demo_tools import get_customer_info as get_customer_info, get_employee_data as get_employee_data, get_user_profile as get_user_profile
+from aip_agents.examples.tools.random_chart_tool import RandomChartTool as RandomChartTool
+from aip_agents.examples.tools.table_generator_tool import TableGeneratorTool as TableGeneratorTool
+from aip_agents.utils.logger import get_logger as get_logger
+
+logger: Incomplete
+SERVER_AGENT_NAME: str
+
+def main(host: str, port: int):
+    """Runs the streaming comparison A2A server.
+
+    Args:
+        host: Host to bind the server to.
+        port: Port to bind the server to.
+    """

aip_agents/examples/demo_memory_recall.pyi

@@ -0,0 +1,58 @@
+from _typeshed import Incomplete
+from aip_agents.agent import LangGraphAgent as LangGraphAgent
+from datetime import datetime
+
+class MemoryTestResult:
+    """Container for test scenario results."""
+    JUDGE_SUCCESS_THRESHOLD: float
+    MEMORY_DISPLAY_LIMIT: int
+    scenario_name: Incomplete
+    query: Incomplete
+    response: Incomplete
+    memory_fetched: Incomplete
+    error: Incomplete
+    success: Incomplete
+    judge_verdict: Incomplete
+    judge_score: Incomplete
+    def __init__(self, scenario_name: str, query: str, response: str = None, memory_fetched: str = None, error: str = None) -> None:
+        """Initialize a memory test result.
+
+        Args:
+            scenario_name: Name of the test scenario.
+            query: The input query for the scenario.
+            response: The agent's response (optional).
+            memory_fetched: The raw memory retrieved from the tool (optional).
+            error: Any error message (optional).
+        """
+    def set_success_from_judge(self, verdict: str, score: float | None):
+        '''Set success using judge verdict and confidence score.
+
+        Success criteria: (verdict == "grounded" or verdict == "general_knowledge") and
+        score >= JUDGE_SUCCESS_THRESHOLD.
+
+        Args:
+            verdict (str): The judge\'s verdict ("grounded", "hallucination", or "general_knowledge").
+            score (float | None): The confidence score (0.0 to 1.0).
+        '''
+
+def print_test_summary(results: list[MemoryTestResult], start_time: datetime):
+    """Print a clean summary of all test results.
+
+    Args:
+        results (list[MemoryTestResult]): List of test results to summarize.
+        start_time (datetime): The start time of the tests.
+    """
+async def evaluate_with_judge(judge_agent: LangGraphAgent, query: str, response: str, memory_fetched: str = None) -> tuple[str | None, float | None]:
+    '''Evaluate if the response is grounded in memory using the judge agent.
+
+    Args:
+        judge_agent (LangGraphAgent): The judge agent to use for evaluation.
+        query (str): The original user query.
+        response (str): The agent\'s response to evaluate.
+        memory_fetched (str, optional): The memory that was retrieved. Defaults to None.
+
+    Returns:
+        tuple[str | None, float | None]: Tuple of (verdict, score) where verdict is "grounded", "hallucination", or "general_knowledge", and score is a confidence score (0.0-1.0).
+    '''
+async def run_all_scenarios() -> None:
+    """Run all memory recall test scenarios with a single reused agent and show clean summary."""

aip_agents/examples/hello_world_a2a_google_adk_client.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from aip_agents.agent.google_adk_agent import GoogleADKAgent as GoogleADKAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 

aip_agents/examples/hello_world_a2a_google_adk_client_agent.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from aip_agents.agent.google_adk_agent import GoogleADKAgent as GoogleADKAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 

aip_agents/examples/hello_world_a2a_google_adk_client_streaming.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from aip_agents.agent.google_adk_agent import GoogleADKAgent as GoogleADKAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 

aip_agents/examples/hello_world_a2a_google_adk_server.pyi

@@ -1,10 +1,15 @@
 from _typeshed import Incomplete
 from aip_agents.agent.google_adk_agent import GoogleADKAgent as GoogleADKAgent
 from aip_agents.examples.tools.adk_weather_tool import weather_tool as weather_tool
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 SERVER_AGENT_NAME: str
 
 def main(host: str, port: int):
-    """Runs the Google ADK Weather A2A server."""
+    """Runs the Google ADK Weather A2A server.
+
+    Args:
+        host (str): Host to bind the server to.
+        port (int): Port to bind the server to.
+    """

aip_agents/examples/hello_world_a2a_langchain_client.pyi

@@ -1,5 +1,5 @@
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
 
 def main() -> None:
     """Main function demonstrating the General Assistant agent with A2A capabilities."""

aip_agents/examples/hello_world_a2a_langchain_client_agent.pyi

@@ -1,5 +1,5 @@
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
 
 def main() -> None:
     """Main function demonstrating the General Assistant agent with A2A capabilities."""

aip_agents/examples/hello_world_a2a_langchain_client_lm_invoker.pyi

@@ -1,5 +1,5 @@
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
 
 def main() -> None:
     """Main function demonstrating the General Assistant agent with A2A capabilities."""

aip_agents/examples/hello_world_a2a_langchain_client_streaming.pyi

@@ -1,5 +1,5 @@
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
 
 async def main() -> None:
     """Main function demonstrating the General Assistant agent with streaming A2A capabilities."""

aip_agents/examples/hello_world_a2a_langchain_reference_client_streaming.pyi

@@ -1,5 +1,5 @@
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.agent.types import A2AClientConfig as A2AClientConfig
+from aip_agents.schema.agent import A2AClientConfig as A2AClientConfig
 
 async def main() -> None:
     """Main function demonstrating streaming A2A client for the reference server."""

aip_agents/examples/hello_world_a2a_langchain_reference_server.pyi

@@ -1,10 +1,15 @@
 from _typeshed import Incomplete
 from aip_agents.agent import LangChainAgent as LangChainAgent
-from aip_agents.examples.tools import google_serper_tool as google_serper_tool, time_tool as time_tool
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.examples.tools import google_serper_tool as google_serper_tool, mock_retrieval_tool as mock_retrieval_tool, time_tool as time_tool
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 SERVER_AGENT_NAME: str
 
 def main(host: str, port: int):
-    """Runs the LangChain Weather A2A server."""
+    """Runs the LangChain Mock Retrieval and Google Search A2A server.
+
+    Args:
+        host (str): Host to bind the server to.
+        port (int): Port to bind the server to.
+    """

aip_agents/examples/hello_world_a2a_langchain_server.pyi

@@ -1,10 +1,15 @@
 from _typeshed import Incomplete
 from aip_agents.agent import LangChainAgent as LangChainAgent
 from aip_agents.examples.tools.langchain_weather_tool import weather_tool as weather_tool
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
+from aip_agents.utils.logger import get_logger as get_logger
 
 logger: Incomplete
 SERVER_AGENT_NAME: str
 
 def main(host: str, port: int):
-    """Runs the LangChain Weather A2A server."""
+    """Runs the LangChain Weather A2A server.
+
+    Args:
+        host (str): Host to bind the server to.
+        port (int): Port to bind the server to.
+    """