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

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

aip_agents/examples/tools/pii_demo_tools.pyi

@@ -0,0 +1,54 @@
+from _typeshed import Incomplete
+from aip_agents.utils.logger import LoggerManager as LoggerManager
+from pydantic import BaseModel
+
+logger: Incomplete
+
+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
+
+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.
+    """
+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.
+    """
+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.
+    """

aip_agents/examples/tools/random_chart_tool.py

@@ -0,0 +1,142 @@
+"""Tool that generates random bar chart artifacts for streaming demos."""
+
+import io
+import random
+from typing import Any
+
+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
+
+DEFAULT_RANDOM_CHART_TITLE = "Random Insights"
+
+try:
+    from PIL import Image, ImageDraw  # type: ignore
+
+    PIL_AVAILABLE = True
+except ImportError:  # pragma: no cover - PIL optional
+    PIL_AVAILABLE = False
+    Image = None
+    ImageDraw = None
+
+
+class RandomChartInput(BaseModel):
+    """Input schema for random chart generation."""
+
+    title: str = Field(
+        default=DEFAULT_RANDOM_CHART_TITLE,
+        description="Title rendered at the top of the chart.",
+    )
+    num_bars: int = Field(
+        default=5,
+        ge=3,
+        le=10,
+        description="How many bars to render.",
+    )
+    min_value: int = Field(
+        default=10,
+        ge=0,
+        description="Minimum possible value for a bar.",
+    )
+    max_value: int = Field(
+        default=100,
+        gt=10,
+        description="Maximum possible value for a bar.",
+    )
+
+
+class RandomChartTool(BaseTool):
+    """Generate random bar chart images without relying on upstream data."""
+
+    name: str = "random_chart_tool"
+    description: str = "Create a random bar chart image artifact to showcase image streaming."
+    args_schema: type[BaseModel] = RandomChartInput
+
+    def _run(  # noqa: D401
+        self,
+        title: str = DEFAULT_RANDOM_CHART_TITLE,
+        num_bars: int = 5,
+        min_value: int = 10,
+        max_value: int = 100,
+        **kwargs: Any,
+    ) -> Command:
+        """Generate the chart synchronously."""
+        if min_value >= max_value:
+            return Command(
+                update={"result": "❌ min_value must be less than max_value", "metadata": {"error": "invalid_range"}}
+            )
+
+        rng = random.Random()
+        values = [rng.randint(min_value, max_value) for _ in range(num_bars)]
+        labels = [f"Bar {i + 1}" for i in range(num_bars)]
+
+        if PIL_AVAILABLE:
+            image_bytes = self._create_image_with_pil(title, labels, values)
+        else:
+            image_bytes = self._create_text_artifact(title, labels, values)
+
+        artifact_mime_type = MimeType.IMAGE_PNG if PIL_AVAILABLE else MimeType.TEXT_PLAIN
+        return create_artifact_command(
+            result=f"📊 Generated random bar chart '{title}' with {num_bars} bars",
+            artifact_data=image_bytes,
+            artifact_name=f"{title.lower().replace(' ', '_')}_random_chart.png",
+            artifact_description=f"Random bar chart ({num_bars} bars)",
+            mime_type=artifact_mime_type,
+            metadata_update={
+                "visualization": {
+                    "chart_type": "bar",
+                    "title": title,
+                    "data_points": num_bars,
+                    "randomized": True,
+                }
+            },
+        )
+
+    async def _arun(
+        self,
+        title: str = DEFAULT_RANDOM_CHART_TITLE,
+        num_bars: int = 5,
+        min_value: int = 10,
+        max_value: int = 100,
+        **kwargs: Any,
+    ) -> Command:
+        """Async wrapper for random chart generation."""
+        return self._run(title=title, num_bars=num_bars, min_value=min_value, max_value=max_value, **kwargs)
+
+    def _create_image_with_pil(self, title: str, labels: list[str], values: list[int]) -> bytes:
+        width, height = 400, 300
+        image = Image.new("RGB", (width, height), "white")
+        draw = ImageDraw.Draw(image)
+        draw.text((10, 10), title[:40], fill="black")
+
+        chart_x, chart_y = 40, 50
+        chart_width = width - 80
+        chart_height = height - 100
+
+        max_val = max(values)
+        min_val = min(values)
+        val_range = max_val - min_val if max_val != min_val else 1
+        bar_width = chart_width // len(values)
+
+        for idx, (label, value) in enumerate(zip(labels, values, strict=False)):
+            x = chart_x + idx * bar_width
+            normalized = (value - min_val) / val_range
+            y = chart_y + chart_height - (normalized * chart_height)
+
+            draw.rectangle([x + 2, y, x + bar_width - 4, chart_y + chart_height], fill="#4C78A8", outline="black")
+            draw.text((x + 2, chart_y + chart_height + 5), label[:8], fill="black")
+
+        buf = io.BytesIO()
+        image.save(buf, format="PNG")
+        return buf.getvalue()
+
+    def _create_text_artifact(self, title: str, labels: list[str], values: list[int]) -> bytes:
+        max_val = max(values)
+        lines = [f"=== {title} ===", "", "Random bar chart artifact", ""]
+        for label, value in zip(labels, values, strict=False):
+            length = int((value / max_val) * 30) if max_val else 0
+            lines.append(f"{label:>8}: {'█' * length} ({value})")
+        return "\n".join(lines).encode("utf-8")

aip_agents/examples/tools/random_chart_tool.pyi

@@ -0,0 +1,20 @@
+from aip_agents.a2a.types import MimeType as MimeType
+from aip_agents.utils.artifact_helpers import create_artifact_command as create_artifact_command
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel
+
+DEFAULT_RANDOM_CHART_TITLE: str
+PIL_AVAILABLE: bool
+
+class RandomChartInput(BaseModel):
+    """Input schema for random chart generation."""
+    title: str
+    num_bars: int
+    min_value: int
+    max_value: int
+
+class RandomChartTool(BaseTool):
+    """Generate random bar chart images without relying on upstream data."""
+    name: str
+    description: str
+    args_schema: type[BaseModel]

aip_agents/examples/tools/serper_tool.py

@@ -0,0 +1,202 @@
+"""Tool to search Google Serper API.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+import json
+import logging
+from typing import Any
+
+from gllm_core.schema import Chunk
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleSerperInput(BaseModel):
+    """Input schema for the GoogleSerperTool."""
+
+    query: str = Field(..., description="Search query")
+
+
+class MockGoogleSerperTool(BaseTool):
+    """Mock Tool to simulate Google Serper API results for testing."""
+
+    name: str = "google_serper"
+    description: str = """
+    Useful for searching the web using the Google Serper API (mocked).
+    Input should be a search query.
+    """
+    save_output_history: bool = True
+    args_schema: type[BaseModel] = GoogleSerperInput
+
+    def _run(self, query: str) -> str:
+        """Return a hardcoded mock response as a JSON string.
+
+        Args:
+            query (str): The search query to process.
+
+        Returns:
+            str: JSON string containing mock search results.
+        """
+        normalized_query = query.lower()
+
+        # Special-case oncology AI drug discovery queries so the agent
+        # naturally sees stock-oriented company context and can call
+        # stock tools (get_stock_price, get_stock_news) afterwards.
+        if "oncology" in normalized_query or "drug discovery" in normalized_query:
+            mock_result = {
+                "organic": [
+                    {
+                        "title": "NVIDIA (NVDA) - AI infrastructure for healthcare and drug discovery",
+                        "link": "https://nvidia.example.com/healthcare-ai",
+                        "snippet": (
+                            "NVIDIA (ticker: NVDA) provides GPU-accelerated platforms widely used in AI-driven "
+                            "drug discovery and oncology research. Investors closely watch NVDA stock as demand "
+                            "for healthcare and life-sciences AI workloads grows."
+                        ),
+                    },
+                    {
+                        "title": "Microsoft (MSFT) - Azure AI partnerships with pharma and oncology",
+                        "link": "https://microsoft.example.com/azure-health",
+                        "snippet": (
+                            "Microsoft (MSFT) collaborates with leading pharma companies to apply Azure AI to "
+                            "oncology and drug discovery workflows. Analysts frequently reference MSFT stock when "
+                            "discussing enterprise AI in healthcare."
+                        ),
+                    },
+                    {
+                        "title": "Apple (AAPL) - Devices and AI ecosystems in digital health",
+                        "link": "https://apple.example.com/health-ai",
+                        "snippet": (
+                            "Apple (AAPL) integrates AI into health and wellness features that support oncology "
+                            "patients and clinical research. Market commentary often links AAPL stock to long-term "
+                            "growth in digital health and AI."
+                        ),
+                    },
+                ],
+                "news": [
+                    {
+                        "title": "NVDA rallies on new healthcare AI partnerships",
+                        "link": "https://news.example.com/nvda-oncology-ai",
+                        "snippet": (
+                            "NVIDIA (NVDA) announced expanded collaborations focused on oncology and AI-driven "
+                            "drug discovery, prompting renewed interest in NVDA stock among growth investors."
+                        ),
+                    },
+                    {
+                        "title": "MSFT deepens AI oncology collaborations on Azure",
+                        "link": "https://news.example.com/msft-azure-oncology",
+                        "snippet": (
+                            "Microsoft (MSFT) reported new Azure AI partnerships with oncology research centers, "
+                            "and analysts highlighted MSFT stock as a key AI infrastructure play in healthcare."
+                        ),
+                    },
+                    {
+                        "title": "AAPL explores AI-enabled health insights for cancer care",
+                        "link": "https://news.example.com/aapl-health-oncology",
+                        "snippet": (
+                            "Apple (AAPL) is reportedly piloting AI-enabled health features that could support "
+                            "oncology patient monitoring, adding another angle to the long-term AAPL stock story."
+                        ),
+                    },
+                ],
+            }
+        else:
+            # Default mock result used for generic queries.
+            mock_result = {
+                "organic": [
+                    {
+                        "title": "NeoAI - Artificial Intelligence Research",
+                        "link": "https://neoai.example.com/",
+                        "snippet": (
+                            "NeoAI is an AI research and deployment company. Our mission is to ensure that "
+                            "artificial general intelligence benefits all of humanity."
+                        ),
+                    },
+                    {
+                        "title": "Wikipedia - NeoAI",
+                        "link": "https://en.wikipedia.org/wiki/NeoAI",
+                        "snippet": (
+                            "NeoAI is a fictional artificial intelligence research organization consisting of the "
+                            "for-profit NeoAI LP and its parent company, the non-profit NeoAI Foundation."
+                        ),
+                    },
+                ],
+                "news": [
+                    {
+                        "title": "NeoAI unveils new LLM model",
+                        "link": "https://news.example.com/neoai-llm",
+                        "snippet": (
+                            "NeoAI has announced the release of NeoAI-LLM, a new large multimodal model that accepts "
+                            "image and text inputs."
+                        ),
+                    }
+                ],
+            }
+
+        return json.dumps(mock_result)
+
+    def _format_agent_reference(self, tool_output: str) -> list[Chunk]:
+        parsed_output = self._parse_tool_output(tool_output)
+        if not parsed_output:
+            return []
+        formatted_chunks = []
+        file_id_counter = 0
+        for section_name, section_data in parsed_output.items():
+            if not isinstance(section_data, list):
+                continue
+            results = self._process_section_items(section_name, section_data, file_id_counter)
+            formatted_chunks.extend(results["chunks"])
+            file_id_counter = results["counter"]
+        return formatted_chunks
+
+    def _parse_tool_output(self, tool_output: str) -> dict[str, Any]:
+        if isinstance(tool_output, str):
+            try:
+                return json.loads(tool_output)
+            except json.JSONDecodeError:
+                logger.error("Error: Unable to parse tool_output as JSON when formatting agent references.")
+                return {}
+        return tool_output
+
+    def _process_section_items(self, section_name: str, section_data: list[dict], start_counter: int) -> dict[str, Any]:
+        chunks = []
+        counter = start_counter
+        for item in section_data:
+            if not isinstance(item, dict):
+                continue
+            try:
+                chunk = self._create_chunk_from_item(section_name, item, counter)
+                if chunk:
+                    chunks.append(chunk)
+                    counter += 1
+            except Exception as e:
+                logger.error(f"Error processing {section_name} result: {e}")
+        return {"chunks": chunks, "counter": counter}
+
+    def _create_chunk_from_item(self, section_name: str, item: dict[str, Any], file_id: int) -> Chunk | None:
+        link = item.get("link")
+        if not link:
+            logger.warning(f"Skipping item {file_id} from {section_name} result: Missing link")
+            return None
+        content = ""
+        if "snippet" in item:
+            content = item["snippet"]
+        elif "title" in item:
+            content = item["title"]
+        if not content:
+            logger.warning(f"Skipping item {file_id} from {section_name} result: Missing content")
+            return None
+        metadata = {
+            "source": item.get("title", "Untitled Source"),
+            "section_type": section_name,
+            "source_type": "website",
+            "title": item.get("title", "Untitled"),
+            "link": link,
+            "file_id": str(file_id),
+        }
+        return Chunk(content=content, metadata=metadata)

aip_agents/examples/tools/serper_tool.pyi

@@ -0,0 +1,16 @@
+from _typeshed import Incomplete
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel
+
+logger: Incomplete
+
+class GoogleSerperInput(BaseModel):
+    """Input schema for the GoogleSerperTool."""
+    query: str
+
+class MockGoogleSerperTool(BaseTool):
+    """Mock Tool to simulate Google Serper API results for testing."""
+    name: str
+    description: str
+    save_output_history: bool
+    args_schema: type[BaseModel]