import asyncio

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_ollama import ChatOllama

# Any tool with a complex schema works — using a minimal stand-in for
# Deep Agents' TodoListMiddleware.write_todos tool, which has a large
# JSON schema with nested objects/enums.
from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import Literal


class TodoItem(BaseModel):
    id: str = Field(description="Unique identifier")
    title: str = Field(description="Short title")
    status: Literal["not_started", "in_progress", "completed"] = "not_started"
    priority: Literal["high", "medium", "low"] = "medium"
    details: str = Field(default="", description="Extended description")


class TodoList(BaseModel):
    todos: list[TodoItem] = Field(description="Full list of todo items")


@tool(args_schema=TodoList)
def write_todos(todos: list[TodoItem]) -> str:
    """Write and replace the full todo list."""
    return "ok"


# A large system prompt (~2-3k tokens), similar in size to Deep Agents'
# BASE_AGENT_PROMPT. The exact content doesn't matter — length does.
LARGE_PROMPT = "You are a helpful assistant.\n\n" + (
    "## Guidelines\n\n"
    + "\n".join(
        f"- Rule {i}: Follow best practices for step {i}." for i in range(200)
    )
)


async def count(label, runnable, messages):
    counts = {"reasoning": 0, "text": 0}
    async for chunk in runnable.astream(messages):
        ak = getattr(chunk, "additional_kwargs", None)
        rc = ak.get("reasoning_content") if isinstance(ak, dict) else None
        content = getattr(chunk, "content", None)
        if rc:
            counts["reasoning"] += 1
        elif content:
            counts["text"] += 1
    print(label, counts)


async def main():
    llm = ChatOllama(
        model="gemma4:26b",
        base_url="http://localhost:11434",
        reasoning=True,
        num_ctx=8192,
        temperature=0,
    )

    small = [
        SystemMessage(content="You are helpful."),
        HumanMessage(content="hola"),
    ]
    big = [
        SystemMessage(content=LARGE_PROMPT),
        HumanMessage(content="hola"),
    ]

    await count("no tools, small prompt", llm, small)
    await count("no tools, big prompt", llm, big)
    await count("tools, small prompt", llm.bind_tools([write_todos]), small)
    await count("tools, big prompt", llm.bind_tools([write_todos]), big)


asyncio.run(main())

---

no tools, small prompt {'reasoning': 12, 'text': 5}
no tools, big prompt {'reasoning': 15, 'text': 6}
tools, small prompt {'reasoning': 8, 'text': 4}
tools, big prompt {'reasoning': 0, 'text': 7}   # <-- reasoning gone