Cohere Command R agent budget limit: adding a cost cap and loop detector to Command R tool-use agents

import cohere
from runguard import guard_async, LoopDetectedError, BudgetExceededError
import asyncio

co = cohere.AsyncClient(api_key="YOUR_COHERE_API_KEY")

TOOLS = [
    {
        "name": "search_documents",
        "description": "Search a document corpus for passages relevant to a query.",
        "parameter_definitions": {
            "query": {
                "description": "The search query string.",
                "type": "str",
                "required": True,
            }
        },
    },
    {
        "name": "get_document_by_id",
        "description": "Retrieve a specific document by its unique identifier.",
        "parameter_definitions": {
            "doc_id": {
                "description": "The document identifier.",
                "type": "str",
                "required": True,
            }
        },
    },
]


async def _call_command_r(chat_history: list, message: str, **kwargs) -> dict:
    response = await co.chat(
        model="command-r-plus",
        message=message,
        chat_history=chat_history,
        tools=TOOLS,
        **kwargs,
    )

    meta = response.meta
    # Command R+ pricing: $3/$15 per 1M input/output tokens
    input_tokens = (meta.tokens.input_tokens if meta and meta.tokens else 0)
    output_tokens = (meta.tokens.output_tokens if meta and meta.tokens else 0)
    usd = (input_tokens * 3.0 + output_tokens * 15.0) / 1_000_000

    # Signature: first tool call name, or "end_turn"
    tool_calls = response.tool_calls or []
    sig = tool_calls[0].name if tool_calls else "end_turn"

    return {"response": response, "usd": usd, "sig": sig}


guarded_call = guard_async(
    _call_command_r,
    budget={"max_usd": 1.0},
    loop={"repeats": 3, "max_cycle_len": 4},
)


async def run_command_r_agent(user_query: str) -> str:
    chat_history = []
    current_message = user_query

    while True:
        try:
            result = await guarded_call(chat_history, current_message)
        except LoopDetectedError as e:
            return f"Loop detected (pattern: {e.pattern!r}). Aborting after {len(chat_history)} turns."
        except BudgetExceededError as e:
            return f"Budget cap reached (${e.spent:.4f}). Aborting."

        response = result["response"]

        # Append this turn to history
        chat_history.append({
            "role": "USER",
            "message": current_message,
        })
        chat_history.append({
            "role": "CHATBOT",
            "message": response.text or "",
            "tool_calls": response.tool_calls or [],
        })

        if not response.tool_calls:
            # No more tool calls — the model has produced a final answer
            return response.text or ""

        # Execute tools and build tool results
        tool_results = []
        for tc in response.tool_calls:
            try:
                output = await dispatch_tool(tc.name, tc.parameters)
                tool_results.append(
                    cohere.ToolResult(
                        call=tc,
                        outputs=[{"result": str(output)}],
                    )
                )
            except Exception as exc:
                # Surface as structured error, not a plain string
                tool_results.append(
                    cohere.ToolResult(
                        call=tc,
                        outputs=[{"error": f"{type(exc).__name__}: {exc}"}],
                    )
                )

        # In Command R's tool-use pattern, the next API call passes tool_results
        # and uses the previous CHATBOT turn's tool_calls
        chat_history.append({
            "role": "TOOL",
            "tool_results": tool_results,
        })

        # Continue with an empty message — Cohere convention for tool-result continuation
        current_message = ""


async def dispatch_tool(name: str, parameters: dict) -> str:
    if name == "search_documents":
        return await search_documents(parameters["query"])
    elif name == "get_document_by_id":
        return await get_document_by_id(parameters["doc_id"])
    raise ValueError(f"Unknown tool: {name}")


# Stub implementations
async def search_documents(query: str) -> str: ...
async def get_document_by_id(doc_id: str) -> str: ...