Evaluating your MCP and A2A agents with W&B Weave

uv run . --host 0.0.0.0 --port 1001

import asyncio
import os
import uuid
import logging
import base64
from typing import Optional, List, Dict

import httpx
from litellm import completion

from mcp.server.fastmcp import FastMCP
import weave

weave.init("wv_mcp")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

AGENT_CARD_PATH = "/.well-known/agent.json"
HOST = "localhost"
START_PORT = 1000
END_PORT = 1010
MODEL_NAME = "gemini/gemini-2.0-flash"



# --- export/set API key for downstream libraries (if not already set elsewhere) ---
os.environ["GEMINI_API_KEY"] = "your_gemini_api_key"

def create_routing_prompt(history: List[Dict], agents: List[Dict]) -> List[Dict]:
    agent_list = "\n".join(
        f"- {a['card'].get('name', 'Unknown')}: {a['card'].get('description', '')}" for a in agents
    )
    convo = []
    for msg in history:
        role = "User" if msg["role"] == "user" else "Agent"
        convo.append(f"{role}: {msg['text']}")
    convo_str = "\n".join(convo)
    content = (
        "You are an expert router for user requests. "
        "Given the following agents:\n"
        f"{agent_list}\n\n"
        "Below is the full conversation so far. "
        "Pay closest attention to the user's most recent request at the end. "
        "Choose the single best matching agent by exact name to handle the user's request, IN CONTEXT of the conversation.\n\n"
        f"Conversation:\n{convo_str}\n\nAgent name:"
    )
    return [{"role": "user", "content": content}]

def make_agent_instruction(history: List[Dict], agent_card: Dict) -> str:
    convo = "\n".join(
        f"User: {msg['text']}" if msg["role"] == "user" else f"Agent: {msg['text']}"
        for msg in history
    )
    agent_desc = agent_card.get('description', 'No agent description provided.')
    agent_name = agent_card.get('name', '[Agent name missing]')
    return (
        f"You are the agent: {agent_name}.\n"
        f"Your job: {agent_desc}\n\n"
        "Conversation so far:\n"
        f"{convo}\n\n"
        "Based on this conversation, take the user's intended action or answer their last request, using all necessary info from the dialog.\n"
        "Return only the task/instruction that the agent will accept ."
    )

async def fetch_agent_card(host: str, port: int) -> Optional[Dict]:
    url = f"http://{host}:{port}{AGENT_CARD_PATH}"
    try:
        async with httpx.AsyncClient(timeout=3) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            card = resp.json()
            logger.info(f"Agent found: {card.get('name', '[no-name]')} on {host}:{port}")
            return {"host": host, "port": port, "card": card}
    except Exception as e:
        logger.debug(f"No agent at {host}:{port} ({e})")
        return None


async def discover_agents(host: str, start_port: int, end_port: int) -> List[Dict]:
    tasks = [fetch_agent_card(host, port) for port in range(start_port, end_port + 1)]
    results = await asyncio.gather(*tasks)
    return [r for r in results if r]


async def send_instruction_to_agent(agent: Dict, instruction: str) -> Optional[Dict]:
    url = f"http://{agent['host']}:{agent['port']}/"
    jsonrpc_id = str(uuid.uuid4())
    payload = {
        "jsonrpc": "2.0",
        "id": jsonrpc_id,
        "method": "tasks/send",
        "params": {
            "id": jsonrpc_id,
            "message": {
                "role": "user",
                "parts": [
                    {"type": "text", "text": instruction}
                ]
            },
            "acceptedOutputModes": ["text/plain",],
        }
    }
    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()
            return data
    except Exception as e:
        logger.error(f"Error sending query to agent: {e}")
        return None

def extract_agent_text(agent_response: Dict) -> str:
    """Extracts text from agent response artifacts."""
    result = agent_response.get("result", {})
    artifacts = result.get("artifacts", [])
    texts = []
    for artifact in artifacts:
        for part in artifact.get("parts", []):
            if part.get("type") == "text":
                texts.append(part.get("text", ""))
    return "\n".join(texts)

# # ------ MCP TOOL WRAPPER ------


def extract_agent_content(agent_response: Dict) -> Dict:
    """
    Extracts texts and images from agent response artifacts.

    Returns a dict:
        {
            "texts": [ ... ],
            "images": [ { "filename": ..., "bytes": ..., "type": ... }, ... ]
        }
    Images are decoded but also stored as bytes (and filenames if saved locally).
    """
    result = agent_response.get("result", {})
    artifacts = result.get("artifacts", [])
    texts = []
    images = []

    for artifact in artifacts:
        for part in artifact.get("parts", []):
            ptype = part.get("type")
            if ptype == "text":
                texts.append(part.get("text", ""))
            elif ptype in ("image/png", "image/jpeg", "image/jpg"):
                b64data = part.get("text") or part.get("bytes") or ""
                if not b64data:
                    logger.warning("Image artifact part missing data")
                    continue
                try:
                    image_bytes = base64.b64decode(b64data)
                    ext = "png" if "png" in ptype else "jpg"
                    filename = f"agent_image_{uuid.uuid4().hex}.{ext}"
                    # Optional: save to file
                    with open(filename, "wb") as f:
                        f.write(image_bytes)
                    images.append({"filename": filename, "bytes": image_bytes, "type": ptype})
                    texts.append(filename)
                except Exception as e:
                    logger.error(f"Failed to decode/save image: {e}")
            elif ptype == "file":
                file_info = part.get("file", {})
                b64data = file_info.get("bytes")
                mime = file_info.get("mimeType", "")
                if not b64data:
                    logger.warning("'file' artifact missing 'bytes' data")
                    continue
                try:
                    file_bytes = base64.b64decode(b64data)
                    ext = "bin"
                    if "png" in mime:
                        ext = "png"
                    elif "jpeg" in mime or "jpg" in mime:
                        ext = "jpg"
                    filename = f"agent_file_{uuid.uuid4().hex}.{ext}"
                    with open(filename, "wb") as f:
                        f.write(file_bytes)
                    images.append({"filename": filename, "bytes": file_bytes, "type": mime})
                    texts.append(filename)
                except Exception as e:
                    logger.error(f"Failed to decode/save file artifact: {e}")
            # (You can add more media/part types as needed)

    return {"texts": texts}


# `history`: [{"role": "user"|"agent", "text": "..."} ..]
# Returns: Dict with "agent":..., "response": ..., and updated "history"

async def run_a2a_router(query: str, history: Optional[List[Dict]] = None) -> Dict:
    """
    Runs a2a multi-agent router: discovers agents, routes, relays, and returns agent response.
    """
    # 1. Discover agents
    agents = await discover_agents(HOST, START_PORT, END_PORT)
    if not agents:
        return {"error": "No agents discovered."}

    if history is None:
        history = []

    history = history.copy()
    history.append({"role": "user", "text": query})

    # 2. Route to best agent via LLM
    prompt_messages = create_routing_prompt(history, agents)
    response = completion(
        model=MODEL_NAME,
        messages=prompt_messages,
        max_tokens=16,
        temperature=0.0,
    )
    agent_name = response.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
    matched_agent = None
    for agent in agents:
        if agent["card"].get("name", "").lower() == agent_name.lower():
            matched_agent = agent
            break
    if not matched_agent:
        matched_agent = agents[0]

    # 3. Build agent instruction and relay
    agent_instruction = make_agent_instruction(history, matched_agent["card"])
    agent_response = await send_instruction_to_agent(matched_agent, agent_instruction)
    if not agent_response:
        return {"error": f"Agent '{matched_agent['card'].get('name')}' did not return a valid response."}

    reply_text = extract_agent_text(agent_response)
    if reply_text:
        history.append({"role": "agent", "text": reply_text})

    return {
        "agent": matched_agent["card"].get("name"),
        "response": reply_text,
        "history": history,
        "raw_response": agent_response,
    }



mcp = FastMCP("mcp_server")

@mcp.tool()
async def handle_route_to_agents_with_a2a(query: str, history: Optional[List[Dict]] = None) -> Dict:
    """use this tool for any queries requiring up to date information"""
    return await run_a2a_router(query, history)



if __name__ == "__main__":
    mcp.run(transport='stdio')


import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np
import wandb
from collections import defaultdict

def get_fewshot_loader(dataset, num_per_class, batch_size=32):
    """Returns a DataLoader with num_per_class images per class."""
    targets = np.array(dataset.targets)
    idxs = []
    class_counts = defaultdict(int)
    for idx, label in enumerate(targets):
        if class_counts[label] < num_per_class:
            idxs.append(idx)
            class_counts[label] += 1
        if all(class_counts[k] >= num_per_class for k in range(10)):
            break
    sampler = torch.utils.data.Subset(dataset, idxs)
    return torch.utils.data.DataLoader(sampler, batch_size=batch_size, shuffle=True, num_workers=2)

def evaluate(model, dataloader, device):
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            pred = logits.argmax(1)
            correct += (pred == y).sum().item()
            total += len(y)
    return correct / total

def main():
    wandb.init(project="cifar10-fewshot-eval")

    # Use GPU if available
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Pretrained models and their final layer info
    models_info = {
        'resnet18': {
            'model': torchvision.models.resnet18(weights='DEFAULT'),
            'final_layer': ('fc', 512)
        },
        'densenet121': {
            'model': torchvision.models.densenet121(weights='DEFAULT'),
            'final_layer': ('classifier', 1024)
        },
        'vgg16': {
            'model': torchvision.models.vgg16(weights='DEFAULT'),
            'final_layer': ('classifier', 4096, 6)  # index 6 in classifier
        },
        'efficientnet_b0': {
            'model': torchvision.models.efficientnet_b0(weights='DEFAULT'),
            'final_layer': ('classifier', 1280)
        }
    }

    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406], 
            std=[0.229, 0.224, 0.225]
        )
    ])

    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.CIFAR10(
        root='./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=128, shuffle=False, num_workers=2)

    fewshot_loader = get_fewshot_loader(trainset, num_per_class=10, batch_size=16)

    for name, info in models_info.items():
        print(f"\nTraining and evaluating {name} ...")
        model = info['model']
        # Replace final layer for CIFAR-10
        if name == 'vgg16':
            clf = model.classifier
            clf[6] = nn.Linear(info['final_layer'][1], 10)
            model.classifier = clf
        else:
            setattr(model, info['final_layer'][0], nn.Linear(info['final_layer'][1], 10))
        model.to(device)

        # Only train the head
        for param in model.parameters():
            param.requires_grad = False
        if name == 'vgg16':
            for param in model.classifier[6].parameters():
                param.requires_grad = True
        else:
            for param in getattr(model, info['final_layer'][0]).parameters():
                param.requires_grad = True

        optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01)
        criterion = nn.CrossEntropyLoss()

        # Quick training (2 epochs)
        model.train()
        for epoch in range(2):
            total_loss = 0
            for x, y in fewshot_loader:
                x, y = x.to(device), y.to(device)
                optimizer.zero_grad()
                logits = model(x)
                loss = criterion(logits, y)
                loss.backward()
                optimizer.step()
                total_loss += loss.item() * len(y)
            avg_loss = total_loss / len(fewshot_loader.dataset)
            print(f"  Epoch {epoch+1}: avg loss {avg_loss:.4f}")

        # Eval
        acc = evaluate(model, testloader, device)
        print(f"{name}: Test set accuracy after few-shot training: {acc:.4f}")
        wandb.log({f"{name}_fewshot_acc": acc})

    wandb.finish()

if __name__ == "__main__":
    main()