from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.agents import create_agent

async def init_agent_with_mcp():
    client = MultiServerMCPClient(
        servers={
            "browser": {
                "command": "mcp-browser-server",
                "args": [],
            },
            # 也可以挂更多 MCP 服务器
        }
    )

    tools = await client.get_tools()

    agent = create_agent(
        model="claude-3-5-sonnet-20251001",
        tools=tools,
    )
    return agent

result = await agent.ainvoke(
    {"messages": [{"role": "user", "content": "截一张当前页面的图，并简单说明。"}]}
)

for message in result["messages"]:
    if message.type == "tool":
        # 原始内容（provider 格式）
        print("Raw content:", message.content)

        # 标准 content_blocks
        for block in message.content_blocks:
            if block["type"] == "text":
                print("Text:", block["text"])
            elif block["type"] == "image":
                print("Image URL:", block.get("url"))
                print("Image base64 prefix:", block.get("base64", "")[:50], "...")

from typing import Any, Dict, List, Tuple
from langchain_core.messages import BaseMessage

def collect_tool_messages(messages: List[BaseMessage]) -> Dict[str, List[BaseMessage]]:
    """按工具名聚合 ToolMessage，便于后续处理。"""
    tool_messages: Dict[str, List[BaseMessage]] = {}
    for msg in messages:
        if msg.type == "tool":
            tool_name = msg.name or "unknown_tool"
            tool_messages.setdefault(tool_name, []).append(msg)
    return tool_messages

def parse_tool_multimodal_blocks(
    tool_msg: BaseMessage,
) -> Tuple[List[str], List[Dict[str, Any]]]:
    """解析单条工具消息中的文本与图片块。"""
    texts = []
    images = []

    for block in tool_msg.content_blocks:
        if block.get("type") == "text":
            texts.append(block.get("text") or "")
        elif block.get("type") == "image":
            images.append(block)

    return texts, images

### 4.2 进一步处理：用 deepseek-ocr + ChatOllama 为截图生成文本

在很多场景里，我们不仅想“看到截图”，还想：

- 把截图里的文字内容纳入后续检索 / 审计；
- 在故障回溯时，能直接搜索到「某次截图中出现的报错信息」；
- 为后续的模型调用提供更丰富的上下文。

这时候，就可以在上面的基础上再加一层 **OCR 处理**：

```python
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage

def build_ocr_model() -> ChatOllama:
    base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
    model_id = os.getenv("OCR_MODEL", "deepseek-ocr")
    return ChatOllama(
        model=model_id,
        base_url=base_url,
        temperature=0,
        output_version="v1",
    )

async def ocr_images_from_tool_message(tool_msg: BaseMessage) -> str:
    """对 ToolMessage 中的图片块做 OCR，返回合并后的文本。"""
    _, images = parse_tool_multimodal_blocks(tool_msg)
    if not images:
        return ""

    ocr_model = build_ocr_model()
    ocr_texts = []

    for img in images:
        base64_data = img.get("base64")
        if not base64_data:
            # 只有 URL 的情况可以按需自行下载，这里先跳过
            continue

        prompt = (
            "请逐行识别图片中的文字，保持原有行结构，"
            "只输出纯文本，不要额外解释。"
        )

        message = HumanMessage(
            content=[
                {"type": "text", "text": prompt},
                {
                    "type": "image",
                    "base64": base64_data,
                    "mime_type": img.get("mime_type") or "image/png",
                },
            ]
        )

        resp = await ocr_model.ainvoke([message])

        blocks = getattr(resp, "content_blocks", None)
        if blocks:
            text_blocks = [
                b.get("text") or ""
                for b in blocks
                if b.get("type") == "text" and b.get("text")
            ]
            if text_blocks:
                ocr_texts.append("\n".join(text_blocks).strip())
                continue

        # 回退到 text / content
        if hasattr(resp, "text") and isinstance(resp.text, str):
            ocr_texts.append(resp.text.strip())
        elif isinstance(resp.content, str):
            ocr_texts.append(resp.content.strip())

    return "\n\n".join(t for t in ocr_texts if t)

ocr_text = await ocr_images_from_tool_message(msg)
if ocr_text:
    print("=== OCR 识别结果（来自 deepseek-ocr + ChatOllama） ===")
    print(ocr_text)

async def run_with_screenshot(agent):
    result = await agent.ainvoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "截一张当前页面的图，并告诉我当前页面的核心信息。",
                }
            ]
        }
    )

    messages = result["messages"]
    tool_messages = collect_tool_messages(messages)

    for tool_name, msgs in tool_messages.items():
        print(f"工具 {tool_name} 共被调用 {len(msgs)} 次。")

        for idx, msg in enumerate(msgs):
            print(f"--- 第 {idx+1} 次调用 ---")
            texts, images = parse_tool_multimodal_blocks(msg)

            print("文本部分：")
            print("\n".join(texts))

            print("图片部分：")
            for img in images:
                print("  mime_type:", img.get("mime_type"))
                print("  url:", img.get("url"))
                print("  base64 前 50 字符:", img.get("base64", "")[:50], "...")

def log_tool_multimodal(tool_msg: BaseMessage):
    for block in tool_msg.content_blocks:
        if block.get("type") == "text":
            log_text(block["text"])
        elif block.get("type") == "image":
            url = block.get("url")
            if not url:
                url = upload_image_base64(block["base64"])
            log_image_reference(url, meta={"tool": tool_msg.name})