diff --git a/src/openjarvis/core/types.py b/src/openjarvis/core/types.py index 62cdc9f4..add45385 100644 --- a/src/openjarvis/core/types.py +++ b/src/openjarvis/core/types.py @@ -63,7 +63,7 @@ class Message: """A single chat message (OpenAI-compatible structure).""" role: Role - content: str = "" + content: str | None = "" name: Optional[str] = None tool_calls: Optional[List[ToolCall]] = None tool_call_id: Optional[str] = None @@ -73,6 +73,11 @@ class Message: # empty for text-only messages (the common case). images: Optional[List[str]] = None + @property + def text(self) -> str: + """Return message content as text, treating ``None`` as empty.""" + return self.content or "" + @dataclass(slots=True) class Conversation: diff --git a/src/openjarvis/engine/_base.py b/src/openjarvis/engine/_base.py index 97f9a23e..6ecaf81a 100644 --- a/src/openjarvis/engine/_base.py +++ b/src/openjarvis/engine/_base.py @@ -13,6 +13,22 @@ class EngineConnectionError(Exception): """Raised when an engine is unreachable.""" +_REASONING_METADATA_KEYS = ("reasoning_content", "thinking") + + +def _message_estimated_chars(message: Message) -> int: + parts = [message.text] + for key in _REASONING_METADATA_KEYS: + value = message.metadata.get(key) + if isinstance(value, str): + parts.append(value) + for tc in message.tool_calls or []: + parts.extend((tc.id, tc.name, tc.arguments)) + if message.tool_call_id: + parts.append(message.tool_call_id) + return sum(len(part) for part in parts) + + def messages_to_dicts(messages: Sequence[Message]) -> List[Dict[str, Any]]: """Convert ``Message`` objects to OpenAI-format dicts.""" out: List[Dict[str, Any]] = [] @@ -53,9 +69,11 @@ def estimate_prompt_tokens(messages: Sequence[Message]) -> int: provider would charge. Uses ~4 characters per token (standard BPE average for English) plus - a small per-message overhead for role markers and separators. + a small per-message overhead for role markers and separators. Counts + content, reasoning metadata, tool-call payloads, and tool result IDs + because all are replayed into later prompt turns when present. """ - total_chars = sum(len(m.content or "") for m in messages) + total_chars = sum(_message_estimated_chars(m) for m in messages) # ~4 tokens overhead per message for role markers / separators overhead = len(messages) * 4 return max(1, total_chars // 4 + overhead) diff --git a/tests/core/test_types.py b/tests/core/test_types.py index 2d80ed34..b91817f9 100644 --- a/tests/core/test_types.py +++ b/tests/core/test_types.py @@ -35,9 +35,15 @@ class TestMessage: msg = Message(role=Role.USER, content="hello") assert msg.role == Role.USER assert msg.content == "hello" + assert msg.text == "hello" assert msg.tool_calls is None assert msg.metadata == {} + def test_none_content_text_helper(self) -> None: + msg = Message(role=Role.ASSISTANT, content=None) + assert msg.content is None + assert msg.text == "" + def test_tool_calls(self) -> None: tc = ToolCall(id="1", name="calc", arguments='{"x": 1}') msg = Message(role=Role.ASSISTANT, content="", tool_calls=[tc]) diff --git a/tests/engine/test_base.py b/tests/engine/test_base.py index 397c18eb..f466b4da 100644 --- a/tests/engine/test_base.py +++ b/tests/engine/test_base.py @@ -9,9 +9,52 @@ def test_estimate_prompt_tokens_handles_none_content_tool_call_turn() -> None: Message(role=Role.USER, content="hi"), Message( role=Role.ASSISTANT, - content=None, # type: ignore[arg-type] + content=None, tool_calls=[ToolCall(id="call_1", name="lookup", arguments="{}")], ), ] - assert estimate_prompt_tokens(messages) == 8 + assert estimate_prompt_tokens(messages) == 12 + + +def test_estimate_prompt_tokens_counts_tool_call_arguments() -> None: + base = [ + Message(role=Role.USER, content="hi"), + Message(role=Role.ASSISTANT, content=None), + ] + with_tool_call = [ + Message(role=Role.USER, content="hi"), + Message( + role=Role.ASSISTANT, + content=None, + tool_calls=[ToolCall(id="", name="", arguments="abcdefgh")], + ), + ] + + assert estimate_prompt_tokens(with_tool_call) - estimate_prompt_tokens(base) == 2 + + +def test_estimate_prompt_tokens_counts_reasoning_metadata() -> None: + base = [ + Message(role=Role.USER, content="hi"), + Message(role=Role.ASSISTANT, content=None), + ] + with_reasoning = [ + Message(role=Role.USER, content="hi"), + Message( + role=Role.ASSISTANT, + content=None, + metadata={"reasoning_content": "abcdefgh"}, + ), + ] + + assert estimate_prompt_tokens(with_reasoning) - estimate_prompt_tokens(base) == 2 + + +def test_estimate_prompt_tokens_counts_tool_result_ids() -> None: + messages = [ + Message(role=Role.USER, content="hi"), + Message(role=Role.TOOL, content="ok", tool_call_id="abcdefgh"), + ] + + assert estimate_prompt_tokens(messages) == 11