Developing Tool Plugins¶

from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from pydantic import BaseModel, Field


class MyToolInput(BaseModel):
    query: str = Field(description="The search query")
    limit: int = Field(default=10, description="Max results to return")


@tool("my_tool_name", args_schema=MyToolInput)
async def my_tool(query: str, limit: int, config: RunnableConfig) -> str:
    """
    Short description of what this tool does.

    The LLM reads this docstring to decide when to use the tool.
    Be specific and concise.
    """
    # Access user/chat context from the config
    user_id = config.get("configurable", {}).get("logged_user_id")
    group_id = config.get("configurable", {}).get("group_id")
    chat_id = config.get("configurable", {}).get("chat_id")

    # Your implementation
    result = f"Found results for '{query}' (limit={limit})"
    return result

config.get("configurable", {}).get("logged_user_id")   # int — current user ID
config.get("configurable", {}).get("group_id")          # int — current group ID
config.get("configurable", {}).get("chat_id")           # int — current chat ID
config.get("configurable", {}).get("chat_uuid")         # str — chat UUID

# In src/virtual_assistant_lc/tool_registrations.py

from src.virtual_assistant_lc.tools.my_module.my_tool import my_tool

register_tools("my_capability_code", [my_tool],
    name="My Tool Group",
    description="Short human-friendly description of this capability"
)

# In src/virtual_assistant_lc/tools/my_module/my_tool.py

from src.virtual_assistant_lc.tool_registry import capability
from langchain_core.tools import tool


@capability("my_capability_code")
@tool("my_tool_name")
async def my_tool(config: RunnableConfig) -> str:
    """Tool description."""
    return "result"

# At the bottom of your toolkit file
my_toolkit = [tool_a, tool_b, tool_c]

from src.virtual_assistant_lc.tools.my_module.my_toolkit import my_toolkit

register_tools("my_capability_code", my_toolkit,
    name="My Toolkit",
    description="Group of related tools for X"
)

from src.virtual_assistant_lc.tool_registry import capability

my_toolkit = capability("my_capability_code")([tool_a, tool_b, tool_c])

# Via decorator — the tool is available when EITHER code is active
@capability("base", "rag")
@tool("get_document_text")
async def get_document_text(...): ...

# Via register_tools — register the same tool under multiple codes
register_tools("base", [shared_tool], name="Base Tools", description="Core agent tools")
register_tools("rag", [shared_tool], name="RAG", description="Retrieval-augmented generation")

# src/assist_manage_toolkit/manage_toolkit.py

from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from pydantic import BaseModel, Field


class SearchMattersInput(BaseModel):
    query: str = Field(description="Search text for matters")
    limit: int = Field(default=20, description="Maximum results")


@tool("search_matters", args_schema=SearchMattersInput)
async def search_matters(query: str, limit: int, config: RunnableConfig) -> str:
    """
    Search for matters/cases by text query in the management system.
    Returns matching matters with their reference numbers and titles.
    """
    # Your implementation here
    return "results..."


class CreateTimerInput(BaseModel):
    matter_id: int = Field(description="The matter ID to log time against")
    description: str = Field(description="Description of the work performed")


@tool("create_timer", args_schema=CreateTimerInput)
async def create_timer(matter_id: int, description: str, config: RunnableConfig) -> str:
    """
    Create a new active timer for a matter in the management system.
    The timer starts immediately and can be stopped later.
    """
    # Your implementation here
    return "Timer created"


# Export as a toolkit list
manage_toolkit = [
    search_matters,
    create_timer,
]

# src/assist_manage_toolkit/__init__.py


def register(registry_fn, context=None):
    """Called by PrimeThink API at startup via entry-point discovery.

    Args:
        registry_fn: The register_tools(code, tools_list, *, name, description)
                     function from the main API's tool_registry module.
        context: Optional object providing PT services (settings, document
                 storage, etc.). See "Plugin Context" section below.
    """
    from .manage_toolkit import manage_toolkit

    registry_fn("obviously_manage", manage_toolkit,
        name="Manage",
        description="Search and manage matters, timers, and billing"
    )

def register(registry_fn, context=None):
    from .manage_toolkit import manage_toolkit
    from .search_tools import search_table_tool

    registry_fn("obviously_manage", manage_toolkit,
        name="Manage", description="Full management toolkit")
    registry_fn("manage_search", [search_table_tool],
        name="Manage Search", description="Search tables in the management system")

# src/my_plugin/__init__.py

def register(registry_fn, context=None):
    from .my_toolkit import core_toolkit
    from .settings import set_settings_provider, PrimethinkSettingsProvider

    # Use context for per-user settings resolution
    if context is not None and hasattr(context, "get_user_setting_value"):
        set_settings_provider(PrimethinkSettingsProvider(context.get_user_setting_value))

    # Always register core tools
    registry_fn("my_code", core_toolkit,
        name="My Plugin", description="Core plugin tools")

    # Conditionally register tools that need PT services
    if context is not None and hasattr(context, "document_service") and context.document_service is not None:
        from .pt_tools import pt_toolkit, set_document_service
        set_document_service(context.document_service)
        registry_fn("my_code", pt_toolkit)

from typing import Protocol, Any, Tuple, Optional

class DocumentServiceProtocol(Protocol):
    async def get_document_by_id(self, document_id: int) -> Optional[Any]: ...
    async def fetch_document_bytes(self, document: Any) -> Tuple[str, bytes]: ...
    def process_content_by_format(self, content: str, format: Any) -> Tuple[str, str, bytes]: ...

# settings.py in your plugin

class ToolkitSettings:
    manage_base_url: str
    manage_token: str

class SettingsProvider(Protocol):
    async def get_settings(self, config: dict) -> ToolkitSettings: ...

class PrimethinkSettingsProvider:
    """Reads settings from the PT DB via the injected context."""
    def __init__(self, get_user_setting_value):
        self._get_setting = get_user_setting_value

    async def get_settings(self, config):
        configurable = config.get("configurable", {})
        user_id = configurable.get("logged_user_id")
        group_id = configurable.get("group_id")
        base_url = await self._get_setting("manage_base_url", user_id, group_id)
        token = await self._get_setting("manage_token", user_id, group_id)
        return ToolkitSettings(manage_base_url=base_url, manage_token=token)

class StaticSettingsProvider:
    """Returns fixed settings — useful for MCP servers or testing."""
    def __init__(self, base_url, token):
        self._settings = ToolkitSettings(manage_base_url=base_url, manage_token=token)

    async def get_settings(self, config):
        return self._settings

def register(registry_fn):
    from .my_toolkit import my_toolkit
    registry_fn("my_code", my_toolkit,
        name="My Plugin", description="What this plugin does")

def register(registry_fn):
    """Entry point called by PrimeThink API at startup."""
    from .crm_toolkit import crm_toolkit

    registry_fn("crm", crm_toolkit,
        name="CRM Integration",
        description="Search and manage contacts in the CRM system"
    )

import httpx
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from pydantic import BaseModel, Field


CRM_BASE_URL = "https://api.example-crm.com/v2"


class SearchContactsInput(BaseModel):
    query: str = Field(description="Name, email, or phone number to search for")
    limit: int = Field(default=10, description="Maximum number of results")


@tool("crm_search_contacts", args_schema=SearchContactsInput)
async def crm_search_contacts(query: str, limit: int, config: RunnableConfig) -> str:
    """
    Search for contacts in the CRM by name, email, or phone number.
    Returns a list of matching contacts with their basic details.
    """
    user_id = config.get("configurable", {}).get("logged_user_id")

    async with httpx.AsyncClient(timeout=15.0) as client:
        resp = await client.get(
            f"{CRM_BASE_URL}/contacts/search",
            params={"q": query, "limit": limit},
            headers={"Authorization": f"Bearer {_get_api_key()}"},
        )
        resp.raise_for_status()
        return resp.text


class CreateContactInput(BaseModel):
    first_name: str = Field(description="Contact's first name")
    last_name: str = Field(description="Contact's last name")
    email: str = Field(description="Contact's email address")
    company: str = Field(default="", description="Company name (optional)")


@tool("crm_create_contact", args_schema=CreateContactInput)
async def crm_create_contact(
    first_name: str, last_name: str, email: str, company: str,
    config: RunnableConfig,
) -> str:
    """
    Create a new contact in the CRM system.
    Returns the created contact's ID and a confirmation message.
    """
    async with httpx.AsyncClient(timeout=15.0) as client:
        resp = await client.post(
            f"{CRM_BASE_URL}/contacts",
            json={
                "first_name": first_name,
                "last_name": last_name,
                "email": email,
                "company": company,
            },
            headers={"Authorization": f"Bearer {_get_api_key()}"},
        )
        resp.raise_for_status()
        return resp.text


def _get_api_key() -> str:
    """Read API key from environment. In production, prefer user settings."""
    import os
    return os.environ.get("CRM_API_KEY", "")


# Export as a toolkit list
crm_toolkit = [
    crm_search_contacts,
    crm_create_contact,
]

"""Basic import and registration test."""


def test_register_function_exists():
    from primethink_crm_tools import register
    assert callable(register)


def test_register_calls_registry():
    """Verify register() calls registry_fn with the expected code and tools."""
    calls = []

    def mock_registry_fn(code, tools, **kwargs):
        calls.append((code, tools, kwargs))

    from primethink_crm_tools import register
    register(mock_registry_fn)

    assert len(calls) == 1
    code, tools, kwargs = calls[0]
    assert code == "crm"
    assert len(tools) == 2
    assert tools[0].name == "crm_search_contacts"
    assert tools[1].name == "crm_create_contact"
    assert kwargs["name"] == "CRM Integration"

def test_plugin_registration():
    calls = []

    def mock_registry(code, tools, **kwargs):
        calls.append((code, tools, kwargs))

    from assist_manage_toolkit import register
    register(mock_registry)

    assert len(calls) >= 1
    for code, tools, kwargs in calls:
        print(f"Code: {code}, Tools: {[t.name for t in tools]}, Meta: {kwargs}")
        assert isinstance(code, str)
        assert isinstance(tools, list)
        assert len(tools) > 0
        # Verify metadata is provided
        assert kwargs.get("name"), f"Missing name for code '{code}'"

import asyncio
from unittest.mock import MagicMock


async def test_my_tool():
    from assist_manage_toolkit.manage_toolkit import search_matters

    config = {"configurable": {"logged_user_id": 1, "group_id": 1, "chat_id": 1}}
    result = await search_matters.ainvoke({"query": "test", "limit": 5}, config=config)
    print(result)

asyncio.run(test_my_tool())