diff --git a/tests/test_portfolio_context.py b/tests/test_portfolio_context.py index 4b89637d2..32e0ede6d 100644 --- a/tests/test_portfolio_context.py +++ b/tests/test_portfolio_context.py @@ -45,6 +45,20 @@ def test_a_ticker_held_under_another_spelling_is_matched(): assert "No current position" not in text and "5" in text +@pytest.mark.unit +@pytest.mark.parametrize("held, analyzed, canonical", [ + ("BTCUSDT", "BTC-USD", "BTC-USD"), + ("XAUUSD+", "GC=F", "GC=F"), + ("700.HK", "0700.HK", "0700.HK"), + ("0700.HK", "700.HK", "0700.HK"), +]) +def test_a_holding_under_a_broker_alias_is_the_analyzed_instrument(held, analyzed, canonical): + book = PortfolioContext.model_validate({"positions": [{"ticker": held, "quantity": 2}, {"ticker": "ethusdt", "quantity": 1}]}) + text = book.render(analyzed) + assert f"Current position in {canonical}: 2 units" in text + assert "Other positions: ETHUSDT 1" in text # as listed + + @pytest.mark.unit def test_absent_context_is_reported_as_not_provided(): notice = get_portfolio_context_from_state({"company_of_interest": "AAPL"}) @@ -66,6 +80,22 @@ def test_load_rejects_a_malformed_file_with_a_clear_error(tmp_path): load_portfolio(bad) +@pytest.mark.unit +def test_one_instrument_held_under_two_spellings_shows_both_holdings(): + book = PortfolioContext.model_validate({"positions": [ + {"ticker": "BTCUSDT", "quantity": 1}, {"ticker": "BTC-USD", "quantity": 2}]}) + text = book.render("BTC-USD") + assert "Current position in BTC-USD: 1 units (listed as BTCUSDT)" in text + assert "Current position in BTC-USD: 2 units" in text and "Other positions" not in text + + +@pytest.mark.unit +def test_a_holding_keeps_the_instrument_it_was_listed_as(): + # WTI is also an alias of crude oil futures; listed beside AAPL it is the stock as written. + text = PortfolioContext.model_validate({"positions": [{"ticker": "WTI", "quantity": 100}]}).render("AAPL") + assert "Other positions: WTI 100" in text and "CL=F" not in text + + @pytest.mark.unit def test_load_reads_a_valid_file(tmp_path): good = tmp_path / "p.json" diff --git a/tradingagents/portfolio.py b/tradingagents/portfolio.py index c4f42392e..1040e8fb1 100644 --- a/tradingagents/portfolio.py +++ b/tradingagents/portfolio.py @@ -19,6 +19,8 @@ from pathlib import Path from pydantic import BaseModel, Field, ValidationError +from tradingagents.dataflows.symbols import normalize_symbol + class Position(BaseModel): ticker: str = Field(description="Instrument symbol, e.g. AAPL") @@ -31,21 +33,32 @@ class PortfolioContext(BaseModel): currency: str | None = Field(default=None, description="Currency label for cash and prices") positions: list[Position] = Field(default_factory=list) + def positions_in(self, ticker: str) -> list[Position]: + """The holdings of the instrument the run prices for ``ticker``, under any + broker spelling (BTCUSDT, 700.HK). Holdings keep the spelling they were + given: a symbol can also name another instrument (WTI the stock).""" + symbol = normalize_symbol(ticker) + return [p for p in self.positions if normalize_symbol(p.ticker) == symbol] + def position_in(self, ticker: str) -> Position | None: - return next((p for p in self.positions if p.ticker.upper() == ticker.strip().upper()), None) + held = self.positions_in(ticker) + return held[0] if held else 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}"] + symbol = normalize_symbol(ticker) + held = self.positions_in(symbol) + lines = [] + for position in held: + price = f", average price {position.average_price:,.2f}" if position.average_price is not None else "" + spelled = position.ticker.strip().upper() + listed = f" (listed as {position.ticker})" if spelled != symbol else "" + lines.append(f"- Current position in {symbol}: {position.quantity:,.4g} units{price}{listed}") + if not held: + lines.append(f"- No current position in {symbol}") 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] + others = [p for p in self.positions if all(p is not h for h in 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)