How to Give a Pydantic AI Agent B2B Data (REST Tools and MCP)

pip install "pydantic-ai-slim[anthropic]" httpx
export LEADOCEAN_API_KEY="..."   # from app.leadocean.io
export ANTHROPIC_API_KEY="..."

import asyncio
import os
from dataclasses import dataclass

import httpx
from pydantic_ai import ModelRetry, RunContext, ToolFailed

BASE_URL = "https://api.leadocean.io"


@dataclass
class Deps:
    http: httpx.AsyncClient
    budget: int = 25  # records this run may spend
    spent: int = 0


def make_client() -> httpx.AsyncClient:
    return httpx.AsyncClient(
        base_url=BASE_URL,
        headers={"x-api-key": os.environ["LEADOCEAN_API_KEY"]},
        timeout=30,
    )


async def call(ctx: RunContext[Deps], method: str, path: str, **kw) -> dict:
    """One LeadOcean request with the error handling the docs describe."""
    for attempt in range(3):
        r = await ctx.deps.http.request(method, path, **kw)
        wait = int(r.headers.get("Retry-After", "1"))
        if r.status_code == 429 and wait <= 5:  # above the plan's per-second limit
            await asyncio.sleep(wait)
            continue
        if r.status_code == 503:  # transient, never billed
            await asyncio.sleep(2**attempt)
            continue
        break

    if r.status_code == 200:
        body = r.json()
        ctx.deps.spent += body.get("meta", {}).get("credits", 0)
        return body
    if r.status_code == 404:  # nobody matched, costs nothing
        return {"found": False}

    message = r.json().get("error", {}).get("message", r.text)
    if r.status_code == 400:  # the model can fix its own arguments
        raise ModelRetry(f"LeadOcean rejected the request: {message}. "
                         "Check enum values with list_values.")
    if r.status_code == 402:
        raise ToolFailed("Out of records on the Free plan. Stop and tell the user.")
    if r.status_code == 429:  # still limited after the short waits
        raise ToolFailed(f"Rate limited, retry after {wait} seconds. Do not retry now.")
    raise ToolFailed(f"LeadOcean returned {r.status_code}: {message}")

import asyncio

from pydantic_ai import Agent, RunContext, ToolFailed
from pydantic_ai.usage import UsageLimits

from tools import Deps, call, make_client

agent = Agent(
    "anthropic:claude-sonnet-4-6",
    deps_type=Deps,
    instructions=(
        "You build small B2B prospect lists with LeadOcean. Always call "
        "count_people first. Search only when the count is under 100,000. "
        "Enrich at most 3 people, and only when asked for contact details. "
        "Look up enum values with list_values instead of guessing."
    ),
)


@agent.tool
async def list_values(ctx: RunContext[Deps], name: str, q: str = "") -> list[str]:
    """Free. Exact values for a filter: job_level, job_function, country, seniority, industry."""
    params = {"limit": 50, **({"q": q} if q else {})}
    body = await call(ctx, "GET", f"/v1/enums/{name}", params=params)
    return body["data"]["values"]


@agent.tool
async def count_people(ctx: RunContext[Deps], filters: dict[str, list[str]]) -> int:
    """Free. Size a segment. Filter values are arrays, e.g. {"jobLevel": ["VP"], "country": ["CA"]}."""
    body = await call(ctx, "POST", "/v1/people/search",
                      params={"count": "true"}, json={**filters, "limit": 1})
    return body["meta"]["total"]


def trim(row: dict) -> dict:
    """Keep six fields. The model does not need the whole record."""
    contact = row.get("contact_data", {})
    job = (contact.get("contact_current_experiences") or [{}])[0]
    return {
        "person_id": row.get("meta", {}).get("source_ids", {}).get("person_id"),
        "name": row.get("profile_data", {}).get("profile_full_name"),
        "title": job.get("job_title"),
        "company": job.get("company_name"),
        "domain": job.get("company_domain"),
        "email_status": contact.get("email_status"),
    }


@agent.tool
async def search_people(ctx: RunContext[Deps], filters: dict[str, list[str]], limit: int = 5) -> list[dict]:
    """Costs 1 record per person returned. Returns thin rows: name, title, company, person_id."""
    left = ctx.deps.budget - ctx.deps.spent
    if left <= 0:
        raise ToolFailed("Record budget for this run is spent. Stop.")
    body = await call(ctx, "POST", "/v1/people/search",
                      json={**filters, "limit": min(limit, left, 25)})
    return [trim(row) for row in body["data"]]


@agent.tool
async def enrich_person(ctx: RunContext[Deps], person_id: str) -> dict:
    """Costs 1 record. Returns the profile and revealed emails for one person_id."""
    if ctx.deps.budget - ctx.deps.spent < 1:
        raise ToolFailed("Record budget for this run is spent. Stop.")
    body = await call(ctx, "POST", "/v1/people/enrich",
                      json={"person_id": person_id, "reveal_email": True})
    return body.get("data", body)


async def main() -> None:
    async with make_client() as http:
        deps = Deps(http=http, budget=25)
        result = await agent.run(
            "Find 5 VPs of Sales in Canada and tell me how many exist in total.",
            deps=deps,
            usage_limits=UsageLimits(request_limit=12),
        )
        print(result.output)
        print(f"LeadOcean records spent: {deps.spent}")


if __name__ == "__main__":
    asyncio.run(main())

python agent.py

import asyncio

from fastmcp import Client
from pydantic_ai import Agent
from pydantic_ai.mcp import MCPToolset

leadocean = MCPToolset(Client("https://api.leadocean.io/mcp", auth="oauth"))
agent = Agent("anthropic:claude-sonnet-4-6", toolsets=[leadocean])


async def main() -> None:
    result = await agent.run("Count VPs of Sales in Canada. Do not search.")
    print(result.output)


asyncio.run(main())

{"name":"job_level","total":6,"values":["C-Team","Director","Manager","Other","Staff","VP"]}

{
  "success": true,
  "data": [
    {
      "profile_data": { "profile_full_name": "Jane Doe" },
      "contact_data": {
        "has_email": true,
        "email_status": "verified",
        "contact_current_experiences": [
          { "job_title": "VP Sales", "company_name": "Acme", "company_domain": "acme.com" }
        ]
      },
      "meta": { "source_ids": { "person_id": "12345678" } }
    }
  ],
  "meta": { "count": 1, "limit": 5, "credits": 1, "nextCursor": "..." }
}