Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 29 additions & 9 deletions agent/title_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,22 @@
},
}


def _title_request_extra_body(
response_format_supported: Optional[bool],
) -> dict[str, Any]:
"""Return structured-output request fields when the route supports them.

Unknown capability is intentionally treated like unsupported capability:
title parsing already has JSON/prose fallbacks, while an optimistic
response_format parameter can make an otherwise valid auxiliary request
fail at the provider boundary.
"""
if response_format_supported is True:
return {"response_format": _TITLE_RESPONSE_FORMAT}
return {}


# Control-tag wrappers that surround machine-authored content inside what is
# nominally a "user" message. Titling from these is what produces a session
# named after a slash command or an injected reminder rather than the user's
Expand Down Expand Up @@ -336,6 +352,7 @@ def generate_title(
failure_callback: Optional[FailureCallback] = None,
main_runtime: dict = None,
runtime_validator: Optional[RuntimeValidator] = None,
response_format_supported: Optional[bool] = None,
) -> Optional[str]:
"""Generate a session title from the user's opening message.

Expand Down Expand Up @@ -392,17 +409,20 @@ def generate_title(
]

try:
response = call_llm(
task="title_generation",
messages=messages,
request_extra_body = _title_request_extra_body(response_format_supported)
request_kwargs = {
"task": "title_generation",
"messages": messages,
# A title is a handful of tokens. The old 500-token ceiling let a
# chatty model burn seconds generating prose we then threw away.
max_tokens=64,
temperature=0.3,
timeout=timeout,
main_runtime=main_runtime,
extra_body={"response_format": _TITLE_RESPONSE_FORMAT},
)
"max_tokens": 64,
"temperature": 0.3,
"timeout": timeout,
"main_runtime": main_runtime,
}
if request_extra_body:
request_kwargs["extra_body"] = request_extra_body
response = call_llm(**request_kwargs)
content = response.choices[0].message.content or ""
return _clean_title(_extract_title_text(content))
except Exception as e:
Expand Down
48 changes: 48 additions & 0 deletions tests/agent/test_title_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,54 @@ def mock_call_llm(**kwargs):
assert captured_kwargs["task"] == "title_generation"
assert captured_kwargs["timeout"] is None

@pytest.mark.parametrize("capability", [None, False])
def test_title_generation_omits_response_format_without_capability(self, capability):
captured_kwargs = {}

def mock_call_llm(**kwargs):
captured_kwargs.update(kwargs)
response = MagicMock()
response.choices = [MagicMock()]
response.choices[0].message.content = '{"title": "Fix login"}'
return response

with patch("agent.title_generator.call_llm", side_effect=mock_call_llm):
assert generate_title(
"fix login", response_format_supported=capability
) == "Fix login"

assert "extra_body" not in captured_kwargs
assert "response_format" not in captured_kwargs

def test_title_generation_sends_response_format_when_capability_supported(self):
captured_kwargs = {}

def mock_call_llm(**kwargs):
captured_kwargs.update(kwargs)
response = MagicMock()
response.choices = [MagicMock()]
response.choices[0].message.content = '{"title": "Fix login"}'
return response

with patch("agent.title_generator.call_llm", side_effect=mock_call_llm):
assert generate_title(
"fix login", response_format_supported=True
) == "Fix login"

response_format = captured_kwargs["extra_body"]["response_format"]
assert response_format["type"] == "json_schema"
assert response_format["json_schema"]["name"] == "session_title"
assert response_format["json_schema"]["schema"]["required"] == ["title"]

def test_title_generation_parses_plain_prose_without_response_format(self):
response = MagicMock()
response.choices = [MagicMock()]
response.choices[0].message.content = "Fix login flow on mobile"

with patch("agent.title_generator.call_llm", return_value=response):
assert generate_title(
"fix login", response_format_supported=None
) == "Fix login flow on mobile"


def test_strips_think_blocks(self):
Expand Down