Spaces:

smolagents
/

ml-intern

Running on CPU Upgrade

App Files Files Community

lewtun HF Staff OpenAI Codex commited on Jun 3

Commit

e6f00a2

unverified ·

1 Parent(s): b103510

Restore prompt caching for FAL router models (#286)

Browse files

* Restore prompt caching for FAL router models

Co-authored-by: OpenAI Codex <codex@openai.com>

* Tune prompt caching hints for GPT-5.5

Co-authored-by: OpenAI Codex <codex@openai.com>

* Address backlog prompt caching import review

Co-authored-by: OpenAI Codex <codex@openai.com>

---------

Co-authored-by: OpenAI Codex <codex@openai.com>

Files changed (6) hide show

agent/context_manager/manager.py +8 -0
agent/core/agent_loop.py +19 -6
agent/core/prompt_caching.py +186 -0
agent/tools/research_tool.py +13 -4
scripts/prioritize_backlog.py +8 -1
tests/unit/test_prompt_caching.py +220 -0

agent/context_manager/manager.py CHANGED Viewed

@@ -13,6 +13,8 @@ import yaml
 from jinja2 import Template
 from litellm import Message, acompletion
 logger = logging.getLogger(__name__)
 _HF_WHOAMI_URL = "https://huggingface.co/api/whoami-v2"
@@ -145,6 +147,12 @@ async def summarize_messages(
         reasoning_effort="high",
         bill_to_user=getattr(session, "premium_user_billed", False),
     )
     _t0 = time.monotonic()
     response = await acompletion(
         messages=prompt_messages,

 from jinja2 import Template
 from litellm import Message, acompletion
+from agent.core.prompt_caching import with_prompt_cache_params, with_prompt_caching
 logger = logging.getLogger(__name__)
 _HF_WHOAMI_URL = "https://huggingface.co/api/whoami-v2"
         reasoning_effort="high",
         bill_to_user=getattr(session, "premium_user_billed", False),
     )
+    llm_params = with_prompt_cache_params(
+        llm_params, session_id=getattr(session, "session_id", None)
+    )
+    prompt_messages, tool_specs = with_prompt_caching(
+        prompt_messages, tool_specs, llm_params
+    )
     _t0 = time.monotonic()
     response = await acompletion(
         messages=prompt_messages,

agent/core/agent_loop.py CHANGED Viewed

@@ -27,6 +27,7 @@ from agent.messaging.gateway import NotificationGateway
 from agent.core import telemetry
 from agent.core.doom_loop import check_for_doom_loop
 from agent.core.llm_params import _resolve_llm_params
 from agent.core.session import DEFAULT_SESSION_LOG_DIR, Event, OpType, Session
 from agent.core.tools import ToolRouter
 from agent.tools.jobs_tool import CPU_FLAVORS
@@ -819,14 +820,20 @@ async def _call_llm_streaming(
     t_start = time.monotonic()
     for _llm_attempt in range(_MAX_LLM_RETRIES):
         try:
             response = await acompletion(
-                messages=messages,
-                tools=tools,
                 tool_choice="auto",
                 stream=True,
                 stream_options={"include_usage": True},
                 timeout=600,
-                **llm_params,
             )
             break
         except ContextWindowExceededError:
@@ -958,13 +965,19 @@ async def _call_llm_non_streaming(
     t_start = time.monotonic()
     for _llm_attempt in range(_MAX_LLM_RETRIES):
         try:
             response = await acompletion(
-                messages=messages,
-                tools=tools,
                 tool_choice="auto",
                 stream=False,
                 timeout=600,
-                **llm_params,
             )
             break
         except ContextWindowExceededError:

 from agent.core import telemetry
 from agent.core.doom_loop import check_for_doom_loop
 from agent.core.llm_params import _resolve_llm_params
+from agent.core.prompt_caching import with_prompt_cache_params, with_prompt_caching
 from agent.core.session import DEFAULT_SESSION_LOG_DIR, Event, OpType, Session
 from agent.core.tools import ToolRouter
 from agent.tools.jobs_tool import CPU_FLAVORS
     t_start = time.monotonic()
     for _llm_attempt in range(_MAX_LLM_RETRIES):
         try:
+            request_llm_params = with_prompt_cache_params(
+                llm_params, session_id=getattr(session, "session_id", None)
+            )
+            cached_messages, cached_tools = with_prompt_caching(
+                messages, tools, request_llm_params
+            )
             response = await acompletion(
+                messages=cached_messages,
+                tools=cached_tools,
                 tool_choice="auto",
                 stream=True,
                 stream_options={"include_usage": True},
                 timeout=600,
+                **request_llm_params,
             )
             break
         except ContextWindowExceededError:
     t_start = time.monotonic()
     for _llm_attempt in range(_MAX_LLM_RETRIES):
         try:
+            request_llm_params = with_prompt_cache_params(
+                llm_params, session_id=getattr(session, "session_id", None)
+            )
+            cached_messages, cached_tools = with_prompt_caching(
+                messages, tools, request_llm_params
+            )
             response = await acompletion(
+                messages=cached_messages,
+                tools=cached_tools,
                 tool_choice="auto",
                 stream=False,
                 timeout=600,
+                **request_llm_params,
             )
             break
         except ContextWindowExceededError:

agent/core/prompt_caching.py ADDED Viewed

	@@ -0,0 +1,186 @@

+"""Prompt-cache helpers for HF Router FAL requests.
+The HF Router/OpenRouter path uses provider-native prompt caching. Anthropic
+models need explicit JSON ``cache_control`` content blocks; OpenAI models cache
+eligible prefixes automatically and accept routing/retention hints in the body.
+Headers like ``X-OpenRouter-Cache`` control response caching, not prompt
+caching through this route.
+"""
+from typing import Any
+from agent.core.model_ids import HF_ROUTER_BASE_URL
+_CACHE_CONTROL = {"type": "ephemeral"}
+_CACHEABLE_ROLES = {"system", "user"}
+_OPENROUTER_SESSION_ID_MAX_LENGTH = 256
+def _is_fal_router_request(llm_params: dict[str, Any]) -> bool:
+    api_base = str(llm_params.get("api_base") or "").rstrip("/")
+    return api_base == HF_ROUTER_BASE_URL and ":fal" in _router_model(llm_params)
+def _router_model(llm_params: dict[str, Any]) -> str:
+    model = str(llm_params.get("model") or "")
+    return model.removeprefix("openai/")
+def _uses_explicit_cache_control(llm_params: dict[str, Any]) -> bool:
+    if not _is_fal_router_request(llm_params):
+        return False
+    return _router_model(llm_params).startswith("anthropic/")
+def _is_openai_gpt55(llm_params: dict[str, Any]) -> bool:
+    if not _is_fal_router_request(llm_params):
+        return False
+    return _router_model(llm_params).startswith("openai/gpt-5.5")
+def _merge_extra_body(
+    llm_params: dict[str, Any], updates: dict[str, Any]
+) -> dict[str, Any]:
+    if not updates:
+        return llm_params
+    cached_params = dict(llm_params)
+    extra_body = dict(cached_params.get("extra_body") or {})
+    extra_body.update(updates)
+    cached_params["extra_body"] = extra_body
+    return cached_params
+def with_prompt_cache_params(
+    llm_params: dict[str, Any],
+    *,
+    session_id: str | None = None,
+) -> dict[str, Any]:
+    """Return LiteLLM params with provider-native prompt-cache body hints."""
+    if not _is_fal_router_request(llm_params):
+        return llm_params
+    updates: dict[str, Any] = {}
+    if session_id:
+        stable_session_id = session_id[:_OPENROUTER_SESSION_ID_MAX_LENGTH]
+        updates["session_id"] = stable_session_id
+        if _is_openai_gpt55(llm_params):
+            updates["prompt_cache_key"] = stable_session_id
+    if _is_openai_gpt55(llm_params):
+        updates["prompt_cache_retention"] = "24h"
+    return _merge_extra_body(llm_params, updates)
+def _message_role(message: Any) -> str | None:
+    if isinstance(message, dict):
+        role = message.get("role")
+    else:
+        role = getattr(message, "role", None)
+    return role if isinstance(role, str) else None
+def _message_content(message: Any) -> Any:
+    if isinstance(message, dict):
+        return message.get("content")
+    return getattr(message, "content", None)
+def _message_to_dict(message: Any) -> dict[str, Any]:
+    if isinstance(message, dict):
+        return dict(message)
+    if hasattr(message, "model_dump"):
+        return message.model_dump(exclude_none=True)
+    raise TypeError(f"Unsupported message type for prompt caching: {type(message)!r}")
+def _has_cacheable_text(content: Any) -> bool:
+    if isinstance(content, str):
+        return bool(content)
+    if not isinstance(content, list):
+        return False
+    return any(
+        isinstance(block, dict)
+        and block.get("type") == "text"
+        and isinstance(block.get("text"), str)
+        and bool(block.get("text"))
+        for block in content
+    )
+def _cache_target_index(messages: list[Any]) -> int | None:
+    if len(messages) < 2:
+        return None
+    for idx in range(len(messages) - 2, -1, -1):
+        message = messages[idx]
+        if _message_role(message) not in _CACHEABLE_ROLES:
+            continue
+        if _has_cacheable_text(_message_content(message)):
+            return idx
+    return None
+def _content_with_cache_control(content: Any) -> list[dict[str, Any]]:
+    if isinstance(content, str):
+        return [
+            {"type": "text", "text": content, "cache_control": dict(_CACHE_CONTROL)}
+        ]
+    blocks = [dict(block) if isinstance(block, dict) else block for block in content]
+    for idx in range(len(blocks) - 1, -1, -1):
+        block = blocks[idx]
+        if (
+            isinstance(block, dict)
+            and block.get("type") == "text"
+            and isinstance(block.get("text"), str)
+            and bool(block.get("text"))
+        ):
+            cached = dict(block)
+            cached["cache_control"] = dict(_CACHE_CONTROL)
+            blocks[idx] = cached
+            break
+    return blocks
+def _tools_with_cache_control(tools: list[dict] | None) -> list[dict] | None:
+    if not tools:
+        return tools
+    cached_tools = list(tools)
+    last_tool = dict(cached_tools[-1])
+    last_tool["cache_control"] = dict(_CACHE_CONTROL)
+    cached_tools[-1] = last_tool
+    return cached_tools
+def with_prompt_caching(
+    messages: list[Any],
+    tools: list[dict] | None,
+    llm_params: dict[str, Any],
+) -> tuple[list[Any], list[dict] | None]:
+    """Return outgoing messages with explicit cache breakpoints when needed.
+    The newest message is treated as dynamic. For Anthropic FAL models, the
+    cache breakpoint is placed on the closest earlier system/user text block so
+    provider-side caching covers the stable prefix without changing persisted
+    conversation history. The final tool spec is also marked so stable tool
+    definitions are cached.
+    """
+    if not _uses_explicit_cache_control(llm_params):
+        return messages, tools
+    cached_tools = _tools_with_cache_control(tools)
+    idx = _cache_target_index(messages)
+    if idx is None:
+        return messages, cached_tools
+    cached_message = _message_to_dict(messages[idx])
+    cached_message["content"] = _content_with_cache_control(
+        cached_message.get("content")
+    )
+    cached_messages = list(messages)
+    cached_messages[idx] = cached_message
+    return cached_messages, cached_tools

agent/tools/research_tool.py CHANGED Viewed

@@ -18,6 +18,7 @@ from agent.core import telemetry
 from agent.core.doom_loop import check_for_doom_loop
 from agent.core.llm_params import _resolve_llm_params
 from agent.core.model_ids import strip_huggingface_model_prefix
 from agent.core.session import Event
 logger = logging.getLogger(__name__)
@@ -259,6 +260,9 @@ async def research_handler(
         reasoning_effort=_capped,
         bill_to_user=getattr(session, "premium_user_billed", False),
     )
     # Get read-only tool specs from the session's tool router
     tool_specs = [
@@ -336,8 +340,9 @@ async def research_handler(
             )
             try:
                 _t0 = time.monotonic()
                 response = await acompletion(
-                    messages=messages,
                     tools=None,  # no tools — force text response
                     stream=False,
                     timeout=120,
@@ -383,9 +388,12 @@ async def research_handler(
         try:
             _t0 = time.monotonic()
             response = await acompletion(
-                messages=messages,
-                tools=tool_specs if tool_specs else None,
                 tool_choice="auto",
                 stream=False,
                 timeout=120,
@@ -499,8 +507,9 @@ async def research_handler(
     )
     try:
         _t0 = time.monotonic()
         response = await acompletion(
-            messages=messages,
             tools=None,
             stream=False,
             timeout=120,

 from agent.core.doom_loop import check_for_doom_loop
 from agent.core.llm_params import _resolve_llm_params
 from agent.core.model_ids import strip_huggingface_model_prefix
+from agent.core.prompt_caching import with_prompt_cache_params, with_prompt_caching
 from agent.core.session import Event
 logger = logging.getLogger(__name__)
         reasoning_effort=_capped,
         bill_to_user=getattr(session, "premium_user_billed", False),
     )
+    llm_params = with_prompt_cache_params(
+        llm_params, session_id=getattr(session, "session_id", None)
+    )
     # Get read-only tool specs from the session's tool router
     tool_specs = [
             )
             try:
                 _t0 = time.monotonic()
+                cached_messages, _ = with_prompt_caching(messages, None, llm_params)
                 response = await acompletion(
+                    messages=cached_messages,
                     tools=None,  # no tools — force text response
                     stream=False,
                     timeout=120,
         try:
             _t0 = time.monotonic()
+            cached_messages, cached_tools = with_prompt_caching(
+                messages, tool_specs if tool_specs else None, llm_params
+            )
             response = await acompletion(
+                messages=cached_messages,
+                tools=cached_tools,
                 tool_choice="auto",
                 stream=False,
                 timeout=120,
     )
     try:
         _t0 = time.monotonic()
+        cached_messages, _ = with_prompt_caching(messages, None, llm_params)
         response = await acompletion(
+            messages=cached_messages,
             tools=None,
             stream=False,
             timeout=120,

scripts/prioritize_backlog.py CHANGED Viewed

@@ -33,6 +33,11 @@ PROJECT_ROOT = Path(__file__).resolve().parent.parent
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 GITHUB_API = "https://api.github.com"
 DEFAULT_GITHUB_REPO = "huggingface/ml-intern"
 DEFAULT_HF_SPACE = "smolagents/ml-intern"
@@ -1297,10 +1302,12 @@ async def _call_json_llm(
         completion_func = acompletion
     attempt_messages = list(messages)
     last_error: Exception | None = None
     for attempt in range(retries + 1):
         response = await completion_func(
-            messages=attempt_messages,
             max_completion_tokens=max_completion_tokens,
             temperature=_temperature_for_params(llm_params),
             **llm_params,

 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
+from agent.core.prompt_caching import (  # noqa: E402
+    with_prompt_cache_params,
+    with_prompt_caching,
+)
 GITHUB_API = "https://api.github.com"
 DEFAULT_GITHUB_REPO = "huggingface/ml-intern"
 DEFAULT_HF_SPACE = "smolagents/ml-intern"
         completion_func = acompletion
     attempt_messages = list(messages)
+    llm_params = with_prompt_cache_params(llm_params)
     last_error: Exception | None = None
     for attempt in range(retries + 1):
+        cached_messages, _ = with_prompt_caching(attempt_messages, None, llm_params)
         response = await completion_func(
+            messages=cached_messages,
             max_completion_tokens=max_completion_tokens,
             temperature=_temperature_for_params(llm_params),
             **llm_params,

tests/unit/test_prompt_caching.py ADDED Viewed

	@@ -0,0 +1,220 @@

+from litellm import Message
+from agent.core.model_ids import HF_ROUTER_BASE_URL
+from agent.core.prompt_caching import with_prompt_cache_params, with_prompt_caching
+def _anthropic_fal_params() -> dict:
+    return {
+        "model": "openai/anthropic/claude-sonnet-4-6:fal-ai",
+        "api_base": HF_ROUTER_BASE_URL,
+    }
+def _gpt55_fal_params() -> dict:
+    return {
+        "model": "openai/openai/gpt-5.5:fal-ai",
+        "api_base": HF_ROUTER_BASE_URL,
+    }
+def test_prompt_caching_marks_system_prefix_and_tools_for_fal_router_model():
+    messages = [
+        Message(role="system", content="stable system prompt"),
+        Message(role="user", content="current question"),
+    ]
+    tools = [
+        {"type": "function", "function": {"name": "read"}},
+        {"type": "function", "function": {"name": "write"}},
+    ]
+    cached_messages, cached_tools = with_prompt_caching(
+        messages, tools, _anthropic_fal_params()
+    )
+    assert cached_tools is not tools
+    assert cached_tools == [
+        {"type": "function", "function": {"name": "read"}},
+        {
+            "type": "function",
+            "function": {"name": "write"},
+            "cache_control": {"type": "ephemeral"},
+        },
+    ]
+    assert "cache_control" not in tools[-1]
+    assert cached_messages is not messages
+    assert cached_messages[0] == {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": "stable system prompt",
+                "cache_control": {"type": "ephemeral"},
+            }
+        ],
+    }
+    assert messages[0].content == "stable system prompt"
+def test_prompt_caching_marks_last_stable_user_before_current_message():
+    messages = [
+        {"role": "system", "content": "stable system"},
+        {"role": "user", "content": "stable reference"},
+        {"role": "assistant", "content": "previous answer"},
+        {"role": "user", "content": "current question"},
+    ]
+    cached_messages, _ = with_prompt_caching(messages, None, _anthropic_fal_params())
+    assert cached_messages[0]["content"] == "stable system"
+    assert cached_messages[1]["content"] == [
+        {
+            "type": "text",
+            "text": "stable reference",
+            "cache_control": {"type": "ephemeral"},
+        }
+    ]
+    assert messages[1]["content"] == "stable reference"
+def test_prompt_caching_marks_last_text_block_in_content_list():
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "stable part one"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.test/i.png"},
+                },
+                {"type": "text", "text": "stable part two"},
+            ],
+        },
+        {"role": "user", "content": "current question"},
+    ]
+    cached_messages, _ = with_prompt_caching(messages, None, _anthropic_fal_params())
+    assert cached_messages[0]["content"][0] == {
+        "type": "text",
+        "text": "stable part one",
+    }
+    assert cached_messages[0]["content"][1] == {
+        "type": "image_url",
+        "image_url": {"url": "https://example.test/i.png"},
+    }
+    assert cached_messages[0]["content"][2] == {
+        "type": "text",
+        "text": "stable part two",
+        "cache_control": {"type": "ephemeral"},
+    }
+    assert "cache_control" not in messages[0]["content"][2]
+def test_prompt_caching_marks_tools_without_message_prefix():
+    messages = [{"role": "user", "content": "current question"}]
+    tools = [{"type": "function", "function": {"name": "read"}}]
+    cached_messages, cached_tools = with_prompt_caching(
+        messages, tools, _anthropic_fal_params()
+    )
+    assert cached_messages is messages
+    assert cached_tools == [
+        {
+            "type": "function",
+            "function": {"name": "read"},
+            "cache_control": {"type": "ephemeral"},
+        }
+    ]
+    assert "cache_control" not in tools[0]
+def test_prompt_caching_is_noop_for_non_fal_router_model():
+    messages = [
+        {"role": "system", "content": "stable system"},
+        {"role": "user", "content": "current question"},
+    ]
+    llm_params = {
+        "model": "openai/moonshotai/Kimi-K2.6",
+        "api_base": HF_ROUTER_BASE_URL,
+    }
+    cached_messages, cached_tools = with_prompt_caching(messages, None, llm_params)
+    assert cached_messages is messages
+    assert cached_tools is None
+def test_prompt_caching_is_noop_for_gpt55_fal_router_model():
+    messages = [
+        {"role": "system", "content": "stable system"},
+        {"role": "user", "content": "current question"},
+    ]
+    tools = [{"type": "function", "function": {"name": "read"}}]
+    cached_messages, cached_tools = with_prompt_caching(
+        messages, tools, _gpt55_fal_params()
+    )
+    assert cached_messages is messages
+    assert cached_tools is tools
+def test_prompt_caching_is_noop_for_non_router_fal_model():
+    messages = [
+        {"role": "system", "content": "stable system"},
+        {"role": "user", "content": "current question"},
+    ]
+    llm_params = {
+        "model": "openai/anthropic/claude-sonnet-4-6:fal-ai",
+        "api_base": "http://localhost:8000/v1",
+    }
+    cached_messages, _ = with_prompt_caching(messages, None, llm_params)
+    assert cached_messages is messages
+def test_prompt_cache_params_add_session_id_for_fal_router_model():
+    llm_params = _anthropic_fal_params()
+    cached_params = with_prompt_cache_params(llm_params, session_id="session-1")
+    assert cached_params is not llm_params
+    assert cached_params["extra_body"] == {"session_id": "session-1"}
+    assert "extra_body" not in llm_params
+def test_prompt_cache_params_merges_gpt55_cache_hints():
+    llm_params = {
+        **_gpt55_fal_params(),
+        "extra_body": {"reasoning_effort": "high"},
+    }
+    cached_params = with_prompt_cache_params(llm_params, session_id="session-1")
+    assert cached_params["extra_body"] == {
+        "reasoning_effort": "high",
+        "session_id": "session-1",
+        "prompt_cache_key": "session-1",
+        "prompt_cache_retention": "24h",
+    }
+    assert llm_params["extra_body"] == {"reasoning_effort": "high"}
+def test_prompt_cache_params_adds_gpt55_retention_without_session():
+    cached_params = with_prompt_cache_params(_gpt55_fal_params())
+    assert cached_params["extra_body"] == {"prompt_cache_retention": "24h"}
+def test_prompt_cache_params_is_noop_for_non_router_model():
+    llm_params = {
+        "model": "openai/openai/gpt-5.5:fal-ai",
+        "api_base": "http://localhost:8000/v1",
+    }
+    cached_params = with_prompt_cache_params(llm_params, session_id="session-1")
+    assert cached_params is llm_params