mirror of
https://github.com/Xe138/AI-Trader.git
synced 2026-04-08 19:57:24 -04:00
feat: store reasoning logs with sessions in model_day_executor
- Add _create_trading_session() method to create session records - Add async _store_reasoning_logs() to store conversation with AI summaries - Add async _update_session_summary() to generate overall session summary - Modify execute() -> execute_async() with async workflow - Add execute_sync() wrapper and keep execute() as sync entry point - Update _write_results_to_db() to accept and use session_id parameter - Modify positions INSERT to include session_id foreign key - Remove old reasoning_logs code block (obsolete schema) - Add comprehensive unit tests for all new functionality All tests pass. Session-based reasoning storage now integrated.
This commit is contained in:
266
tests/unit/test_model_day_executor_reasoning.py
Normal file
266
tests/unit/test_model_day_executor_reasoning.py
Normal file
@@ -0,0 +1,266 @@
|
||||
"""Tests for reasoning log storage in model_day_executor."""
|
||||
|
||||
import pytest
|
||||
import sqlite3
|
||||
from api.model_day_executor import ModelDayExecutor
|
||||
from api.database import initialize_database, get_db_connection
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_db(tmp_path):
|
||||
"""Create test database with job record."""
|
||||
db_path = str(tmp_path / "test.db")
|
||||
initialize_database(db_path)
|
||||
|
||||
# Create a job record to satisfy foreign key constraint
|
||||
conn = get_db_connection(db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
INSERT INTO jobs (job_id, config_path, status, date_range, models, created_at)
|
||||
VALUES ('test-job', 'configs/default_config.json', 'running', '["2025-01-01"]', '["test-model"]', '2025-01-01T00:00:00Z')
|
||||
""")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
return db_path
|
||||
|
||||
|
||||
def test_create_trading_session(test_db):
|
||||
"""Should create trading session record."""
|
||||
executor = ModelDayExecutor(
|
||||
job_id="test-job",
|
||||
date="2025-01-01",
|
||||
model_sig="test-model",
|
||||
config_path="configs/default_config.json",
|
||||
db_path=test_db
|
||||
)
|
||||
|
||||
conn = get_db_connection(test_db)
|
||||
cursor = conn.cursor()
|
||||
|
||||
session_id = executor._create_trading_session(cursor)
|
||||
conn.commit()
|
||||
|
||||
# Verify session created
|
||||
cursor.execute("SELECT * FROM trading_sessions WHERE id = ?", (session_id,))
|
||||
session = cursor.fetchone()
|
||||
|
||||
assert session is not None
|
||||
assert session['job_id'] == "test-job"
|
||||
assert session['date'] == "2025-01-01"
|
||||
assert session['model'] == "test-model"
|
||||
assert session['started_at'] is not None
|
||||
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_reasoning_logs(test_db):
|
||||
"""Should store conversation with summaries."""
|
||||
from agent.mock_provider.mock_langchain_model import MockChatModel
|
||||
from agent.base_agent.base_agent import BaseAgent
|
||||
|
||||
executor = ModelDayExecutor(
|
||||
job_id="test-job",
|
||||
date="2025-01-01",
|
||||
model_sig="test-model",
|
||||
config_path="configs/default_config.json",
|
||||
db_path=test_db
|
||||
)
|
||||
|
||||
# Create mock agent
|
||||
agent = BaseAgent(
|
||||
signature="test-model",
|
||||
basemodel="mock",
|
||||
stock_symbols=["AAPL"],
|
||||
init_date="2025-01-01"
|
||||
)
|
||||
agent.model = MockChatModel(model="test", signature="test")
|
||||
|
||||
# Create conversation
|
||||
conversation = [
|
||||
{"role": "user", "content": "Analyze market", "timestamp": "2025-01-01T10:00:00Z"},
|
||||
{"role": "assistant", "content": "Bought AAPL 10 shares based on strong earnings", "timestamp": "2025-01-01T10:05:00Z"}
|
||||
]
|
||||
|
||||
conn = get_db_connection(test_db)
|
||||
cursor = conn.cursor()
|
||||
session_id = executor._create_trading_session(cursor)
|
||||
|
||||
await executor._store_reasoning_logs(cursor, session_id, conversation, agent)
|
||||
conn.commit()
|
||||
|
||||
# Verify logs stored
|
||||
cursor.execute("SELECT * FROM reasoning_logs WHERE session_id = ? ORDER BY message_index", (session_id,))
|
||||
logs = cursor.fetchall()
|
||||
|
||||
assert len(logs) == 2
|
||||
assert logs[0]['role'] == 'user'
|
||||
assert logs[0]['content'] == 'Analyze market'
|
||||
assert logs[0]['summary'] is None # No summary for user messages
|
||||
|
||||
assert logs[1]['role'] == 'assistant'
|
||||
assert logs[1]['content'] == 'Bought AAPL 10 shares based on strong earnings'
|
||||
assert logs[1]['summary'] is not None # Summary generated for assistant
|
||||
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_session_summary(test_db):
|
||||
"""Should update session with overall summary."""
|
||||
from agent.mock_provider.mock_langchain_model import MockChatModel
|
||||
from agent.base_agent.base_agent import BaseAgent
|
||||
|
||||
executor = ModelDayExecutor(
|
||||
job_id="test-job",
|
||||
date="2025-01-01",
|
||||
model_sig="test-model",
|
||||
config_path="configs/default_config.json",
|
||||
db_path=test_db
|
||||
)
|
||||
|
||||
# Create mock agent
|
||||
agent = BaseAgent(
|
||||
signature="test-model",
|
||||
basemodel="mock",
|
||||
stock_symbols=["AAPL"],
|
||||
init_date="2025-01-01"
|
||||
)
|
||||
agent.model = MockChatModel(model="test", signature="test")
|
||||
|
||||
# Create conversation
|
||||
conversation = [
|
||||
{"role": "user", "content": "Analyze market", "timestamp": "2025-01-01T10:00:00Z"},
|
||||
{"role": "assistant", "content": "Bought AAPL 10 shares", "timestamp": "2025-01-01T10:05:00Z"},
|
||||
{"role": "assistant", "content": "Sold MSFT 5 shares", "timestamp": "2025-01-01T10:10:00Z"}
|
||||
]
|
||||
|
||||
conn = get_db_connection(test_db)
|
||||
cursor = conn.cursor()
|
||||
session_id = executor._create_trading_session(cursor)
|
||||
|
||||
await executor._update_session_summary(cursor, session_id, conversation, agent)
|
||||
conn.commit()
|
||||
|
||||
# Verify session updated
|
||||
cursor.execute("SELECT * FROM trading_sessions WHERE id = ?", (session_id,))
|
||||
session = cursor.fetchone()
|
||||
|
||||
assert session['session_summary'] is not None
|
||||
assert len(session['session_summary']) > 0
|
||||
assert session['completed_at'] is not None
|
||||
assert session['total_messages'] == 3
|
||||
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_reasoning_logs_with_tool_messages(test_db):
|
||||
"""Should store tool messages with tool_name and tool_input."""
|
||||
from agent.mock_provider.mock_langchain_model import MockChatModel
|
||||
from agent.base_agent.base_agent import BaseAgent
|
||||
|
||||
executor = ModelDayExecutor(
|
||||
job_id="test-job",
|
||||
date="2025-01-01",
|
||||
model_sig="test-model",
|
||||
config_path="configs/default_config.json",
|
||||
db_path=test_db
|
||||
)
|
||||
|
||||
# Create mock agent
|
||||
agent = BaseAgent(
|
||||
signature="test-model",
|
||||
basemodel="mock",
|
||||
stock_symbols=["AAPL"],
|
||||
init_date="2025-01-01"
|
||||
)
|
||||
agent.model = MockChatModel(model="test", signature="test")
|
||||
|
||||
# Create conversation with tool message
|
||||
conversation = [
|
||||
{"role": "user", "content": "Get price", "timestamp": "2025-01-01T10:00:00Z"},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "AAPL: $150.00",
|
||||
"tool_name": "get_price",
|
||||
"tool_input": '{"symbol": "AAPL"}',
|
||||
"timestamp": "2025-01-01T10:01:00Z"
|
||||
},
|
||||
{"role": "assistant", "content": "AAPL is $150", "timestamp": "2025-01-01T10:02:00Z"}
|
||||
]
|
||||
|
||||
conn = get_db_connection(test_db)
|
||||
cursor = conn.cursor()
|
||||
session_id = executor._create_trading_session(cursor)
|
||||
|
||||
await executor._store_reasoning_logs(cursor, session_id, conversation, agent)
|
||||
conn.commit()
|
||||
|
||||
# Verify tool message stored correctly
|
||||
cursor.execute("SELECT * FROM reasoning_logs WHERE session_id = ? AND role = 'tool'", (session_id,))
|
||||
tool_log = cursor.fetchone()
|
||||
|
||||
assert tool_log is not None
|
||||
assert tool_log['tool_name'] == 'get_price'
|
||||
assert tool_log['tool_input'] == '{"symbol": "AAPL"}'
|
||||
assert tool_log['content'] == 'AAPL: $150.00'
|
||||
assert tool_log['summary'] is None # No summary for tool messages
|
||||
|
||||
conn.close()
|
||||
|
||||
|
||||
def test_write_results_includes_session_id(test_db):
|
||||
"""Should include session_id when writing positions."""
|
||||
from agent.mock_provider.mock_langchain_model import MockChatModel
|
||||
from agent.base_agent.base_agent import BaseAgent
|
||||
|
||||
executor = ModelDayExecutor(
|
||||
job_id="test-job",
|
||||
date="2025-01-01",
|
||||
model_sig="test-model",
|
||||
config_path="configs/default_config.json",
|
||||
db_path=test_db
|
||||
)
|
||||
|
||||
# Create mock agent with positions
|
||||
agent = BaseAgent(
|
||||
signature="test-model",
|
||||
basemodel="mock",
|
||||
stock_symbols=["AAPL"],
|
||||
init_date="2025-01-01"
|
||||
)
|
||||
agent.model = MockChatModel(model="test", signature="test")
|
||||
|
||||
# Mock positions data
|
||||
agent.positions = {"AAPL": 10, "CASH": 8500.0}
|
||||
agent.last_trade = {"action": "buy", "symbol": "AAPL", "amount": 10, "price": 150.0}
|
||||
agent.current_prices = {"AAPL": 150.0}
|
||||
|
||||
# Add required methods
|
||||
agent.get_positions = lambda: agent.positions
|
||||
agent.get_last_trade = lambda: agent.last_trade
|
||||
agent.get_current_prices = lambda: agent.current_prices
|
||||
|
||||
conn = get_db_connection(test_db)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Create session
|
||||
session_id = executor._create_trading_session(cursor)
|
||||
conn.commit()
|
||||
|
||||
# Write results
|
||||
executor._write_results_to_db(agent, session_id)
|
||||
|
||||
# Verify position has session_id
|
||||
cursor.execute("SELECT * FROM positions WHERE job_id = ? AND model = ?",
|
||||
("test-job", "test-model"))
|
||||
position = cursor.fetchone()
|
||||
|
||||
assert position is not None
|
||||
assert position['session_id'] == session_id
|
||||
assert position['action_type'] == 'buy'
|
||||
assert position['symbol'] == 'AAPL'
|
||||
|
||||
conn.close()
|
||||
Reference in New Issue
Block a user