feat: optimize TimesFM compute, add offline test suite, and implement YMYL quality gates

app/models/__init__.py

@@ -0,0 +1 @@
+# Models package

app/services/database_service.py

@@ -182,6 +182,15 @@ class DatabaseService:
             # No stocks table exists — rename stocks_v2 to stocks
             await self._execute("ALTER TABLE stocks_v2 RENAME TO stocks")
 
+        # Add new columns for sync tracking (migration)
+        await self._add_column_if_not_exists("stocks", "trading_status", "TEXT DEFAULT 'active'")
+        await self._add_column_if_not_exists("stocks", "last_synced", "INTEGER DEFAULT 0")
+        await self._add_column_if_not_exists("stocks", "market_cap_tier", "TEXT DEFAULT NULL")
+        await self._add_column_if_not_exists("stocks", "data_quality_score", "INTEGER DEFAULT 100")
+        await self._add_column_if_not_exists("stocks", "yfinance_available", "INTEGER DEFAULT 1")
+        await self._add_column_if_not_exists("stocks", "description", "TEXT DEFAULT NULL")
+        await self._add_column_if_not_exists("stocks", "listing_date", "TEXT DEFAULT NULL")
+
         # Create indexes
         await self._execute("""
             CREATE INDEX IF NOT EXISTS idx_stocks_symbol
@@ -213,15 +222,6 @@ class DatabaseService:
             ON forecasts(created_at)
         """)
 
-        # Add new columns for sync tracking (migration)
-        await self._add_column_if_not_exists("stocks", "trading_status", "TEXT DEFAULT 'active'")
-        await self._add_column_if_not_exists("stocks", "last_synced", "INTEGER DEFAULT 0")
-        await self._add_column_if_not_exists("stocks", "market_cap_tier", "TEXT DEFAULT NULL")
-        await self._add_column_if_not_exists("stocks", "data_quality_score", "INTEGER DEFAULT 100")
-        await self._add_column_if_not_exists("stocks", "yfinance_available", "INTEGER DEFAULT 1")
-        await self._add_column_if_not_exists("stocks", "description", "TEXT DEFAULT NULL")
-        await self._add_column_if_not_exists("stocks", "listing_date", "TEXT DEFAULT NULL")
-
         # Create sync log table
         await self._execute("""
             CREATE TABLE IF NOT EXISTS sync_log (
@@ -716,35 +716,35 @@ class DatabaseService:
             logger.error(f"Failed to create sync log: {e}")
             return None
 
-    async def update_sync_log(self, log_id: int, **updates):
+    async def update_sync_log(self, log_id: int, **updates_dict):
         """Update sync log entry"""
         try:
-            updates = []
+            clauses = []
             params = []
             
-            if 'completed_at' in updates:
-                updates.append("completed_at = ?")
-                params.append(updates['completed_at'])
-            if 'stocks_added' in updates:
-                updates.append("stocks_added = ?")
-                params.append(updates['stocks_added'])
-            if 'stocks_removed' in updates:
-                updates.append("stocks_removed = ?")
-                params.append(updates['stocks_removed'])
-            if 'stocks_updated' in updates:
-                updates.append("stocks_updated = ?")
-                params.append(updates['stocks_updated'])
-            if 'errors' in updates:
-                updates.append("errors = ?")
-                params.append(updates['errors'])
-            if 'status' in updates:
-                updates.append("status = ?")
-                params.append(updates['status'])
+            if 'completed_at' in updates_dict:
+                clauses.append("completed_at = ?")
+                params.append(updates_dict['completed_at'])
+            if 'stocks_added' in updates_dict:
+                clauses.append("stocks_added = ?")
+                params.append(updates_dict['stocks_added'])
+            if 'stocks_removed' in updates_dict:
+                clauses.append("stocks_removed = ?")
+                params.append(updates_dict['stocks_removed'])
+            if 'stocks_updated' in updates_dict:
+                clauses.append("stocks_updated = ?")
+                params.append(updates_dict['stocks_updated'])
+            if 'errors' in updates_dict:
+                clauses.append("errors = ?")
+                params.append(updates_dict['errors'])
+            if 'status' in updates_dict:
+                clauses.append("status = ?")
+                params.append(updates_dict['status'])
             
             params.append(log_id)
             
-            if updates:
-                sql = f"UPDATE sync_log SET {', '.join(updates)} WHERE id = ?"
+            if clauses:
+                sql = f"UPDATE sync_log SET {', '.join(clauses)} WHERE id = ?"
                 await self._execute(sql, params)
         except Exception as e:
             logger.error(f"Failed to update sync log: {e}")

requirements.txt

@@ -16,3 +16,5 @@ httpx==0.28.0
 tenacity==9.0.0
 python-multipart==0.0.12
 curl-cffi==0.15.0
+pytest==7.4.3
+pytest-asyncio==0.21.1

scripts/daily_pipeline.py

@@ -111,24 +111,41 @@ async def process_stock(
             print("[WARN] Insufficient data")
             return False
 
-        # Generate forecasts for each horizon
+        # Generate forecast once at the maximum horizon
+        max_horizon = max(HORIZONS)
+        full_forecast = await timesfm.predict(
+            historical_prices=stock_data['Close'].tolist(),
+            horizon=max_horizon,
+        )
+
+        if full_forecast is None:
+            print(f"[WARN] Forecast failed (h={max_horizon})")
+            return False
+
+        # Generate forecasts and charts for each horizon by slicing
         for horizon in HORIZONS:
-            forecast = await timesfm.predict(
-                historical_prices=stock_data['Close'].tolist(),
-                horizon=horizon,
-            )
+            # Slice point forecast and quantiles to the current horizon length
+            sliced_point_forecast = full_forecast['point_forecast'][:horizon]
+            sliced_quantiles = {
+                'p10': full_forecast['quantiles']['p10'][:horizon],
+                'p50': full_forecast['quantiles']['p50'][:horizon],
+                'p90': full_forecast['quantiles']['p90'][:horizon],
+            }
 
-            if forecast is None:
-                print(f"[WARN] Forecast failed (h={horizon})")
-                continue
+            # Reconstruct horizon-specific forecast object for the chart service
+            horizon_forecast = {
+                'point_forecast': sliced_point_forecast,
+                'quantiles': sliced_quantiles,
+                'method': full_forecast.get('method', 'timesfm')
+            }
 
-            # Generate chart for every horizon
+            # Generate chart for every horizon using the sliced data
             chart_svg = None
             try:
                 chart_svg = chart_service.generate_forecast_chart(
                     symbol=clean_symbol,
                     historical_prices=stock_data['Close'].tolist()[-60:],
-                    forecast=forecast,
+                    forecast=horizon_forecast,
                     current_price=float(stock_data['Close'].iloc[-1]),
                 )
             except Exception as ce:
@@ -146,15 +163,15 @@ async def process_stock(
                 'current_price': float(stock_data['Close'].iloc[-1]),
                 'last_updated': stock_data.index[-1].isoformat(),
                 'horizon_days': horizon,
-                'point_forecast': float(forecast['point_forecast'][-1]),
+                'point_forecast': float(sliced_point_forecast[-1]),
                 'percentage_change': float(
-                    ((forecast['point_forecast'][-1] - stock_data['Close'].iloc[-1])
+                    ((sliced_point_forecast[-1] - stock_data['Close'].iloc[-1])
                      / stock_data['Close'].iloc[-1]) * 100
                 ),
                 'quantiles': {
-                    'p10': [float(x) for x in forecast['quantiles']['p10']],
-                    'p50': [float(x) for x in forecast['quantiles']['p50']],
-                    'p90': [float(x) for x in forecast['quantiles']['p90']],
+                    'p10': [float(x) for x in sliced_quantiles['p10']],
+                    'p50': [float(x) for x in sliced_quantiles['p50']],
+                    'p90': [float(x) for x in sliced_quantiles['p90']],
                 },
                 'chart_svg': chart_svg,
                 'methodology_version': 'timesfm-2.5-200m-v1.0',

tests/test_caching.py

@@ -0,0 +1,60 @@
+import sqlite3
+import pytest
+import json
+import time
+from app.services.database_service import DatabaseService
+
+class SQLiteTestDatabaseService(DatabaseService):
+    def __init__(self):
+        super().__init__(url="libsql://dummy", token="dummy")
+        self.conn = sqlite3.connect(":memory:")
+        self.conn.row_factory = sqlite3.Row
+        self.connected = True
+        
+    async def initialize(self):
+        await self._create_tables()
+        
+    async def _execute(self, sql: str, params: list = None) -> list[dict]:
+        cursor = self.conn.cursor()
+        cursor.execute(sql, params or [])
+        self.conn.commit()
+        if cursor.description:
+            cols = [col[0] for col in cursor.description]
+            return [dict(zip(cols, row)) for row in cursor.fetchall()]
+        return []
+
+@pytest.mark.asyncio
+async def test_cache_forecast_lifecycle():
+    db = SQLiteTestDatabaseService()
+    await db.initialize()
+    
+    forecast_data = {
+        "symbol": "AAPL",
+        "horizon_days": 10,
+        "point_forecast": 150.0,
+        "percentage_change": 1.5
+    }
+    
+    await db.cache_forecast("AAPL", 10, forecast_data)
+    
+    cached = await db.get_cached_forecast("AAPL", 10, max_age_hours=24)
+    assert cached is not None
+    assert cached["symbol"] == "AAPL"
+
+@pytest.mark.asyncio
+async def test_cache_forecast_expiration():
+    db = SQLiteTestDatabaseService()
+    await db.initialize()
+    
+    forecast_data = {"symbol": "AAPL", "horizon_days": 10}
+    stale_time = int(time.time()) - (25 * 3600)  # 25 hours ago
+    
+    await db._execute(
+        "INSERT INTO forecasts (symbol, horizon_days, forecast_data, created_at) VALUES (?, ?, ?, ?)",
+        ["AAPL", 10, json.dumps(forecast_data), stale_time]
+    )
+    
+    # Stale cache retrieval should return None for 24h limit
+    assert await db.get_cached_forecast("AAPL", 10, max_age_hours=24) is None
+    # But valid if max age is 48h
+    assert await db.get_cached_forecast("AAPL", 10, max_age_hours=48) is not None

tests/test_forecast_api.py

@@ -0,0 +1,100 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+from fastapi.testclient import TestClient
+import pandas as pd
+import app.main as main
+
+# Create TestClient without context manager to bypass actual lifespan events
+client = TestClient(main.app, raise_server_exceptions=False)
+
+@pytest.fixture(autouse=True)
+def setup_mocks(monkeypatch):
+    mock_db = AsyncMock()
+    mock_db.connected = True
+    
+    mock_timesfm = AsyncMock()
+    mock_timesfm.model = MagicMock()
+    
+    mock_data = AsyncMock()
+    mock_chart = MagicMock()
+    
+    monkeypatch.setattr(main, "db_service", mock_db)
+    monkeypatch.setattr(main, "timesfm_service", mock_timesfm)
+    monkeypatch.setattr(main, "data_service", mock_data)
+    monkeypatch.setattr(main, "chart_service", mock_chart)
+    
+    # Disable rate limiting for testing by mocking the _check method
+    monkeypatch.setattr("app.middleware.RateLimitMiddleware._check", lambda self, ip, path: (False, 0))
+    
+    return mock_db, mock_timesfm, mock_data, mock_chart
+
+def test_health_check_healthy(setup_mocks):
+    mock_db, _, _, _ = setup_mocks
+    mock_db._execute.return_value = [{"1": 1}]
+    
+    response = client.get("/health")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "healthy"
+    assert data["database_connected"] is True
+    assert data["model_loaded"] is True
+
+def test_health_check_db_disconnected(setup_mocks):
+    mock_db, _, _, _ = setup_mocks
+    mock_db.connected = False
+    
+    response = client.get("/health")
+    assert response.status_code == 200
+    assert response.json()["database_connected"] is False
+
+def test_get_forecast_cache_hit(setup_mocks):
+    mock_db, _, _, _ = setup_mocks
+    mock_db.get_cached_forecast.return_value = {
+        "symbol": "AAPL",
+        "name": "Apple Inc.",
+        "exchange": "NASDAQ",
+        "currency": "USD",
+        "current_price": 150.0,
+        "last_updated": "2026-06-19T00:00:00",
+        "horizon_days": 20,
+        "point_forecast": 155.0,
+        "percentage_change": 3.33,
+        "quantiles": {"p10": [148.0], "p50": [155.0], "p90": [162.0]},
+        "chart_svg": "<svg></svg>",
+        "methodology_version": "timesfm-2.5-200m-v1.0"
+    }
+    
+    response = client.get("/api/v1/forecast/AAPL?horizon=20")
+    assert response.status_code == 200
+    assert response.json()["symbol"] == "AAPL"
+    mock_db.get_cached_forecast.assert_called_once_with("AAPL", 20)
+
+def test_get_forecast_cache_miss_and_generate(setup_mocks):
+    mock_db, mock_timesfm, mock_data, mock_chart = setup_mocks
+    mock_db.get_cached_forecast.return_value = None
+    
+    mock_df = pd.DataFrame(
+        {"Close": [140.0 + i for i in range(100)]},
+        index=pd.date_range("2026-01-01", periods=100)
+    )
+    mock_df.attrs["name"] = "Apple Inc."
+    mock_df.attrs["exchange"] = "NASDAQ"
+    mock_df.attrs["currency"] = "USD"
+    mock_data.get_stock_data.return_value = mock_df
+    
+    mock_timesfm.predict.return_value = {
+        "point_forecast": [245.0] * 20,
+        "quantiles": {
+            "p10": [240.0] * 20,
+            "p50": [245.0] * 20,
+            "p90": [250.0] * 20
+        }
+    }
+    mock_chart.generate_forecast_chart.return_value = "<svg>chart</svg>"
+    
+    response = client.get("/api/v1/forecast/AAPL?horizon=20")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["symbol"] == "AAPL"
+    assert data["current_price"] == 239.0
+    mock_db.cache_forecast.assert_called_once()