diff --git a/cli/main.py b/cli/main.py index 9f4549e3..1647ea01 100644 --- a/cli/main.py +++ b/cli/main.py @@ -27,6 +27,7 @@ import time from rich import box from rich.align import Align from rich.rule import Rule +from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn from tradingagents.graph.trading_graph import TradingAgentsGraph from tradingagents.report_paths import get_daily_dir, get_market_dir, get_ticker_dir @@ -1587,10 +1588,8 @@ def run_pipeline( f" [dim]{c.sector} · {c.conviction.upper()} conviction[/dim]" ) console.print() - import time as _time - from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn - pipeline_start = _time.monotonic() + pipeline_start = time.monotonic() with Progress( SpinnerColumn(), @@ -1604,17 +1603,17 @@ def run_pipeline( overall = progress.add_task("[bold]Pipeline progress[/bold]", total=len(candidates)) def on_done(result, done_count, total_count): - ticker_elapsed = _time.monotonic() - pipeline_start + ticker_elapsed = result.elapsed_seconds if result.error: console.print( f" [red]✗ {result.ticker}[/red]" - f" [dim]failed ({ticker_elapsed:.0f}s elapsed) — {result.error[:80]}[/dim]" + f" [dim]failed ({ticker_elapsed:.0f}s) — {result.error[:80]}[/dim]" ) else: decision_preview = str(result.final_trade_decision)[:70].replace("\n", " ") console.print( f" [green]✓ {result.ticker}[/green]" - f" [dim]({done_count}/{total_count}, {ticker_elapsed:.0f}s elapsed)[/dim]" + f" [dim]({done_count}/{total_count}, {ticker_elapsed:.0f}s)[/dim]" f" → {decision_preview}" ) progress.advance(overall) @@ -1630,7 +1629,7 @@ def run_pipeline( console.print(f"[red]Pipeline failed: {e}[/red]") raise typer.Exit(1) - elapsed_total = _time.monotonic() - pipeline_start + elapsed_total = time.monotonic() - pipeline_start console.print( f"\n[bold green]All {len(candidates)} ticker(s) finished in {elapsed_total:.0f}s[/bold green]\n" ) diff --git a/tests/cli/test_stats_handler.py b/tests/cli/test_stats_handler.py index 8bad63fb..c2d4d7f8 100644 --- a/tests/cli/test_stats_handler.py +++ b/tests/cli/test_stats_handler.py @@ -1,8 +1,9 @@ import threading import pytest from cli.stats_handler import StatsCallbackHandler -from langchain_core.outputs import LLMResult, Generation +from langchain_core.outputs import LLMResult, Generation, ChatGeneration from langchain_core.messages import AIMessage +from langchain_core.messages.ai import UsageMetadata def test_stats_handler_initial_state(): handler = StatsCallbackHandler() @@ -35,11 +36,10 @@ def test_stats_handler_on_tool_start(): def test_stats_handler_on_llm_end_with_usage(): handler = StatsCallbackHandler() - # Mock usage metadata - usage_metadata = {"input_tokens": 10, "output_tokens": 20} - message = AIMessage(content="test response") - message.usage_metadata = usage_metadata - generation = Generation(message=message, text="test response") + # ChatGeneration wraps chat messages; Generation (plain text) has no .message attr. + usage_metadata = UsageMetadata(input_tokens=10, output_tokens=20, total_tokens=30) + message = AIMessage(content="test response", usage_metadata=usage_metadata) + generation = ChatGeneration(message=message) response = LLMResult(generations=[[generation]]) handler.on_llm_end(response) @@ -83,11 +83,10 @@ def test_stats_handler_thread_safety(): handler.on_llm_start({}, []) handler.on_tool_start({}, "") - # Mock usage metadata for on_llm_end - usage_metadata = {"input_tokens": 1, "output_tokens": 1} - message = AIMessage(content="x") - message.usage_metadata = usage_metadata - generation = Generation(message=message, text="x") + # ChatGeneration wraps chat messages with usage_metadata + usage_metadata = UsageMetadata(input_tokens=1, output_tokens=1, total_tokens=2) + message = AIMessage(content="x", usage_metadata=usage_metadata) + generation = ChatGeneration(message=message) response = LLMResult(generations=[[generation]]) handler.on_llm_end(response) diff --git a/tradingagents/dataflows/y_finance.py b/tradingagents/dataflows/y_finance.py index c837594a..a49cd12d 100644 --- a/tradingagents/dataflows/y_finance.py +++ b/tradingagents/dataflows/y_finance.py @@ -168,7 +168,7 @@ def get_stock_stats_indicators_window( ind_string += f"{date_str}: {value}\n" except Exception as e: - print(f"Error getting bulk stockstats data: {e}") + logger.warning("Bulk stockstats failed for %s/%s, falling back to per-day loop: %s", symbol, indicator, e) # Fallback to original implementation if bulk method fails ind_string = "" curr_date_dt = datetime.strptime(curr_date, "%Y-%m-%d") diff --git a/tradingagents/pipeline/macro_bridge.py b/tradingagents/pipeline/macro_bridge.py index e18c6ef9..637c1256 100644 --- a/tradingagents/pipeline/macro_bridge.py +++ b/tradingagents/pipeline/macro_bridge.py @@ -70,6 +70,7 @@ class TickerResult: final_trade_decision: str = "" error: str | None = None + elapsed_seconds: float = 0.0 # ─── Parsing ────────────────────────────────────────────────────────────────── @@ -207,6 +208,7 @@ def run_ticker_analysis( result.final_trade_decision = decision elapsed = time.monotonic() - t0 + result.elapsed_seconds = elapsed logger.info( "[%s] ✓ Analysis complete in %.0fs — decision: %s", candidate.ticker, elapsed, str(decision)[:80], @@ -214,6 +216,7 @@ def run_ticker_analysis( except Exception as exc: elapsed = time.monotonic() - t0 + result.elapsed_seconds = elapsed logger.error( "[%s] ✗ Analysis FAILED after %.0fs: %s", candidate.ticker, elapsed, exc, exc_info=True,