aip-agents-binary 0.5.20__py3-none-any.whl

aip_agents/examples/tools/langchain_graph_artifact_tool.py

@@ -0,0 +1,172 @@
+"""This is a tool that generates a tiny bar/line chart and returns it as a Command artifact.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+"""
+
+# Optional imports for chart rendering
+import io
+
+from langchain_core.tools import BaseTool
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from aip_agents.a2a.types import MimeType
+from aip_agents.utils.artifact_helpers import create_artifact_command
+
+try:
+    from PIL import Image, ImageDraw  # type: ignore
+
+    PIL_AVAILABLE = True
+except ImportError:
+    PIL_AVAILABLE = False
+    Image = None
+    ImageDraw = None
+
+
+class GraphCommandInput(BaseModel):
+    """Input schema for a tiny bar/line chart artifact tool.
+
+    Note: width/height are clamped to a small size to keep artifacts lightweight.
+    """
+
+    chart_type: str = Field(
+        default="bar",
+        description="Type of chart to generate (bar or line).",
+        pattern="^(bar|line)$",
+    )
+    labels: list[str] = Field(
+        default_factory=lambda: ["A", "B", "C"],
+        description="Labels for categories",
+    )
+    values: list[float] = Field(
+        default_factory=lambda: [10.0, 20.0, 15.0],
+        description="Values corresponding to labels",
+    )
+    title: str = Field(default="Chart", description="Chart title (metadata only)")
+    width: int = Field(default=64, ge=16, le=128, description="Desired image width in px (will be clamped)")
+    height: int = Field(default=64, ge=16, le=128, description="Desired image height in px (will be clamped)")
+    image_name: str = Field(default="chart", description="Base filename (without extension)")
+
+
+class GraphArtifactCommandTool(BaseTool):
+    """Generate a very small PNG chart and return it as a Command artifact."""
+
+    name: str = "graph_generator_command"
+    description: str = "Generate a tiny bar/line chart and return a lightweight PNG artifact using Command."
+    args_schema: type[BaseModel] = GraphCommandInput
+
+    def _run(  # noqa: PLR0913
+        self,
+        chart_type: str = "bar",
+        labels: list[str] | None = None,
+        values: list[float] | None = None,
+        title: str = "Chart",
+        width: int = 64,
+        height: int = 64,
+        image_name: str = "chart",
+    ) -> Command:
+        labels = labels or ["A", "B", "C"]
+        values = values or [10.0, 20.0, 15.0]
+
+        if not labels or len(labels) != len(values):
+            return Command(update={"result": "Invalid input: labels and values must be non-empty and equal length."})
+
+        # Clamp to very small size to keep base64 short
+        img_w = max(16, min(int(width), 64))
+        img_h = max(16, min(int(height), 64))
+
+        png_bytes = self._render_tiny_chart(chart_type, labels, values, img_w, img_h)
+
+        result_text = f"Generated tiny {chart_type} chart '{title}' with {len(labels)} points as {image_name}.png"
+        metadata_delta = {
+            "last_chart": {
+                "name": f"{image_name}.png",
+                "chart_type": chart_type,
+                "title": title,
+                "num_points": len(labels),
+                "size": f"{img_w}x{img_h}",
+            }
+        }
+
+        return create_artifact_command(
+            result=result_text,
+            artifact_data=png_bytes,
+            artifact_name=f"{image_name}.png",
+            artifact_description=f"Tiny {chart_type} chart: {title}",
+            mime_type=MimeType.IMAGE_PNG,
+            metadata_update=metadata_delta,
+        )
+
+    @staticmethod
+    def _render_tiny_chart(
+        chart_type: str,
+        labels: list[str],
+        values: list[float],
+        width: int,
+        height: int,
+    ) -> bytes:
+        """Render a tiny chart using only PIL to keep artifact size minimal.
+
+        Args:
+            chart_type (str): Type of chart to render ("bar" or "line").
+            labels (list[str]): Labels for the data points.
+            values (list[float]): Values to plot.
+            width (int): Width of the chart image in pixels.
+            height (int): Height of the chart image in pixels.
+
+        Returns:
+            bytes: PNG image data as bytes.
+        """
+        if not PIL_AVAILABLE:
+            raise ImportError("PIL is not available for chart rendering")
+
+        # Canvas and basic layout
+        image = Image.new("RGB", (width, height), color="white")
+        draw = ImageDraw.Draw(image)
+
+        # Small margins for axes
+        margin_left = max(3, width // 12)
+        margin_right = max(2, width // 20)
+        margin_top = max(3, height // 12)
+        margin_bottom = max(6, height // 8)
+
+        plot_w = max(1, width - margin_left - margin_right)
+        plot_h = max(1, height - margin_top - margin_bottom)
+
+        # Axes
+        x0, y0 = margin_left, height - margin_bottom
+        x1, y1 = width - margin_right, margin_top
+        draw.line((x0, y0, x1, y0), fill="black")
+        draw.line((x0, y0, x0, y1), fill="black")
+
+        # Normalize values to plot height
+        max_val = max(values + [1.0])
+        n = len(values)
+        if chart_type == "bar":
+            # Thin bars to keep readable at tiny sizes
+            bar_spacing = plot_w / max(n, 1)
+            bar_w = max(1, int(bar_spacing * 0.6))
+            for i, val in enumerate(values):
+                x_center = x0 + int((i + 0.5) * bar_spacing)
+                bar_h = int((val / max_val) * plot_h)
+                x_left = x_center - bar_w // 2
+                x_right = x_center + bar_w // 2
+                y_top = y0 - bar_h
+                draw.rectangle((x_left, y_top, x_right, y0), fill="#4C78A8")
+        else:
+            # Polyline with small points
+            points = []
+            step = plot_w / max(n, 1)
+            for i, val in enumerate(values):
+                x = x0 + int((i + 0.5) * step)
+                y = y0 - int((val / max_val) * plot_h)
+                points.append((x, y))
+            if len(points) >= 2:  # noqa: PLR2004
+                draw.line(points, fill="#F58518")
+            for px, py in points:
+                draw.ellipse((px - 1, py - 1, px + 1, py + 1), fill="#F58518")
+
+        buf = io.BytesIO()
+        image.save(buf, format="PNG", optimize=True)
+        return buf.getvalue()

aip_agents/examples/tools/langchain_weather_tool.py

@@ -0,0 +1,48 @@
+"""Defines a weather tool that can be used to get the weather for a specified city.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+"""
+
+from langchain_core.tools import tool
+from pydantic import BaseModel
+
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class WeatherToolInputSchema(BaseModel):
+    """Schema for weather tool input."""
+
+    city: str
+
+
+@tool(args_schema=WeatherToolInputSchema)
+def weather_tool(city: str) -> str:
+    """Gets the weather for a specified city.
+
+    Args:
+        city: The name of the city to get weather for.
+
+    Returns:
+        A string describing the weather conditions.
+    """
+    weather_data = {
+        "Jakarta": "32°C, Partly cloudy with high humidity",
+        "Singapore": "30°C, Scattered thunderstorms",
+        "Tokyo": "25°C, Clear skies",
+        "London": "18°C, Light rain",
+        "New York": "22°C, Sunny",
+    }
+
+    city_name = city.strip().title()
+
+    weather = weather_data.get(city_name)
+    if weather:
+        logger.info(f"Found weather for {city_name}: {weather}")
+        return weather
+    else:
+        message = f"Weather data not available for {city}"
+        logger.warning(message)
+        return message

aip_agents/examples/tools/langgraph_streaming_tool.py

@@ -0,0 +1,130 @@
+"""Tool that wraps a LangGraph agent with time and weather tools.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+import os
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field, PrivateAttr
+
+from aip_agents.agent.langgraph_react_agent import LangGraphReactAgent
+from aip_agents.examples.tools.langchain_weather_tool import weather_tool
+from aip_agents.examples.tools.time_tool import TimeTool
+
+
+class StreamingToolConfig(BaseModel):
+    """Tool configuration schema for the LangGraph streaming tool."""
+
+    time_format: str = Field(
+        default="%m/%d/%y %H:%M:%S",
+        description="DateTime format for time-related queries (e.g., '%Y-%m-%d %H:%M:%S', '%A, %B %d, %Y')",
+    )
+
+
+class LangGraphStreamingToolInput(BaseModel):
+    """Input schema for the LangGraphStreamingTool."""
+
+    query: str = Field(..., description="Query prompt for the LangGraph agent to execute")
+
+
+class LangGraphStreamingTool(BaseTool):
+    """Tool that wraps a LangGraph agent with time and weather forecast capabilities."""
+
+    name: str = "langgraph_streaming_tool"
+    description: str = (
+        "Execute tasks using a LangGraph agent with time and weather capabilities. "
+        "Can get current time in various formats and weather information for specific cities. "
+        "Supports configuration for time format settings."
+    )
+    args_schema: type[BaseModel] = LangGraphStreamingToolInput
+    tool_config_schema: type[BaseModel] = StreamingToolConfig
+    _agent: Any = PrivateAttr()
+
+    def __init__(self, model: Any = None, **kwargs):
+        """Initialize the LangGraphStreamingTool.
+
+        Args:
+            model: The model to use for the agent.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(**kwargs)
+
+        if model is None:
+            model = os.getenv("DEFAULT_MODEL", "openai/gpt-4o")
+
+        self._agent = LangGraphReactAgent(
+            name="internal_time_weather_agent",
+            instruction="You are a helpful assistant with access to time and weather tools. "
+            "When users ask for both time and weather information, use BOTH tools: "
+            "1. Use the time_tool to get current time information "
+            "2. Use the weather_tool to get weather for specific cities "
+            "Always provide complete answers using all relevant tools.",
+            model=model,
+            tools=[TimeTool(), weather_tool],
+            description="Internal agent for time and weather tasks",
+        )
+
+    @property
+    def agent(self):
+        """Access the internal agent."""
+        return self._agent
+
+    def _run(self, query: str) -> str:
+        """Run the tool synchronously.
+
+        Args:
+            query: The query to execute.
+
+        Returns:
+            The final output from the tool execution.
+        """
+        result = self.agent.run(query)
+        if isinstance(result, dict) and "output" in result:
+            return str(result["output"])
+        return str(result)
+
+    async def _arun(self, query: str) -> str:
+        """Run the tool asynchronously.
+
+        Args:
+            query: The query to execute.
+
+        Returns:
+            The final output from the tool execution.
+        """
+        result = await self.agent.arun(query)
+        if isinstance(result, dict) and "output" in result:
+            return str(result["output"])
+        return str(result)
+
+    async def arun_streaming(self, query: str = None, config: RunnableConfig = None, **kwargs):
+        """Execute the LangGraph agent asynchronously with A2A streaming output.
+
+        Args:
+            query: The query to execute.
+            config: Tool configuration containing user context and preferences.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            An async generator that yields streaming events.
+        """
+        if query is None:
+            query = kwargs.get("query") or kwargs.get("task")
+
+        effective_config = self.get_tool_config(config)
+
+        if not effective_config:
+            effective_config = StreamingToolConfig(time_format="%m/%d/%y %H:%M:%S")
+
+        try:
+            enhanced_query = f"[Time format: {effective_config.time_format}] {query}"
+
+            async for event in self.agent.arun_a2a_stream(enhanced_query, **kwargs):
+                if event is not None and isinstance(event, dict):
+                    yield event
+        except Exception as e:
+            yield {"event_type": "error", "content": f"Error in agent wrapper streaming: {str(e)}", "is_final": True}

aip_agents/examples/tools/mock_retrieval_tool.py

@@ -0,0 +1,56 @@
+"""Simple mock tool that returns references for testing reference propagation.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from gllm_core.schema import Chunk
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+
+class MockRetrievalInput(BaseModel):
+    """Input schema for the MockRetrievalTool."""
+
+    query: str = Field(..., description="Search query")
+
+
+class MockRetrievalTool(BaseTool):
+    """Mock tool that returns hardcoded references for testing."""
+
+    name: str = "mock_retrieval"
+    description: str = "Mock retrieval tool that returns sample references for testing reference propagation."
+    args_schema: type[BaseModel] = MockRetrievalInput
+    save_output_history: bool = True
+
+    def _run(self, query: str) -> str:
+        """Return a simple mock response.
+
+        Args:
+            query (str): The search query to process.
+
+        Returns:
+            str: Mock retrieval result for the query.
+        """
+        return f"Mock retrieval result for: {query}"
+
+    def _format_agent_reference(self, tool_output: str) -> list[Chunk]:
+        """Return hardcoded references for testing.
+
+        Args:
+            tool_output (str): The tool output to format as references.
+
+        Returns:
+            list[Chunk]: List of formatted chunks for the query.
+        """
+        query = tool_output.split(":")[-1].strip()
+        return [
+            Chunk(
+                content=f"Mock reference 1 for query: {query}",
+                metadata={"tool_name": "mock_retrieval", "source": "Mock Database", "title": "Mock Reference 1"},
+            ),
+            Chunk(
+                content=f"Mock reference 2 for query: {query}",
+                metadata={"tool_name": "mock_retrieval", "source": "Mock Database", "title": "Mock Reference 2"},
+            ),
+        ]

aip_agents/examples/tools/pii_demo_tools.py

@@ -0,0 +1,189 @@
+"""Mock tools for demonstrating PII handling in parallel tool execution.
+
+This module provides tools that return personally identifiable information (PII)
+to demonstrate how the PII handler automatically anonymizes and deanonymizes data
+during tool execution.
+
+Tools included:
+- get_customer_info: Returns customer information including email and phone
+- get_employee_data: Returns employee data including name and salary
+- get_user_profile: Returns user profile with personal details
+
+Authors:
+    Cascade AI Assistant
+"""
+
+from langchain_core.tools import tool
+from pydantic import BaseModel
+
+from aip_agents.utils.logger import LoggerManager
+
+logger = LoggerManager().get_logger(__name__)
+
+
+class CustomerIdSchema(BaseModel):
+    """Schema for customer ID input."""
+
+    customer_id: str
+
+
+class EmployeeIdSchema(BaseModel):
+    """Schema for employee ID input."""
+
+    employee_id: str
+
+
+class UserIdSchema(BaseModel):
+    """Schema for user ID input."""
+
+    user_id: str
+
+
+@tool(args_schema=CustomerIdSchema)
+def get_customer_info(customer_id: str) -> str:
+    """Gets customer information including email and phone number.
+
+    This tool demonstrates PII handling by returning sensitive customer data
+    that will be automatically anonymized by the PII handler.
+
+    Args:
+        customer_id: The ID of the customer to retrieve information for.
+
+    Returns:
+        A string containing customer information with PII.
+    """
+    customer_data = {
+        "C001": {
+            "name": "John Smith",
+            "email": "john.smith@example.com",
+            "phone": "+1-555-0101",
+            "address": "123 Main St, New York, NY 10001",
+        },
+        "C002": {
+            "name": "Alice Johnson",
+            "email": "alice.johnson@example.com",
+            "phone": "+1-555-0102",
+            "address": "456 Oak Ave, Los Angeles, CA 90001",
+        },
+        "C003": {
+            "name": "Bob Williams",
+            "email": "bob.williams@example.com",
+            "phone": "+1-555-0103",
+            "address": "789 Pine Rd, Chicago, IL 60601",
+        },
+    }
+
+    customer = customer_data.get(customer_id.upper())
+    if customer:
+        logger.info(f"Retrieved customer info for {customer_id}")
+        return (
+            f"Customer {customer_id}: Name: {customer['name']}, "
+            f"Email: {customer['email']}, Phone: {customer['phone']}, "
+            f"Address: {customer['address']}"
+        )
+    else:
+        message = f"Customer {customer_id} not found"
+        logger.warning(message)
+        return message
+
+
+@tool(args_schema=EmployeeIdSchema)
+def get_employee_data(employee_id: str) -> str:
+    """Gets employee data including name, email, and salary information.
+
+    This tool demonstrates PII handling with employee-specific data that includes
+    sensitive information like salary and personal email addresses.
+
+    Args:
+        employee_id: The ID of the employee to retrieve data for.
+
+    Returns:
+        A string containing employee information with PII.
+    """
+    employee_data = {
+        "E001": {
+            "name": "Carol Davis",
+            "email": "carol.davis@company.com",
+            "phone": "+1-555-0201",
+            "salary": "$85,000",
+            "department": "Engineering",
+        },
+        "E002": {
+            "name": "David Miller",
+            "email": "david.miller@company.com",
+            "phone": "+1-555-0202",
+            "salary": "$92,000",
+            "department": "Product",
+        },
+        "E003": {
+            "name": "Emma Wilson",
+            "email": "emma.wilson@company.com",
+            "phone": "+1-555-0203",
+            "salary": "$78,000",
+            "department": "Marketing",
+        },
+    }
+
+    employee = employee_data.get(employee_id.upper())
+    if employee:
+        logger.info(f"Retrieved employee data for {employee_id}")
+        return (
+            f"Employee {employee_id}: Name: {employee['name']}, "
+            f"Email: {employee['email']}, Phone: {employee['phone']}, "
+            f"Salary: {employee['salary']}, Department: {employee['department']}"
+        )
+    else:
+        message = f"Employee {employee_id} not found"
+        logger.warning(message)
+        return message
+
+
+@tool(args_schema=UserIdSchema)
+def get_user_profile(user_id: str) -> str:
+    """Gets user profile information with personal details.
+
+    This tool demonstrates PII handling with user profile data that includes
+    personal information like email, phone, and date of birth.
+
+    Args:
+        user_id: The ID of the user to retrieve profile for.
+
+    Returns:
+        A string containing user profile information with PII.
+    """
+    user_data = {
+        "U001": {
+            "username": "john_doe",
+            "email": "john.doe@personal.com",
+            "phone": "+1-555-0301",
+            "dob": "1990-05-15",
+            "city": "Seattle, WA",
+        },
+        "U002": {
+            "username": "jane_smith",
+            "email": "jane.smith@personal.com",
+            "phone": "+1-555-0302",
+            "dob": "1992-08-22",
+            "city": "Portland, OR",
+        },
+        "U003": {
+            "username": "bob_jones",
+            "email": "bob.jones@personal.com",
+            "phone": "+1-555-0303",
+            "dob": "1988-03-10",
+            "city": "San Francisco, CA",
+        },
+    }
+
+    user = user_data.get(user_id.upper())
+    if user:
+        logger.info(f"Retrieved user profile for {user_id}")
+        return (
+            f"User {user_id}: Username: {user['username']}, "
+            f"Email: {user['email']}, Phone: {user['phone']}, "
+            f"DOB: {user['dob']}, City: {user['city']}"
+        )
+    else:
+        message = f"User {user_id} not found"
+        logger.warning(message)
+        return message