from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route

from agents import duck_agent_card, text_agent_card, wikipedia_agent_card

AGENT_CALLER_ROLE = "agent.caller"
agent_cards = [wikipedia_agent_card, duck_agent_card, text_agent_card]


async def list_agent_cards(request: Request) -> JSONResponse:
  roles = getattr(request.user, "claims", {}).get("roles", [])
  restricted_names = {duck_agent_card.name}
  visible = [
    c for c in agent_cards
    if c.name not in restricted_names or AGENT_CALLER_ROLE in roles
  ]
  return JSONResponse([card.model_dump(mode="json", exclude_none=True) for card in visible])
app = Starlette(
  routes=[Route("/agents", list_agent_cards)],
  # ...
)

import uuid

from a2a.client import Client as AgentClient
from a2a.client.client_factory import ClientFactory, ClientConfig
from a2a.types import (
  AgentCard,
  Message,
  Part,
  TaskArtifactUpdateEvent,
  TextPart,
)
from httpx import AsyncClient as HttpxClient
from langchain.agents import create_agent
from langchain_core.tools import tool


async def create_client(cards: list[AgentCard], http_client: HttpxClient) -> AgentClient:
  # simplified
  client = await ClientFactory.connect(
    cards[0],  # pick one of your client cards returned from /agent
    client_config=ClientConfig(httpx_client=http_client))
  return client


def make_message(text: str) -> Message:
  return Message(
    role="user",
    message_id=str(uuid.uuid4()),
    parts=[Part(root=TextPart(text=text))],
  )


@tool
async def wiki_agent(query: str, client: HttpxClient) -> str:
  """Search Wikipedia for information."""
  parts = []
  async for _, event in client.send_message(make_message(query)):
    if isinstance(event, TaskArtifactUpdateEvent):
      for part in event.artifact.parts:
        if isinstance(part.root, TextPart):
          parts.append(part.root.text)
  return "\n".join(parts) or "No results."


agent = create_agent(tools=[wiki_agent])

from langchain_core.tools import BaseTool, tool
from pandas import DataFrame

example_frame = DataFrame(
  {
    "row_access": ["admin", "user", "guest", "admin", "user"],
    "value": [100, 200, 300, 400, 500],
  }
)


def make_data_fetch_tool(is_authenticated: bool, user_roles: list[str]) -> BaseTool:
  """Create a data fetch tool with row-level access control."""

  allowed = ["guest"]
  if is_authenticated:
    allowed.append("user")
  if "data_fetch_admin.caller" in user_roles:
    allowed.append("admin")

  @tool
  def fetch_data() -> str:
    """Fetch data from the dataset. Returns rows the current user is allowed to see based on their access level."""
    filtered = example_frame[example_frame["row_access"].isin(allowed)]
    return filtered.to_string(index=False)

  return fetch_data