feat(graph): accept the caller's portfolio as run input (#1304, #1166)

- PortfolioContext input, rendered once into state and read by the trader, risk and portfolio agents
- --portfolio JSON file on the CLI
- a run without it is never treated as a flat book
- the checkpoint signature keys on the portfolio
This commit is contained in:
Yijia-Xiao
2026-09-16 20:19:49 +00:00
parent dffff22951
commit 6436d1ff30
14 changed files with 391 additions and 21 deletions

View File

@@ -1001,7 +1001,7 @@ def _build_run_config(selections: dict, checkpoint: bool | None) -> dict:
return config return config
def run_analysis(checkpoint: bool | None = None): def run_analysis(checkpoint: bool | None = None, portfolio=None):
# First get all user selections # First get all user selections
selections = get_user_selections() selections = get_user_selections()
@@ -1113,7 +1113,7 @@ def run_analysis(checkpoint: bool | None = None):
# The same initial state propagate() builds: settled decision log, past # The same initial state propagate() builds: settled decision log, past
# context and resolved instrument identity. # context and resolved instrument identity.
init_agent_state = graph.create_run_state( init_agent_state = graph.create_run_state(
selections["ticker"], selections["analysis_date"], selections["asset_type"] selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
) )
# Pass callbacks to graph config for tool execution tracking # Pass callbacks to graph config for tool execution tracking
# (LLM tracking is handled separately via LLM constructor) # (LLM tracking is handled separately via LLM constructor)
@@ -1123,7 +1123,7 @@ def run_analysis(checkpoint: bool | None = None):
# actually saves and resumes on the CLI path (#1249); a no-op when # actually saves and resumes on the CLI path (#1249); a no-op when
# checkpointing is disabled. Torn down in the finally below. # checkpointing is disabled. Torn down in the finally below.
checkpoint_tid = graph.begin_checkpoint( checkpoint_tid = graph.begin_checkpoint(
selections["ticker"], selections["analysis_date"], selections["asset_type"] selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
) )
if checkpoint_tid is not None: if checkpoint_tid is not None:
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
@@ -1246,7 +1246,7 @@ def run_analysis(checkpoint: bool | None = None):
# the checkpoint for resume. # the checkpoint for resume.
graph.record_decision(selections["ticker"], selections["analysis_date"], final_state) graph.record_decision(selections["ticker"], selections["analysis_date"], final_state)
graph.clear_checkpoint_on_success( graph.clear_checkpoint_on_success(
selections["ticker"], selections["analysis_date"], selections["asset_type"] selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
) )
finally: finally:
# Always restore the plain uncheckpointed graph, even on failure. # Always restore the plain uncheckpointed graph, even on failure.
@@ -1308,13 +1308,28 @@ def analyze(
"--clear-checkpoints", "--clear-checkpoints",
help="Delete all saved checkpoints before running (force fresh start).", help="Delete all saved checkpoints before running (force fresh start).",
), ),
portfolio: str = typer.Option(
None,
"--portfolio",
help="JSON file with current holdings and cash, so the trader, risk and "
"portfolio agents size against your actual position.",
),
): ):
if clear_checkpoints: if clear_checkpoints:
from tradingagents.graph.checkpointer import clear_all_checkpoints from tradingagents.graph.checkpointer import clear_all_checkpoints
n = clear_all_checkpoints(DEFAULT_CONFIG["data_cache_dir"]) n = clear_all_checkpoints(DEFAULT_CONFIG["data_cache_dir"])
console.print(f"[yellow]Cleared {n} checkpoint(s).[/yellow]") console.print(f"[yellow]Cleared {n} checkpoint(s).[/yellow]")
portfolio_context = None
if portfolio:
from tradingagents.portfolio import load_portfolio
try: try:
run_analysis(checkpoint=checkpoint) portfolio_context = load_portfolio(portfolio)
except ValueError as exc:
console.print(f"[red]{exc}[/red]")
raise typer.Exit(code=1) from None
try:
run_analysis(checkpoint=checkpoint, portfolio=portfolio_context)
except _NO_CONSOLE_ERRORS: except _NO_CONSOLE_ERRORS:
# A terminal with no console buffer cannot host the interactive prompts. # A terminal with no console buffer cannot host the interactive prompts.
# Emit one actionable line on stderr instead of a prompt_toolkit # Emit one actionable line on stderr instead of a prompt_toolkit

View File

@@ -69,7 +69,7 @@ class _FakeGraph:
self.graph = self self.graph = self
self.propagator = self self.propagator = self
def create_run_state(self, ticker, trade_date, asset_type="stock"): def create_run_state(self, ticker, trade_date, asset_type="stock", portfolio=None):
self.calls.append(("create_run_state", ticker, trade_date)) self.calls.append(("create_run_state", ticker, trade_date))
return {"messages": [], "company_of_interest": ticker} return {"messages": [], "company_of_interest": ticker}

View File

@@ -0,0 +1,240 @@
"""Portfolio context: what the caller holds, threaded into the decision agents.
Decisions were made with no knowledge of the current book, so "add to a full
position" and "open a new one" read alike. The context is optional and carries
three distinct states: a position, a flat book, and no context at all. Nothing
may present the third as the second. The research team stays blind so the bull
and bear cases are not anchored by the caller's position.
"""
from __future__ import annotations
import json
import pytest
from tradingagents.agents.utils.agent_utils import get_portfolio_context_from_state
from tradingagents.portfolio import PortfolioContext, load_portfolio
HOLDING = {
"cash": 25000.0,
"currency": "USD",
"positions": [
{"ticker": "AAPL", "quantity": 120, "average_price": 150.0},
{"ticker": "MSFT", "quantity": 10},
],
}
@pytest.mark.unit
def test_position_in_the_analyzed_instrument_leads_the_render():
text = PortfolioContext.model_validate(HOLDING).render("AAPL")
assert "120" in text and "150" in text
assert "MSFT" in text and "25,000" in text and "USD" in text
@pytest.mark.unit
def test_flat_book_says_no_position_rather_than_omitting_it():
text = PortfolioContext.model_validate({"cash": 1000.0, "positions": []}).render("AAPL")
assert "No current position in AAPL" in text
@pytest.mark.unit
def test_a_ticker_held_under_another_spelling_is_matched():
text = PortfolioContext.model_validate({"positions": [{"ticker": "aapl", "quantity": 5}]}).render("AAPL")
assert "No current position" not in text and "5" in text
@pytest.mark.unit
def test_absent_context_is_reported_as_not_provided():
notice = get_portfolio_context_from_state({"company_of_interest": "AAPL"})
assert "not provided" in notice.lower()
assert "no position" not in notice.lower() # missing must not read as flat
@pytest.mark.unit
def test_rendered_context_reaches_the_agents_from_state():
block = get_portfolio_context_from_state({"portfolio_context": "Portfolio: flat", "company_of_interest": "AAPL"})
assert block == "Portfolio: flat"
@pytest.mark.unit
def test_load_rejects_a_malformed_file_with_a_clear_error(tmp_path):
bad = tmp_path / "p.json"
bad.write_text(json.dumps({"positions": [{"quantity": 5}]}))
with pytest.raises(ValueError, match="portfolio"):
load_portfolio(bad)
@pytest.mark.unit
def test_load_reads_a_valid_file(tmp_path):
good = tmp_path / "p.json"
good.write_text(json.dumps(HOLDING))
assert load_portfolio(good).positions[0].ticker == "AAPL"
# --- threading through the graph --------------------------------------------
def _bare_graph(tmp_path):
from tradingagents.agents.utils.memory import TradingMemoryLog
from tradingagents.graph.propagation import Propagator
from tradingagents.graph.trading_graph import TradingAgentsGraph
graph = object.__new__(TradingAgentsGraph)
graph.config = {"memory_log_path": str(tmp_path / "m.md"), "max_debate_rounds": 1,
"max_risk_discuss_rounds": 1}
graph.memory_log = TradingMemoryLog(graph.config)
graph.propagator = Propagator()
graph.selected_analysts = ["market"]
graph._resolve_pending_entries = lambda t: None
graph.resolve_instrument_context = lambda t, a="stock": ""
graph._memory_as_of = lambda d: None
return graph
@pytest.mark.unit
def test_create_run_state_renders_the_portfolio_once(tmp_path):
graph = _bare_graph(tmp_path)
state = graph.create_run_state("AAPL", "2026-08-14", portfolio=PortfolioContext.model_validate(HOLDING))
assert "120" in state["portfolio_context"]
assert graph.create_run_state("AAPL", "2026-08-14")["portfolio_context"] == ""
@pytest.mark.unit
def test_checkpoint_signature_changes_with_the_portfolio(tmp_path):
graph = _bare_graph(tmp_path)
none = graph._run_signature("stock")
flat = graph._run_signature("stock", PortfolioContext())
held = graph._run_signature("stock", PortfolioContext.model_validate(HOLDING))
assert len({none, flat, held}) == 3
@pytest.mark.unit
@pytest.mark.parametrize("module, factory", [
("tradingagents.agents.trader.trader", "create_trader"),
("tradingagents.agents.managers.portfolio_manager", "create_portfolio_manager"),
("tradingagents.agents.risk_mgmt.aggressive_debator", "create_aggressive_debator"),
("tradingagents.agents.risk_mgmt.conservative_debator", "create_conservative_debator"),
("tradingagents.agents.risk_mgmt.neutral_debator", "create_neutral_debator"),
])
def test_decision_agents_see_the_portfolio(module, factory, monkeypatch):
"""The prompt each decision agent sends carries the portfolio block."""
import importlib
mod = importlib.import_module(module)
seen = []
class _LLM:
def invoke(self, prompt, *a, **k):
seen.append(prompt if isinstance(prompt, str) else json.dumps(str(prompt)))
from langchain_core.messages import AIMessage
return AIMessage("Rating: Hold\n\nnothing to do")
def with_structured_output(self, *a, **k):
raise NotImplementedError # force the free-text path
state = {
"company_of_interest": "AAPL", "trade_date": "2026-08-14", "asset_type": "stock",
"instrument_context": "", "market_report": "M", "sentiment_report": "S",
"news_report": "N", "fundamentals_report": "F", "investment_plan": "P",
"trader_investment_plan": "T", "past_context": "",
"portfolio_context": "PORTFOLIO_BLOCK_MARKER",
"investment_debate_state": {"history": "", "judge_decision": "", "count": 0},
"risk_debate_state": {"history": "", "latest_speaker": "", "count": 0,
"aggressive_history": "", "conservative_history": "", "neutral_history": "",
"current_aggressive_response": "", "current_conservative_response": "",
"current_neutral_response": "", "judge_decision": ""},
}
node = getattr(mod, factory)(_LLM())
node(state)
assert any("PORTFOLIO_BLOCK_MARKER" in p for p in seen), f"{factory} prompt lacks the portfolio block"
@pytest.mark.unit
def test_research_team_stays_blind_to_the_portfolio():
import inspect
from tradingagents.agents.researchers import bear_researcher, bull_researcher
for mod in (bull_researcher, bear_researcher):
assert "portfolio_context" not in inspect.getsource(mod)
@pytest.mark.unit
def test_completed_run_clears_the_checkpoint_it_wrote(tmp_path, monkeypatch):
"""The clear must key on the same portfolio the run was checkpointed under.
Keyed on a different one it deletes nothing, and the next identical call
resumes the finished thread and returns the old decision without running.
"""
import tradingagents.graph.trading_graph as tg
graph = _bare_graph(tmp_path)
graph.config.update({"checkpoint_enabled": True, "data_cache_dir": str(tmp_path),
"results_dir": str(tmp_path)})
graph.debug = False
graph._resuming = False
graph.propagator.get_graph_args = lambda callbacks=None: {}
graph.process_signal = lambda d: "Hold"
graph._log_state = lambda *a, **k: None
graph.graph = type("G", (), {"invoke": lambda self, i, **k: {"final_trade_decision": "Rating: Hold\n\nx"}})()
book = PortfolioContext.model_validate(HOLDING)
written = graph._run_signature("stock", book) # what begin_checkpoint keys on
cleared = []
monkeypatch.setattr(tg, "clear_checkpoint", lambda d, t, dt, signature: cleared.append(signature))
graph._run_graph("AAPL", "2026-08-14", "stock", checkpoint_thread_id=None, portfolio=book)
assert cleared == [written]
@pytest.mark.unit
def test_research_layer_sizes_against_a_standard_allocation():
"""The research team is blind to the book, so its plan cannot promise
position-relative sizing: it sizes against a standard allocation instead."""
from tradingagents.agents.schemas import ResearchPlan
description = ResearchPlan.model_fields["strategic_actions"].description
assert "standard allocation" in description
assert "does not see the caller's holdings" in description
@pytest.mark.unit
def test_partial_portfolio_states_only_what_it_was_given():
"""Cash omitted is not cash zero; the line is absent rather than invented."""
text = PortfolioContext.model_validate({"positions": [{"ticker": "AAPL", "quantity": 5}]}).render("AAPL")
assert "Cash" not in text
assert "5" in text
@pytest.mark.unit
def test_cli_rejects_an_unusable_portfolio_file_before_running(tmp_path, monkeypatch):
from typer.testing import CliRunner
import cli.main as m
bad = tmp_path / "bad.json"
bad.write_text('{"positions": [{"quantity": 5}]}')
ran = []
monkeypatch.setattr(m, "run_analysis", lambda **k: ran.append(k))
result = CliRunner().invoke(m.app, ["--portfolio", str(bad)])
assert result.exit_code == 1 and ran == []
@pytest.mark.unit
def test_cli_passes_a_valid_portfolio_into_the_run(tmp_path, monkeypatch):
from typer.testing import CliRunner
import cli.main as m
good = tmp_path / "good.json"
good.write_text(json.dumps(HOLDING))
ran = []
monkeypatch.setattr(m, "run_analysis", lambda **k: ran.append(k))
result = CliRunner().invoke(m.app, ["--portfolio", str(good)])
assert result.exit_code == 0
assert ran[0]["portfolio"].position_in("AAPL").quantity == 120

View File

@@ -14,6 +14,7 @@ from tradingagents.agents.schemas import PortfolioDecision, render_pm_decision
from tradingagents.agents.utils.agent_utils import ( from tradingagents.agents.utils.agent_utils import (
get_instrument_context_from_state, get_instrument_context_from_state,
get_language_instruction, get_language_instruction,
get_portfolio_context_from_state,
) )
from tradingagents.agents.utils.structured import ( from tradingagents.agents.utils.structured import (
NO_EXTERNAL_TOOLS, NO_EXTERNAL_TOOLS,
@@ -27,6 +28,7 @@ def create_portfolio_manager(llm):
def portfolio_manager_node(state) -> dict: def portfolio_manager_node(state) -> dict:
instrument_context = get_instrument_context_from_state(state) instrument_context = get_instrument_context_from_state(state)
portfolio_context = get_portfolio_context_from_state(state)
history = state["risk_debate_state"]["history"] history = state["risk_debate_state"]["history"]
risk_debate_state = state["risk_debate_state"] risk_debate_state = state["risk_debate_state"]
@@ -44,6 +46,8 @@ def create_portfolio_manager(llm):
{instrument_context} {instrument_context}
{portfolio_context}
--- ---
**Rating Scale** (use exactly one): **Rating Scale** (use exactly one):

View File

@@ -1,6 +1,7 @@
from tradingagents.agents.utils.agent_utils import ( from tradingagents.agents.utils.agent_utils import (
get_instrument_context_from_state, get_instrument_context_from_state,
get_language_instruction, get_language_instruction,
get_portfolio_context_from_state,
opponent_argument_or_opening, opponent_argument_or_opening,
) )
@@ -23,6 +24,7 @@ def create_aggressive_debator(llm):
news_report = state["news_report"] news_report = state["news_report"]
fundamentals_report = state["fundamentals_report"] fundamentals_report = state["fundamentals_report"]
instrument_context = get_instrument_context_from_state(state) instrument_context = get_instrument_context_from_state(state)
portfolio_context = get_portfolio_context_from_state(state)
trader_decision = state["trader_investment_plan"] trader_decision = state["trader_investment_plan"]
@@ -33,6 +35,7 @@ def create_aggressive_debator(llm):
Your task is to create a compelling case for the trader's decision by questioning and critiquing the conservative and neutral stances to demonstrate why your high-reward perspective offers the best path forward. Incorporate insights from the following sources into your arguments: Your task is to create a compelling case for the trader's decision by questioning and critiquing the conservative and neutral stances to demonstrate why your high-reward perspective offers the best path forward. Incorporate insights from the following sources into your arguments:
{instrument_context} {instrument_context}
{portfolio_context}
Market Research Report: {market_research_report} Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report} Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report} Latest World Affairs Report: {news_report}

View File

@@ -1,6 +1,7 @@
from tradingagents.agents.utils.agent_utils import ( from tradingagents.agents.utils.agent_utils import (
get_instrument_context_from_state, get_instrument_context_from_state,
get_language_instruction, get_language_instruction,
get_portfolio_context_from_state,
opponent_argument_or_opening, opponent_argument_or_opening,
) )
@@ -23,6 +24,7 @@ def create_conservative_debator(llm):
news_report = state["news_report"] news_report = state["news_report"]
fundamentals_report = state["fundamentals_report"] fundamentals_report = state["fundamentals_report"]
instrument_context = get_instrument_context_from_state(state) instrument_context = get_instrument_context_from_state(state)
portfolio_context = get_portfolio_context_from_state(state)
trader_decision = state["trader_investment_plan"] trader_decision = state["trader_investment_plan"]
@@ -33,6 +35,7 @@ def create_conservative_debator(llm):
Your task is to actively counter the arguments of the Aggressive and Neutral Analysts, highlighting where their views may overlook potential threats or fail to prioritize sustainability. Respond directly to their points, drawing from the following data sources to build a convincing case for a low-risk approach adjustment to the trader's decision: Your task is to actively counter the arguments of the Aggressive and Neutral Analysts, highlighting where their views may overlook potential threats or fail to prioritize sustainability. Respond directly to their points, drawing from the following data sources to build a convincing case for a low-risk approach adjustment to the trader's decision:
{instrument_context} {instrument_context}
{portfolio_context}
Market Research Report: {market_research_report} Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report} Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report} Latest World Affairs Report: {news_report}

View File

@@ -1,6 +1,7 @@
from tradingagents.agents.utils.agent_utils import ( from tradingagents.agents.utils.agent_utils import (
get_instrument_context_from_state, get_instrument_context_from_state,
get_language_instruction, get_language_instruction,
get_portfolio_context_from_state,
opponent_argument_or_opening, opponent_argument_or_opening,
) )
@@ -23,6 +24,7 @@ def create_neutral_debator(llm):
news_report = state["news_report"] news_report = state["news_report"]
fundamentals_report = state["fundamentals_report"] fundamentals_report = state["fundamentals_report"]
instrument_context = get_instrument_context_from_state(state) instrument_context = get_instrument_context_from_state(state)
portfolio_context = get_portfolio_context_from_state(state)
trader_decision = state["trader_investment_plan"] trader_decision = state["trader_investment_plan"]
@@ -33,6 +35,7 @@ def create_neutral_debator(llm):
Your task is to challenge both the Aggressive and Conservative Analysts, pointing out where each perspective may be overly optimistic or overly cautious. Use insights from the following data sources to support a moderate, sustainable strategy to adjust the trader's decision: Your task is to challenge both the Aggressive and Conservative Analysts, pointing out where each perspective may be overly optimistic or overly cautious. Use insights from the following data sources to support a moderate, sustainable strategy to adjust the trader's decision:
{instrument_context} {instrument_context}
{portfolio_context}
Market Research Report: {market_research_report} Market Research Report: {market_research_report}
Social Media Sentiment Report: {sentiment_report} Social Media Sentiment Report: {sentiment_report}
Latest World Affairs Report: {news_report} Latest World Affairs Report: {news_report}

View File

@@ -112,7 +112,9 @@ class ResearchPlan(BaseModel):
strategic_actions: str = Field( strategic_actions: str = Field(
description=( description=(
"Concrete steps for the trader to implement the recommendation, " "Concrete steps for the trader to implement the recommendation, "
"including position sizing guidance consistent with the rating." "including sizing guidance relative to a standard allocation. The "
"research team does not see the caller's holdings; the trader and "
"portfolio manager apply the actual position."
), ),
) )

View File

@@ -10,6 +10,7 @@ from tradingagents.agents.schemas import TraderProposal, render_trader_proposal
from tradingagents.agents.utils.agent_utils import ( from tradingagents.agents.utils.agent_utils import (
get_instrument_context_from_state, get_instrument_context_from_state,
get_language_instruction, get_language_instruction,
get_portfolio_context_from_state,
) )
from tradingagents.agents.utils.structured import ( from tradingagents.agents.utils.structured import (
NO_EXTERNAL_TOOLS, NO_EXTERNAL_TOOLS,
@@ -31,6 +32,7 @@ def create_trader(llm):
# report is empty when the user did not select the market analyst, so # report is empty when the user did not select the market analyst, so
# only offer it (and the grounding instruction) when it has content. # only offer it (and the grounding instruction) when it has content.
market_report = (state["market_report"] or "").strip() market_report = (state["market_report"] or "").strip()
portfolio_context = get_portfolio_context_from_state(state)
if market_report: if market_report:
grounding = ( grounding = (
@@ -67,6 +69,7 @@ def create_trader(llm):
f"Here is the research team's investment plan for {company_name}. " f"Here is the research team's investment plan for {company_name}. "
f"{instrument_context}\n\n" f"{instrument_context}\n\n"
f"{report_section}" f"{report_section}"
f"{portfolio_context}\n\n"
f"Proposed Investment Plan:\n{investment_plan}\n\n" f"Proposed Investment Plan:\n{investment_plan}\n\n"
f"Make an informed, strategic trading decision." f"Make an informed, strategic trading decision."
), ),

View File

@@ -74,3 +74,4 @@ class AgentState(MessagesState):
] ]
final_trade_decision: Annotated[str, "Final decision made by the Risk Analysts"] final_trade_decision: Annotated[str, "Final decision made by the Risk Analysts"]
past_context: Annotated[str, "Memory log context injected at run start (same-ticker decisions + cross-ticker lessons)"] past_context: Annotated[str, "Memory log context injected at run start (same-ticker decisions + cross-ticker lessons)"]
portfolio_context: Annotated[str, "Caller-supplied holdings and cash, rendered at run start; empty when not provided"]

View File

@@ -201,6 +201,23 @@ def get_instrument_context_from_state(state: Mapping[str, Any]) -> str:
) )
def get_portfolio_context_from_state(state: Mapping[str, Any]) -> str:
"""Return the caller's portfolio block, or a notice that none was given.
A run without portfolio context must not read as a flat book: the agents
would otherwise size as if the caller held nothing, which is a claim about
an account we were never told about.
"""
context = state.get("portfolio_context")
if isinstance(context, str) and context.strip():
return context
return (
"Portfolio context: not provided. You do not know the caller's current "
"holdings or cash, so do not assume a flat book; give direction and "
"sizing guidance in terms the caller can apply to their own position."
)
def create_msg_delete(): def create_msg_delete():
def delete_messages(state): def delete_messages(state):
"""Clear messages and add a context-anchored placeholder. """Clear messages and add a context-anchored placeholder.

View File

@@ -22,6 +22,7 @@ class Propagator:
asset_type: str = "stock", asset_type: str = "stock",
past_context: str = "", past_context: str = "",
instrument_context: str = "", instrument_context: str = "",
portfolio_context: str = "",
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Create the initial state for the agent graph. """Create the initial state for the agent graph.
@@ -38,6 +39,7 @@ class Propagator:
"instrument_context": instrument_context, "instrument_context": instrument_context,
"trade_date": str(trade_date), "trade_date": str(trade_date),
"past_context": past_context, "past_context": past_context,
"portfolio_context": portfolio_context,
"investment_debate_state": InvestDebateState( "investment_debate_state": InvestDebateState(
{ {
"bull_history": "", "bull_history": "",

View File

@@ -403,7 +403,7 @@ class TradingAgentsGraph:
td = str(trade_date) td = str(trade_date)
return td if td < datetime.now().strftime("%Y-%m-%d") else None return td if td < datetime.now().strftime("%Y-%m-%d") else None
def _run_signature(self, asset_type: str) -> str: def _run_signature(self, asset_type: str, portfolio=None) -> str:
"""Graph-shape inputs that must invalidate a checkpoint if changed. """Graph-shape inputs that must invalidate a checkpoint if changed.
Keyed into the checkpoint thread ID so a resume under a different analyst Keyed into the checkpoint thread ID so a resume under a different analyst
@@ -415,9 +415,11 @@ class TradingAgentsGraph:
f"debate={self.config['max_debate_rounds']}", f"debate={self.config['max_debate_rounds']}",
f"risk={self.config['max_risk_discuss_rounds']}", f"risk={self.config['max_risk_discuss_rounds']}",
f"asset={asset_type}", f"asset={asset_type}",
# None, an empty book and a changed book are three different runs.
f"portfolio={portfolio.fingerprint() if portfolio is not None else 'none'}",
]) ])
def propagate(self, company_name, trade_date, asset_type: str = "stock"): def propagate(self, company_name, trade_date, asset_type: str = "stock", portfolio=None):
"""Run the trading agents graph for a company on a specific date. """Run the trading agents graph for a company on a specific date.
``asset_type`` selects between the stock pipeline (default) and the ``asset_type`` selects between the stock pipeline (default) and the
@@ -436,13 +438,13 @@ class TradingAgentsGraph:
trade_date = _validate_trade_date(trade_date) trade_date = _validate_trade_date(trade_date)
self.ticker = company_name self.ticker = company_name
with self.checkpoint_scope(company_name, trade_date, asset_type) as thread_id_value: with self.checkpoint_scope(company_name, trade_date, asset_type, portfolio) as thread_id_value:
return self._run_graph( return self._run_graph(
company_name, trade_date, asset_type=asset_type, company_name, trade_date, asset_type=asset_type,
checkpoint_thread_id=thread_id_value, checkpoint_thread_id=thread_id_value, portfolio=portfolio,
) )
def begin_checkpoint(self, company_name, trade_date, asset_type: str = "stock") -> str | None: def begin_checkpoint(self, company_name, trade_date, asset_type: str = "stock", portfolio=None) -> str | None:
"""Recompile the graph with a per-ticker checkpointer and return the """Recompile the graph with a per-ticker checkpointer and return the
``thread_id`` to inject into the stream/invoke ``config`` (or ``None`` ``thread_id`` to inject into the stream/invoke ``config`` (or ``None``
when checkpointing is disabled). when checkpointing is disabled).
@@ -456,7 +458,7 @@ class TradingAgentsGraph:
self._resuming = False self._resuming = False
if not self.config.get("checkpoint_enabled"): if not self.config.get("checkpoint_enabled"):
return None return None
signature = self._run_signature(asset_type) signature = self._run_signature(asset_type, portfolio)
self._checkpointer_ctx = get_checkpointer(self.config["data_cache_dir"], company_name) self._checkpointer_ctx = get_checkpointer(self.config["data_cache_dir"], company_name)
saver = self._checkpointer_ctx.__enter__() saver = self._checkpointer_ctx.__enter__()
self.graph = self.workflow.compile(checkpointer=saver) self.graph = self.workflow.compile(checkpointer=saver)
@@ -490,19 +492,19 @@ class TradingAgentsGraph:
self._resuming = False self._resuming = False
@contextmanager @contextmanager
def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock"): def checkpoint_scope(self, company_name, trade_date, asset_type: str = "stock", portfolio=None):
"""Context-manager form of begin/end_checkpoint for the propagate path.""" """Context-manager form of begin/end_checkpoint for the propagate path."""
try: try:
yield self.begin_checkpoint(company_name, trade_date, asset_type) yield self.begin_checkpoint(company_name, trade_date, asset_type, portfolio)
finally: finally:
self.end_checkpoint() self.end_checkpoint()
def clear_checkpoint_on_success(self, company_name, trade_date, asset_type: str = "stock"): def clear_checkpoint_on_success(self, company_name, trade_date, asset_type: str = "stock", portfolio=None):
"""Drop a completed run's checkpoint so a later run starts fresh (#1249).""" """Drop a completed run's checkpoint so a later run starts fresh (#1249)."""
if self.config.get("checkpoint_enabled"): if self.config.get("checkpoint_enabled"):
clear_checkpoint( clear_checkpoint(
self.config["data_cache_dir"], company_name, str(trade_date), self.config["data_cache_dir"], company_name, str(trade_date),
self._run_signature(asset_type), self._run_signature(asset_type, portfolio),
) )
def save_reports(self, final_state, ticker, save_path=None) -> Path: def save_reports(self, final_state, ticker, save_path=None) -> Path:
@@ -520,7 +522,7 @@ class TradingAgentsGraph:
) )
return write_report_tree(final_state, ticker, save_path) return write_report_tree(final_state, ticker, save_path)
def create_run_state(self, company_name, trade_date, asset_type: str = "stock"): def create_run_state(self, company_name, trade_date, asset_type: str = "stock", portfolio=None):
"""Build a run's initial state; propagate() and the CLI both start here. """Build a run's initial state; propagate() and the CLI both start here.
Settles this ticker's pending decisions first, then injects the lessons Settles this ticker's pending decisions first, then injects the lessons
@@ -537,8 +539,19 @@ class TradingAgentsGraph:
company_name, as_of=self._memory_as_of(trade_date) company_name, as_of=self._memory_as_of(trade_date)
), ),
instrument_context=self.resolve_instrument_context(company_name, asset_type), instrument_context=self.resolve_instrument_context(company_name, asset_type),
portfolio_context=portfolio.render(company_name) if portfolio is not None else "",
) )
def settle_pending(self, company_name):
"""Settle this ticker's decisions whose holding window has now traded.
A run settles the ticker's earlier decisions on its way in, so the most
recent one stays pending until the next run for that ticker. A caller
that is done analyzing a ticker (a backtest sweep, a scheduled job) calls
this to settle it now.
"""
self._resolve_pending_entries(company_name)
def record_decision(self, company_name, trade_date, final_state): def record_decision(self, company_name, trade_date, final_state):
"""Log a finished run's decision for reflection on the next same-ticker run.""" """Log a finished run's decision for reflection on the next same-ticker run."""
decision = final_state.get("final_trade_decision") decision = final_state.get("final_trade_decision")
@@ -550,9 +563,9 @@ class TradingAgentsGraph:
) )
def _run_graph(self, company_name, trade_date, asset_type: str = "stock", def _run_graph(self, company_name, trade_date, asset_type: str = "stock",
checkpoint_thread_id: str | None = None): checkpoint_thread_id: str | None = None, portfolio=None):
"""Execute the graph and write the resulting state to disk and memory log.""" """Execute the graph and write the resulting state to disk and memory log."""
init_agent_state = self.create_run_state(company_name, trade_date, asset_type) init_agent_state = self.create_run_state(company_name, trade_date, asset_type, portfolio)
args = self.propagator.get_graph_args() args = self.propagator.get_graph_args()
# Inject the checkpoint thread_id (from checkpoint_scope) so the same # Inject the checkpoint thread_id (from checkpoint_scope) so the same
@@ -593,7 +606,7 @@ class TradingAgentsGraph:
self.record_decision(company_name, trade_date, final_state) self.record_decision(company_name, trade_date, final_state)
# Clear checkpoint on successful completion to avoid stale state. # Clear checkpoint on successful completion to avoid stale state.
self.clear_checkpoint_on_success(company_name, trade_date, asset_type) self.clear_checkpoint_on_success(company_name, trade_date, asset_type, portfolio)
return final_state, self.process_signal(final_state["final_trade_decision"]) return final_state, self.process_signal(final_state["final_trade_decision"])

View File

@@ -0,0 +1,64 @@
"""The caller's book, as the decision agents see it.
Optional input to a run: what is held, at what average price, and how much cash
is free. Without it the agents cannot tell adding to a full position from
opening a new one. Three states are distinct and must stay so: a position, a
flat book, and no context at all, since treating "not provided" as "flat" would
invent a fact about the caller's account.
Broker-neutral by construction: quantities are generic units and the currency is
whatever label the caller passes, so nothing here implies a venue or an
execution path.
"""
from __future__ import annotations
import hashlib
import json
from pathlib import Path
from pydantic import BaseModel, Field, ValidationError
class Position(BaseModel):
ticker: str = Field(description="Instrument symbol, e.g. AAPL")
quantity: float = Field(description="Signed units held; negative is short")
average_price: float | None = Field(default=None, description="Average entry price per unit")
class PortfolioContext(BaseModel):
cash: float | None = Field(default=None, description="Free cash available")
currency: str | None = Field(default=None, description="Currency label for cash and prices")
positions: list[Position] = Field(default_factory=list)
def position_in(self, ticker: str) -> Position | None:
return next((p for p in self.positions if p.ticker.upper() == ticker.strip().upper()), None)
def render(self, ticker: str) -> str:
"""The portfolio block for the decision agents, led by the analyzed instrument."""
symbol = ticker.strip().upper()
held = self.position_in(symbol)
if held is None:
lines = [f"- No current position in {symbol}"]
else:
price = f", average price {held.average_price:,.2f}" if held.average_price is not None else ""
lines = [f"- Current position in {symbol}: {held.quantity:,.4g} units{price}"]
if self.cash is not None:
lines.append(f"- Cash available: {self.cash:,.2f}{' ' + self.currency if self.currency else ''}")
others = [p for p in self.positions if p is not held]
if others:
lines.append("- Other positions: " + ", ".join(f"{p.ticker.upper()} {p.quantity:,.4g}" for p in others))
return "Portfolio at the analysis date:\n" + "\n".join(lines)
def fingerprint(self) -> str:
"""Stable digest of the book, so a changed one cannot resume a stale run."""
return hashlib.sha256(self.model_dump_json().encode()).hexdigest()[:12]
def load_portfolio(path: str | Path) -> PortfolioContext:
"""Read a portfolio JSON file, failing before the run rather than mid-graph."""
try:
data = json.loads(Path(path).read_text(encoding="utf-8"))
return PortfolioContext.model_validate(data)
except (OSError, json.JSONDecodeError, ValidationError) as exc:
raise ValueError(f"portfolio file {path} is not usable: {exc}") from exc