diff --git a/tests/test_tool_date_enforcement.py b/tests/test_tool_date_enforcement.py index 55355823d..8beb479d3 100644 --- a/tests/test_tool_date_enforcement.py +++ b/tests/test_tool_date_enforcement.py @@ -141,6 +141,14 @@ def test_an_unpadded_date_reaches_the_vendor_in_canonical_form(): assert args[-2:] == ("2026-07-01", "2026-07-03") +@pytest.mark.unit +@pytest.mark.parametrize("tool", [tools.get_balance_sheet, tools.get_cashflow, tools.get_income_statement]) +def test_a_statement_frequency_is_one_the_vendors_serve(tool): + # Any other word was served annual figures under the word asked for. + schema = tool.tool_call_schema.model_json_schema()["properties"]["freq"] + assert schema["enum"] == ["annual", "quarterly"] + + @pytest.mark.unit def test_direct_call_without_state_is_unchanged(): with mock.patch.object(tools, "route_to_vendor", return_value="ok") as routed: diff --git a/tradingagents/agents/tools.py b/tradingagents/agents/tools.py index c0e9adf74..e2a7eb362 100644 --- a/tradingagents/agents/tools.py +++ b/tradingagents/agents/tools.py @@ -4,7 +4,7 @@ Each dated tool takes the run's ``trade_date`` from graph state (``InjectedState and never serves data past it, whatever date the model asks for. """ -from typing import Annotated +from typing import Annotated, Literal from langchain_core.tools import tool from langgraph.prebuilt import InjectedState @@ -114,7 +114,7 @@ def get_fundamentals( @tool def get_balance_sheet( ticker: Annotated[str, InjectedState("company_of_interest")], - freq: Annotated[str, "reporting frequency: annual/quarterly"] = "quarterly", + freq: Annotated[Literal["annual", "quarterly"], "reporting frequency"] = "quarterly", curr_date: Annotated[str, "current date you are trading at, yyyy-mm-dd"] = None, trade_date: Annotated[str, InjectedState("trade_date")] = "", ) -> str: @@ -133,7 +133,7 @@ def get_balance_sheet( @tool def get_cashflow( ticker: Annotated[str, InjectedState("company_of_interest")], - freq: Annotated[str, "reporting frequency: annual/quarterly"] = "quarterly", + freq: Annotated[Literal["annual", "quarterly"], "reporting frequency"] = "quarterly", curr_date: Annotated[str, "current date you are trading at, yyyy-mm-dd"] = None, trade_date: Annotated[str, InjectedState("trade_date")] = "", ) -> str: @@ -152,7 +152,7 @@ def get_cashflow( @tool def get_income_statement( ticker: Annotated[str, InjectedState("company_of_interest")], - freq: Annotated[str, "reporting frequency: annual/quarterly"] = "quarterly", + freq: Annotated[Literal["annual", "quarterly"], "reporting frequency"] = "quarterly", curr_date: Annotated[str, "current date you are trading at, yyyy-mm-dd"] = None, trade_date: Annotated[str, InjectedState("trade_date")] = "", ) -> str: