glean-agent-toolkit 0.2.0__py3-none-any.whl

glean/agent_toolkit/decorators.py

@@ -0,0 +1,271 @@
+"""Decorators for creating tool specifications."""
+
+import functools
+import inspect
+from collections.abc import Callable
+from typing import Any, Protocol, TypedDict, TypeVar, cast
+
+from pydantic import BaseModel
+
+from glean.agent_toolkit.registry import get_registry
+from glean.agent_toolkit.spec import ToolSpec
+
+
+class InputSchema(TypedDict):
+    """JSON Schema for tool input."""
+
+    type: str
+    properties: dict[str, Any]
+    required: list[str]
+
+
+CallableT = Callable[..., Any]
+
+
+class ToolSpecFunction(Protocol):
+    """Protocol for functions decorated with tool_spec."""
+
+    tool_spec: ToolSpec
+
+    def as_openai_tool(self) -> dict[str, Any] | Any:
+        """Convert to OpenAI tool format.
+
+        Returns:
+            OpenAI tool specification
+        """
+        ...
+
+    def as_adk_tool(self) -> Any:
+        """Convert to Google ADK tool format.
+
+        Returns:
+            Google ADK tool
+        """
+        ...
+
+    def as_langchain_tool(self) -> Any:
+        """Convert to LangChain tool format.
+
+        Returns:
+            LangChain tool
+        """
+        ...
+
+    def as_crewai_tool(self) -> Any:
+        """Convert to CrewAI tool format.
+
+        Returns:
+            CrewAI tool
+        """
+        ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the function.
+
+        Args:
+            *args: Positional arguments
+            **kwargs: Keyword arguments
+
+        Returns:
+            Function result
+        """
+        ...
+
+    __name__: str
+
+
+def tool_spec(
+    name: str,
+    description: str,
+    output_model: type[BaseModel] | None = None,
+    version: str | None = None,
+) -> Callable[[CallableT], ToolSpecFunction]:
+    """Decorator for registering a function as a tool.
+
+    Args:
+        name: Name of the tool
+        description: Description of the tool
+        output_model: Optional Pydantic model for the output
+        version: Optional version string
+
+    Returns:
+        Decorated function with tool spec attached
+    """
+
+    def decorator(func: CallableT) -> ToolSpecFunction:
+        """Decorator function.
+
+        Args:
+            func: Function to decorate
+
+        Returns:
+            Decorated function
+        """
+        sig = inspect.signature(func)
+        params = {}
+        out_type = None
+
+        for param_name, param in sig.parameters.items():
+            if param.annotation != inspect.Parameter.empty:
+                params[param_name] = param.annotation
+
+        if sig.return_annotation != inspect.Signature.empty:
+            out_type = sig.return_annotation
+
+        input_schema: InputSchema = {
+            "type": "object",
+            "properties": {},
+            "required": [],
+        }
+
+        required_fields: list[str] = []
+
+        for param_name, param in sig.parameters.items():
+            if param.default is param.empty:
+                required_fields.append(param_name)
+
+        input_schema["required"] = required_fields
+
+        if params:
+            for param_name, param_type in params.items():
+                if isinstance(param_type, type) and issubclass(param_type, str):
+                    input_schema["properties"][param_name] = {"type": "string"}
+                elif isinstance(param_type, type) and issubclass(param_type, int):
+                    input_schema["properties"][param_name] = {"type": "integer"}
+                elif isinstance(param_type, type) and issubclass(param_type, float):
+                    input_schema["properties"][param_name] = {"type": "number"}
+                elif isinstance(param_type, type) and issubclass(param_type, bool):
+                    input_schema["properties"][param_name] = {"type": "boolean"}
+                elif param_type is list or param_type is list[str]:
+                    input_schema["properties"][param_name] = {
+                        "type": "array",
+                        "items": {"type": "string"},
+                    }
+                elif param_type is list[int]:
+                    input_schema["properties"][param_name] = {
+                        "type": "array",
+                        "items": {"type": "integer"},
+                    }
+                else:
+                    input_schema["properties"][param_name] = {"type": "string"}
+
+        output_schema: dict[str, Any] = {"type": "object"}
+        if out_type is not None and hasattr(out_type, "model_json_schema"):
+            output_schema = out_type.model_json_schema()
+        elif out_type is int:
+            output_schema = {"type": "integer"}
+        elif out_type is float:
+            output_schema = {"type": "number"}
+        elif out_type is bool:
+            output_schema = {"type": "boolean"}
+        elif out_type is str:
+            output_schema = {"type": "string"}
+        elif out_type is list or out_type is list[str]:
+            output_schema = {
+                "type": "array",
+                "items": {"type": "string"},
+            }
+        elif out_type is list[int]:
+            output_schema = {
+                "type": "array",
+                "items": {"type": "integer"},
+            }
+
+        tool_spec_obj = ToolSpec(
+            name=name,
+            description=description,
+            function=func,
+            input_schema=cast(dict[str, Any], input_schema),
+            output_schema=output_schema,
+            version=version,
+            output_model=(
+                output_model
+                if isinstance(output_model, type) and issubclass(output_model, BaseModel)
+                else None
+            ),
+        )
+
+        get_registry().register(tool_spec_obj)
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            """Wrapper that preserves the original function's call semantics.
+
+            Args:
+                *args: Positional arguments
+                **kwargs: Keyword arguments
+
+            Returns:
+                The result of the function call
+            """
+            return func(*args, **kwargs)
+
+        def as_openai_tool() -> dict[str, Any] | Any:
+            """Convert to OpenAI tool format.
+
+            Returns:
+                OpenAI tool specification
+            """
+            from glean.agent_toolkit.adapters.openai import OpenAIAdapter
+
+            adapter = tool_spec_obj.get_adapter("openai")
+            if adapter is None:
+                adapter = OpenAIAdapter(tool_spec_obj)
+                tool_spec_obj.set_adapter("openai", adapter)
+
+            return adapter.to_tool()
+
+        def as_adk_tool() -> Any:
+            """Convert to Google ADK tool format.
+
+            Returns:
+                Google ADK tool
+            """
+            from glean.agent_toolkit.adapters.adk import ADKAdapter
+
+            adapter = tool_spec_obj.get_adapter("adk")
+            if adapter is None:
+                adapter = ADKAdapter(tool_spec_obj)
+                tool_spec_obj.set_adapter("adk", adapter)
+
+            return adapter.to_tool()
+
+        def as_langchain_tool() -> Any:
+            """Convert to LangChain tool format.
+
+            Returns:
+                LangChain tool
+            """
+            from glean.agent_toolkit.adapters.langchain import LangChainAdapter
+
+            adapter = tool_spec_obj.get_adapter("langchain")
+            if adapter is None:
+                adapter = LangChainAdapter(tool_spec_obj)
+                tool_spec_obj.set_adapter("langchain", adapter)
+
+            return adapter.to_tool()
+
+        def as_crewai_tool() -> Any:
+            """Convert to CrewAI tool format.
+
+            Returns:
+                CrewAI tool
+            """
+            from glean.agent_toolkit.adapters.crewai import CrewAIAdapter
+
+            adapter = tool_spec_obj.get_adapter("crewai")
+            if adapter is None:
+                adapter = CrewAIAdapter(tool_spec_obj)
+                tool_spec_obj.set_adapter("crewai", adapter)
+
+            return adapter.to_tool()
+
+        wrapper.as_openai_tool = as_openai_tool  # type: ignore
+        wrapper.as_adk_tool = as_adk_tool  # type: ignore
+        wrapper.as_langchain_tool = as_langchain_tool  # type: ignore
+        wrapper.as_crewai_tool = as_crewai_tool  # type: ignore
+        wrapper.tool_spec = tool_spec_obj  # type: ignore
+
+        return cast(ToolSpecFunction, wrapper)
+
+    return decorator

glean/agent_toolkit/py.typed

@@ -0,0 +1 @@
+ 

glean/agent_toolkit/registry.py

@@ -0,0 +1,50 @@
+"""Tool registry for managing and retrieving tool specifications."""
+
+from glean.agent_toolkit.spec import ToolSpec
+
+
+class Registry:
+    """Registry for tool specifications."""
+
+    def __init__(self) -> None:
+        """Initialize the registry."""
+        self._tools: dict[str, ToolSpec] = {}
+
+    def register(self, tool_spec: ToolSpec) -> None:
+        """Register a tool specification.
+
+        Args:
+            tool_spec: The tool specification to register
+        """
+        self._tools[tool_spec.name] = tool_spec
+
+    def get(self, name: str) -> ToolSpec | None:
+        """Get a tool specification by name.
+
+        Args:
+            name: The name of the tool
+
+        Returns:
+            The tool specification, or None if not found
+        """
+        return self._tools.get(name)
+
+    def list(self) -> list[ToolSpec]:
+        """List all registered tool specifications.
+
+        Returns:
+            List of all registered tool specifications
+        """
+        return list(self._tools.values())
+
+
+_REGISTRY = Registry()
+
+
+def get_registry() -> Registry:
+    """Get the global registry instance.
+
+    Returns:
+        The global registry instance
+    """
+    return _REGISTRY

glean/agent_toolkit/spec.py

@@ -0,0 +1,51 @@
+"""Tool specification dataclass."""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic import BaseModel
+
+
+@dataclass
+class ToolSpec:
+    """Specification for a tool.
+
+    Attributes:
+        name: The name of the tool
+        description: A description of what the tool does
+        function: The function that implements the tool
+        input_schema: JSON schema for the input parameters
+        output_schema: JSON schema for the output value
+        version: Optional version string
+        output_model: Optional pydantic model for the output
+    """
+
+    name: str
+    description: str
+    function: Callable
+    input_schema: dict[str, Any]
+    output_schema: dict[str, Any]
+    version: str | None = None
+    output_model: type[BaseModel] | None = None
+    _adapters: dict[str, Any] = field(default_factory=dict)
+
+    def get_adapter(self, name: str) -> Any | None:
+        """Get a cached adapter instance.
+
+        Args:
+            name: The name of the adapter
+
+        Returns:
+            The adapter instance
+        """
+        return self._adapters.get(name)
+
+    def set_adapter(self, name: str, adapter: Any) -> None:
+        """Set an adapter instance.
+
+        Args:
+            name: The name of the adapter
+            adapter: The adapter instance
+        """
+        self._adapters[name] = adapter

glean/agent_toolkit/tools/__init__.py

@@ -0,0 +1,43 @@
+"""
+Each tool lives in its own module under :pymod:`glean.agent_toolkit.tools`.
+
+Importing this package will load all available tools.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module as _import_module
+
+_tool_modules: list[str] = [
+    "glean_search",
+    "web_search",
+    "ai_web_search",
+    "calendar_search",
+    "employee_search",
+    "code_search",
+    "gmail_search",
+    "outlook_search",
+]
+
+for _mod in _tool_modules:
+    _import_module(f"{__name__}.{_mod}")
+
+from .ai_web_search import ai_web_search  # noqa: E402
+from .calendar_search import calendar_search  # noqa: E402
+from .code_search import code_search  # noqa: E402
+from .employee_search import employee_search  # noqa: E402
+from .glean_search import glean_search  # noqa: E402
+from .gmail_search import gmail_search  # noqa: E402
+from .outlook_search import outlook_search  # noqa: E402
+from .web_search import web_search  # noqa: E402
+
+__all__: list[str] = [
+    "glean_search",
+    "web_search",
+    "ai_web_search",
+    "calendar_search",
+    "employee_search",
+    "code_search",
+    "gmail_search",
+    "outlook_search",
+]

glean/agent_toolkit/tools/_common.py

@@ -0,0 +1,36 @@
+"""Shared helpers for built-in stub tools."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from glean.api_client import Glean, models
+
+
+def api_client() -> Glean:
+    """Get the Glean API client."""
+    instance = os.getenv("GLEAN_INSTANCE")
+    api_token = os.getenv("GLEAN_API_TOKEN")
+
+    if not api_token or not instance:
+        raise ValueError("GLEAN_API_TOKEN and GLEAN_INSTANCE environment variables are required")
+
+    return Glean(api_token=api_token, instance=instance)
+
+
+def run_tool(
+    tool_display_name: str,
+    parameters: dict[str, models.ToolsCallParameter],
+) -> dict[str, Any]:
+    """Execute a Glean stub tool and wrap the response."""
+    try:
+        with api_client() as g_client:
+            result = g_client.client.tools.run(
+                name=tool_display_name,
+                parameters=parameters,
+            )
+
+            return {"result": result}
+    except Exception as exc:
+        return {"error": str(exc), "result": None}

glean/agent_toolkit/tools/ai_web_search.py

@@ -0,0 +1,40 @@
+"""AI Web Search tool for searching the internet with AI-powered results."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="ai_web_search",
+    description=(
+        "Searches the web for up-to-date external information and generates a well structured "
+        "response to the user query grounded in relevant webpages. Closely evaluate the "
+        "instructions below for the user query to decide whether to use the web agent. If you "
+        "think the scenarios are contradictory, do not use the web agent unless there is clear "
+        "user intent. If you think the user query is ambiguous, use the Conversation history to "
+        "infer the context.\n"
+        "INSTRUCTIONS:\n"
+        "Examples of when to use this tool:\n"
+        "- User Intent: Use this tool if the user is asking you to search the web, look online, "
+        "provide links/sources, or explicitly looking for current external information (outside "
+        "of the company) like weather, news, or latest updates/plans.\n"
+        "- Freshness: Use this tool if you need up-to-date information on time-dependent topics "
+        "or any time you would otherwise refuse to answer a question because your knowledge might "
+        "be out of date.\n"
+        "- Niche Information: If the answer would likely change based on detailed information not "
+        "widely known or understood (which might be found on the internet), use web sources "
+        "directly rather than relying on the distilled knowledge from pretraining.\n"
+        "- Do NOT use this tool for programming related queries. You already know enough about "
+        "those.\n"
+        "- Do NOT use this tool for queries seeking general ideas, which may not benefit from "
+        "specific information."
+    ),
+)
+def ai_web_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search the web for up-to-date external information."""
+    return run_tool("Gemini Web Search", parameters)

glean/agent_toolkit/tools/calendar_search.py

@@ -0,0 +1,18 @@
+"""Calendar Search tool for searching calendar events."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="calendar_search",
+    description="Searches over all the calendar meetings of the company.",
+)
+def calendar_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search the calendar for meetings."""
+    return run_tool("Meeting Lookup", parameters)

glean/agent_toolkit/tools/code_search.py

@@ -0,0 +1,29 @@
+"""Code Search tool for searching code repositories."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="code_search",
+    description=(
+        "Searches over all code changes made in the company.\n"
+        "INSTRUCTIONS:\n"
+        "- Use this tool to help users find information in or about code, add new code, etc. "
+        "Prefer including code snippets in your response.\n"
+        "- This is your primary tool to access knowledge present in the company's code "
+        "repositories.\n"
+        "- The results returned are not exhaustive; we only return the top few most relevant "
+        "results to a query."
+    ),
+)
+def code_search(
+    parameters: dict[str, models.ToolsCallParameter],
+) -> dict[str, Any]:
+    """Search code repositories based on the query."""
+    return run_tool("Code Search", parameters)

glean/agent_toolkit/tools/employee_search.py

@@ -0,0 +1,33 @@
+"""Employee Search tool for finding people profiles in the company."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="employee_search",
+    description=(
+        "Finds people at the company based on their personal information.\n"
+        "INSTRUCTIONS:\n"
+        "- Only use this when the user explicitly wants to find people in the company (e.g.,"
+        ' "who" questions) or for aggregation queries on people.\n'
+        "- You can also use this tool to find personal information about employees (e.g., what is "
+        "person X's phone number or email address).\n"
+        "- Do not use this when the user wants to find people outside of the company, or people "
+        "who are no longer at the company.\n"
+        "- You can find people based on details such as name, email, title, department, and "
+        "location.\n"
+        "- The results returned are not exhaustive; we only return the top few most relevant "
+        "people to a query.\n"
+        '- For analytics questions such as "how many people..." use the "statistics" field '
+        "in the output."
+    ),
+)
+def employee_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search for employees based on the query."""
+    return run_tool("Employee Search", parameters)

glean/agent_toolkit/tools/glean_search.py

@@ -0,0 +1,26 @@
+"""Glean Search tool for searching company documents and data."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="glean_search",
+    description=(
+        "Finds relevant documents in the company.\n"
+        "INSTRUCTIONS:\n"
+        "- This is your primary tool to access all knowledge within the company.\n"
+        "- The results returned are not exhaustive; we only return the top few most relevant "
+        "documents to a query.\n"
+        '- For analytics questions such as "how many documents..." use the "statistics" '
+        "field in the output."
+    ),
+)
+def glean_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search Glean for relevant documents using the query."""
+    return run_tool("Glean Search", parameters)

glean/agent_toolkit/tools/gmail_search.py

@@ -0,0 +1,23 @@
+"""Gmail Search tool for searching email messages."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="gmail_search",
+    description=(
+        "Finds relevant emails in the user's mailbox.\n"
+        "- Only use this tool if the user asks for email.\n"
+        "- Results returned are not exhaustive; we can only return the top 10 emails sorted by "
+        "recency (most recent first)."
+    ),
+)
+def gmail_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search Gmail messages based on the query."""
+    return run_tool("Gmail Search", parameters)

glean/agent_toolkit/tools/outlook_search.py

@@ -0,0 +1,23 @@
+"""Outlook Search tool for searching email messages."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from glean.agent_toolkit.decorators import tool_spec
+from glean.agent_toolkit.tools._common import run_tool
+from glean.api_client import models
+
+
+@tool_spec(
+    name="outlook_search",
+    description=(
+        "Finds relevant emails in the user's mailbox.\n"
+        "- Only use this tool if the user asks for email.\n"
+        "- Results returned are not exhaustive; we can only return the top 10 emails sorted by "
+        "recency (most recent first)."
+    ),
+)
+def outlook_search(parameters: dict[str, models.ToolsCallParameter]) -> dict[str, Any]:
+    """Search Outlook messages based on the query."""
+    return run_tool("Outlook Search", parameters)