131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
"""Tests for robust JSON extraction from LLM output."""
|
|
import pytest
|
|
from tradingagents.agents.utils.json_utils import extract_json
|
|
|
|
|
|
# ─── Happy-path tests ─────────────────────────────────────────────────────────
|
|
|
|
def test_pure_json():
|
|
assert extract_json('{"key": "value"}') == {"key": "value"}
|
|
|
|
|
|
def test_json_with_whitespace():
|
|
assert extract_json(' \n{"key": "value"}\n ') == {"key": "value"}
|
|
|
|
|
|
def test_markdown_fence_json():
|
|
text = '```json\n{"key": "value"}\n```'
|
|
assert extract_json(text) == {"key": "value"}
|
|
|
|
|
|
def test_markdown_fence_no_lang():
|
|
text = '```\n{"key": "value"}\n```'
|
|
assert extract_json(text) == {"key": "value"}
|
|
|
|
|
|
def test_think_preamble_only():
|
|
text = '<think>I need to analyze the macro environment carefully.</think>\n{"key": "value"}'
|
|
assert extract_json(text) == {"key": "value"}
|
|
|
|
|
|
def test_think_plus_fence():
|
|
text = '<think>Some reasoning here.</think>\n```json\n{"key": "value"}\n```'
|
|
assert extract_json(text) == {"key": "value"}
|
|
|
|
|
|
def test_prose_with_json():
|
|
text = 'Here is the result:\n{"key": "value"}\nDone.'
|
|
assert extract_json(text) == {"key": "value"}
|
|
|
|
|
|
def test_nested_json():
|
|
data = {
|
|
"timeframe": "1 month",
|
|
"executive_summary": "Strong growth momentum",
|
|
"macro_context": {
|
|
"economic_cycle": "expansion",
|
|
"central_bank_stance": "hawkish",
|
|
"geopolitical_risks": ["trade tensions", "energy prices"],
|
|
},
|
|
"key_themes": [
|
|
{"theme": "AI Infrastructure", "description": "Data center boom", "conviction": "high", "timeframe": "3-6 months"}
|
|
],
|
|
"stocks_to_investigate": [
|
|
{
|
|
"ticker": "NVDA",
|
|
"name": "NVIDIA Corp",
|
|
"sector": "Technology",
|
|
"rationale": "GPU demand for AI training",
|
|
"thesis_angle": "growth",
|
|
"conviction": "high",
|
|
"key_catalysts": ["H100 demand", "Blackwell launch"],
|
|
"risks": ["Supply constraints", "Competition"],
|
|
}
|
|
],
|
|
"risk_factors": ["Fed rate hikes", "China tensions"],
|
|
}
|
|
import json
|
|
text = json.dumps(data)
|
|
result = extract_json(text)
|
|
assert result["timeframe"] == "1 month"
|
|
assert result["stocks_to_investigate"][0]["ticker"] == "NVDA"
|
|
|
|
|
|
def test_deepseek_r1_realistic():
|
|
"""Simulate a real DeepSeek R1 response with think block and JSON fence."""
|
|
text = (
|
|
"<think>\n"
|
|
"Let me analyze the macro environment. The geopolitical scanner shows tension...\n"
|
|
"I need to identify the top 8-10 stocks.\n"
|
|
"</think>\n"
|
|
"```json\n"
|
|
'{"timeframe": "1 month", "executive_summary": "Bullish macro backdrop", '
|
|
'"macro_context": {"economic_cycle": "expansion", "central_bank_stance": "neutral", "geopolitical_risks": []}, '
|
|
'"key_themes": [], "stocks_to_investigate": [{"ticker": "AAPL", "name": "Apple", "sector": "Technology", '
|
|
'"rationale": "Strong cash flows", "thesis_angle": "value", "conviction": "high", '
|
|
'"key_catalysts": ["Services growth"], "risks": ["China sales"]}], "risk_factors": []}\n'
|
|
"```"
|
|
)
|
|
result = extract_json(text)
|
|
assert result["timeframe"] == "1 month"
|
|
assert result["stocks_to_investigate"][0]["ticker"] == "AAPL"
|
|
|
|
|
|
def test_preamble_and_postamble():
|
|
"""JSON buried in prose before and after."""
|
|
text = 'Based on my analysis of the market data:\n\n{"result": 42}\n\nThis concludes my analysis.'
|
|
assert extract_json(text) == {"result": 42}
|
|
|
|
|
|
# ─── Error cases ──────────────────────────────────────────────────────────────
|
|
|
|
def test_empty_input():
|
|
with pytest.raises(ValueError, match="Empty input"):
|
|
extract_json("")
|
|
|
|
|
|
def test_whitespace_only():
|
|
with pytest.raises(ValueError, match="Empty input"):
|
|
extract_json(" \n\t ")
|
|
|
|
|
|
def test_malformed_json_no_fallback():
|
|
with pytest.raises(ValueError):
|
|
extract_json('{"key": value_without_quotes}')
|
|
|
|
|
|
def test_no_json_at_all():
|
|
with pytest.raises(ValueError):
|
|
extract_json("Just some text with no JSON structure at all")
|
|
|
|
|
|
def test_array_input_raises_value_error():
|
|
"""extract_json rejects JSON arrays — only dicts are accepted.
|
|
|
|
All callers (macro_synthesis, macro_bridge, CLI) call .get() on the result,
|
|
so returning a list would cause AttributeError downstream. The function
|
|
enforces dict-only return at runtime.
|
|
"""
|
|
with pytest.raises(ValueError, match="Expected a JSON object"):
|
|
extract_json('[1, 2, 3]')
|