mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 05:12:26 +00:00
Elliot SluskyandGitHub
420908401c
Closes #608. Make Message.content officially Optional (str | None) with a Message.text accessor that treats None as empty, and count tool-call IDs/names/arguments, tool-result IDs, and reasoning/thinking metadata in estimate_prompt_tokens — all are replayed into later prompt turns, so they belong in the estimate. Extends estimator and message-type regression tests.
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from openjarvis.core.types import Message, Role, ToolCall
|
|
from openjarvis.engine._base import estimate_prompt_tokens
|
|
|
|
|
|
def test_estimate_prompt_tokens_handles_none_content_tool_call_turn() -> None:
|
|
messages = [
|
|
Message(role=Role.USER, content="hi"),
|
|
Message(
|
|
role=Role.ASSISTANT,
|
|
content=None,
|
|
tool_calls=[ToolCall(id="call_1", name="lookup", arguments="{}")],
|
|
),
|
|
]
|
|
|
|
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
|