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

aip_agents/tools/web_search/serper_tool.py

@@ -0,0 +1,187 @@
+"""Tool to search Google Serper API.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+import json
+from json import dumps
+from typing import Any
+
+from gllm_core.schema import Chunk
+from langchain_community.utilities.google_serper import GoogleSerperAPIWrapper
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class GoogleSerperInput(BaseModel):
+    """Input schema for the GoogleSerperTool."""
+
+    query: str = Field(..., description="Search query", min_length=1)
+
+
+class GoogleSerperTool(BaseTool):
+    """Tool to search Google Serper API."""
+
+    name: str = "google_serper"
+    description: str = """
+    Useful for searching the web using the Google Serper API.
+    Input should be a search query.
+    """
+    save_output_history: bool = Field(default=True)
+    args_schema: type[BaseModel] = GoogleSerperInput
+    api_wrapper: GoogleSerperAPIWrapper = Field(default_factory=GoogleSerperAPIWrapper, exclude=True)
+
+    def _run(
+        self,
+        query: str,
+    ) -> str:
+        """Executes a query using the API wrapper and returns the result as a JSON string.
+
+        Args:
+            query (str): The query string to be executed.
+            run_manager (Optional[CallbackManagerForToolRun], optional): An optional callback manager for the tool run.
+                Defaults to None.
+
+        Returns:
+            str: The result of the query execution, serialized as a JSON string.
+        """
+        result = self.api_wrapper.results(query)
+        return dumps(result)
+
+    def _format_agent_reference(self, tool_output: str) -> list[Chunk]:
+        """Format tool output into agent references.
+
+        Args:
+            tool_output (str): The output from the tool, which is expected to be a JSON string.
+
+        Returns:
+            list[Chunk]: Formatted tool output as agent references
+        """
+        # Parse the tool output if it's a string
+        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():
+            # Process only list-type sections
+            if not isinstance(section_data, list):
+                continue
+
+            # Process each item in the section
+            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]:
+        """Parses the tool output to ensure it is in dictionary format.
+
+        This method attempts to convert string outputs to JSON dictionaries and handles
+        potential JSON parsing errors gracefully.
+
+        Args:
+            tool_output (str): The output from a tool, which is expected to be a JSON string.
+
+        Returns:
+            dict[str, Any]: The parsed output as a dictionary. Returns an empty
+            dictionary if parsing fails.
+        """
+        if isinstance(tool_output, str):
+            try:
+                return json.loads(tool_output)
+            except json.JSONDecodeError:
+                logger.warning("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]:
+        """Process items from a specific section of search results and create chunks.
+
+        Args:
+            section_name (str): Name of the section being processed (e.g., 'organic', 'news').
+            section_data (list): List of items to process from the section.
+            start_counter (int): Initial counter value for chunk numbering.
+
+        Returns:
+            dict[str, Any]: Dictionary containing:
+                - chunks (list): List of processed chunks extracted from the section items.
+                - counter (int): Updated counter after processing all items.
+
+        Raises:
+            Exception: Logs errors that occur during processing of individual items.
+        """
+        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.warning(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:
+        """Creates a Chunk object from a search result item.
+
+        This method extracts content and metadata from a search result item and creates
+        a Chunk object. It prioritizes snippet over title for content and includes various
+        metadata fields.
+
+        Args:
+            section_name (str): The name of the section this item belongs to.
+            item (dict[str, Any]): The search result item containing information like
+                snippet, title, and link.
+            file_id (int): The ID of the file associated with this chunk.
+
+        Returns:
+            Optional[Chunk]: A Chunk object containing the extracted content and metadata,
+                or None if the item lacks essential information (link or content).
+        """
+        # First check if link exists
+        link = item.get("link")
+        if not link:
+            logger.warning(f"Skipping item {file_id} from {section_name} result: Missing link")
+            return None
+
+        # Then extract and check content
+        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),
+            "url": link,
+            "start_index": 0,
+            "end_index": 0,
+        }
+
+        return Chunk(content=content, metadata=metadata)

aip_agents/tools/web_search/serper_tool.pyi

@@ -0,0 +1,19 @@
+from _typeshed import Incomplete
+from aip_agents.utils.logger import get_logger as get_logger
+from langchain_community.utilities.google_serper import GoogleSerperAPIWrapper
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel
+
+logger: Incomplete
+
+class GoogleSerperInput(BaseModel):
+    """Input schema for the GoogleSerperTool."""
+    query: str
+
+class GoogleSerperTool(BaseTool):
+    """Tool to search Google Serper API."""
+    name: str
+    description: str
+    save_output_history: bool
+    args_schema: type[BaseModel]
+    api_wrapper: GoogleSerperAPIWrapper

aip_agents/types/__init__.py

@@ -0,0 +1,70 @@
+"""Type definitions for aip_agents package.
+
+This module exports type definitions used throughout the aip_agents package
+for better type safety and code clarity.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from typing import Any, Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+from aip_agents.schema.a2a import (
+    A2AEvent,
+    A2AStreamEventType,
+    ToolCallInfo,
+    ToolResultInfo,
+)
+
+
+class ChatMessage(BaseModel):
+    """Represents a single message in a chat conversation."""
+
+    role: str
+    content: str
+
+
+@runtime_checkable
+class AgentProtocol(Protocol):
+    """Defines the expected interface for an agent."""
+
+    id: str
+    name: str
+    description: str
+
+    def run(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Synchronous execution method.
+
+        Args:
+            query (str): The query to execute.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            dict[str, Any]: The execution result.
+        """
+        ...
+
+    async def arun(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Asynchronous execution method.
+
+        Args:
+            query (str): The query to execute.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            dict[str, Any]: The execution result.
+        """
+        ...
+
+
+__all__ = [
+    "A2AEvent",
+    "A2AStreamEventType",
+    "ToolCallInfo",
+    "ToolResultInfo",
+    "AgentProtocol",
+    "ChatMessage",
+]

aip_agents/types/__init__.pyi

@@ -0,0 +1,36 @@
+from aip_agents.schema.a2a import A2AEvent as A2AEvent, A2AStreamEventType as A2AStreamEventType, ToolCallInfo as ToolCallInfo, ToolResultInfo as ToolResultInfo
+from pydantic import BaseModel
+from typing import Any, Protocol
+
+__all__ = ['A2AEvent', 'A2AStreamEventType', 'ToolCallInfo', 'ToolResultInfo', 'AgentProtocol', 'ChatMessage']
+
+class ChatMessage(BaseModel):
+    """Represents a single message in a chat conversation."""
+    role: str
+    content: str
+
+class AgentProtocol(Protocol):
+    """Defines the expected interface for an agent."""
+    id: str
+    name: str
+    description: str
+    def run(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Synchronous execution method.
+
+        Args:
+            query (str): The query to execute.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            dict[str, Any]: The execution result.
+        """
+    async def arun(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Asynchronous execution method.
+
+        Args:
+            query (str): The query to execute.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            dict[str, Any]: The execution result.
+        """

aip_agents/types/a2a_events.py

@@ -0,0 +1,13 @@
+"""Backward-compatible exports for A2A event schema.
+
+Prefer importing from ``aip_agents.schema.a2a`` in new code.
+"""
+
+from aip_agents.schema.a2a import A2AEvent, A2AStreamEventType, ToolCallInfo, ToolResultInfo
+
+__all__ = [
+    "A2AStreamEventType",
+    "A2AEvent",
+    "ToolCallInfo",
+    "ToolResultInfo",
+]

aip_agents/types/a2a_events.pyi

@@ -0,0 +1,3 @@
+from aip_agents.schema.a2a import A2AEvent as A2AEvent, A2AStreamEventType as A2AStreamEventType, ToolCallInfo as ToolCallInfo, ToolResultInfo as ToolResultInfo
+
+__all__ = ['A2AStreamEventType', 'A2AEvent', 'ToolCallInfo', 'ToolResultInfo']

aip_agents/utils/__init__.py

@@ -0,0 +1,79 @@
+# flake8: noqa: F401
+"""AIP Agents Utils.
+
+This module contains utility functions and classes for the AIP Agents package.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from aip_agents.utils.artifact_helpers import (
+    create_artifact_command,
+    create_artifact_response,
+    create_error_response,
+    create_multiple_artifacts_command,
+    create_text_artifact_response,
+)
+from aip_agents.utils.file_prompt_utils import augment_query_with_file_paths
+from aip_agents.utils.final_response_builder import assemble_final_response
+from aip_agents.utils.logger import LoggerManager, get_logger
+from aip_agents.utils.reference_helper import (
+    add_references_chunks,
+    serialize_references_for_metadata,
+    validate_references,
+)
+from aip_agents.utils.step_limit_manager import (
+    _DELEGATION_CHAIN_CVAR,
+    _DELEGATION_DEPTH_CVAR,
+    _REMAINING_STEP_BUDGET_CVAR,
+    _STEP_LIMIT_CONFIG_CVAR,
+    StepExecutionContext,
+    StepLimitManager,
+)
+
+# Optional PII imports - only import if dependencies are available
+try:
+    from aip_agents.utils.pii.pii_handler import ToolPIIHandler
+    from aip_agents.utils.pii.pii_helper import (
+        add_pii_mappings,
+        extract_pii_mapping_from_agent_response,
+    )
+
+    _PII_AVAILABLE = True
+except ImportError:
+    _PII_AVAILABLE = False
+
+# Build __all__ conditionally based on available dependencies
+from aip_agents.utils.sse_chunk_transformer import SSEChunkTransformer
+
+__all__ = [
+    "get_logger",
+    "LoggerManager",
+    "create_artifact_response",
+    "create_error_response",
+    "create_artifact_command",
+    "create_multiple_artifacts_command",
+    "create_text_artifact_response",
+    "validate_references",
+    "serialize_references_for_metadata",
+    "add_references_chunks",
+    "assemble_final_response",
+    "augment_query_with_file_paths",
+    "SSEChunkTransformer",
+    "StepExecutionContext",
+    "StepLimitManager",
+    "_DELEGATION_CHAIN_CVAR",
+    "_DELEGATION_DEPTH_CVAR",
+    "_REMAINING_STEP_BUDGET_CVAR",
+    "_STEP_LIMIT_CONFIG_CVAR",
+]
+
+if _PII_AVAILABLE:
+    __all__.extend(
+        [
+            "ToolPIIHandler",
+            "add_pii_mappings",
+            "extract_pii_mapping_from_agent_response",
+        ]
+    )

aip_agents/utils/__init__.pyi

@@ -0,0 +1,11 @@
+from aip_agents.utils.artifact_helpers import create_artifact_command as create_artifact_command, create_artifact_response as create_artifact_response, create_error_response as create_error_response, create_multiple_artifacts_command as create_multiple_artifacts_command, create_text_artifact_response as create_text_artifact_response
+from aip_agents.utils.file_prompt_utils import augment_query_with_file_paths as augment_query_with_file_paths
+from aip_agents.utils.final_response_builder import assemble_final_response as assemble_final_response
+from aip_agents.utils.logger import LoggerManager as LoggerManager, get_logger as get_logger
+from aip_agents.utils.pii.pii_handler import ToolPIIHandler as ToolPIIHandler
+from aip_agents.utils.pii.pii_helper import add_pii_mappings as add_pii_mappings, extract_pii_mapping_from_agent_response as extract_pii_mapping_from_agent_response
+from aip_agents.utils.reference_helper import add_references_chunks as add_references_chunks, serialize_references_for_metadata as serialize_references_for_metadata, validate_references as validate_references
+from aip_agents.utils.sse_chunk_transformer import SSEChunkTransformer as SSEChunkTransformer
+from aip_agents.utils.step_limit_manager import StepExecutionContext as StepExecutionContext, StepLimitManager as StepLimitManager, _DELEGATION_CHAIN_CVAR as _DELEGATION_CHAIN_CVAR, _DELEGATION_DEPTH_CVAR as _DELEGATION_DEPTH_CVAR, _REMAINING_STEP_BUDGET_CVAR as _REMAINING_STEP_BUDGET_CVAR, _STEP_LIMIT_CONFIG_CVAR as _STEP_LIMIT_CONFIG_CVAR
+
+__all__ = ['get_logger', 'LoggerManager', 'create_artifact_response', 'create_error_response', 'create_artifact_command', 'create_multiple_artifacts_command', 'create_text_artifact_response', 'validate_references', 'serialize_references_for_metadata', 'add_references_chunks', 'assemble_final_response', 'augment_query_with_file_paths', 'SSEChunkTransformer', 'StepExecutionContext', 'StepLimitManager', '_DELEGATION_CHAIN_CVAR', '_DELEGATION_DEPTH_CVAR', '_REMAINING_STEP_BUDGET_CVAR', '_STEP_LIMIT_CONFIG_CVAR', 'ToolPIIHandler', 'add_pii_mappings', 'extract_pii_mapping_from_agent_response']