305 lines
11 KiB
Python
305 lines
11 KiB
Python
"""MockEngine — streams scripted events for UI testing without real LLM calls.
|
|
|
|
Usage (via POST /api/run/mock):
|
|
params = {
|
|
"mock_type": "pipeline" | "scan" | "auto",
|
|
"ticker": "AAPL", # used for pipeline / auto
|
|
"tickers": ["AAPL","NVDA"], # used for auto (overrides ticker list)
|
|
"date": "2026-03-24",
|
|
"speed": 2.0, # delay divisor — higher = faster
|
|
}
|
|
"""
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import AsyncGenerator, Dict, Any
|
|
|
|
|
|
class MockEngine:
|
|
"""Generates scripted AgentOS events without calling real LLMs."""
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public entry points
|
|
# ------------------------------------------------------------------
|
|
|
|
async def run_mock(
|
|
self, run_id: str, params: Dict[str, Any]
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
mock_type = params.get("mock_type", "pipeline")
|
|
speed = max(float(params.get("speed", 1.0)), 0.1)
|
|
|
|
if mock_type == "scan":
|
|
async for evt in self._run_scan(run_id, params, speed):
|
|
yield evt
|
|
elif mock_type == "auto":
|
|
async for evt in self._run_auto(run_id, params, speed):
|
|
yield evt
|
|
else:
|
|
async for evt in self._run_pipeline(run_id, params, speed):
|
|
yield evt
|
|
|
|
# ------------------------------------------------------------------
|
|
# Pipeline mock
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _run_pipeline(
|
|
self, run_id: str, params: Dict[str, Any], speed: float
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
ticker = params.get("ticker", "AAPL").upper()
|
|
date = params.get("date", time.strftime("%Y-%m-%d"))
|
|
|
|
yield self._log(f"[MOCK] Starting pipeline for {ticker} on {date}")
|
|
await self._sleep(0.3, speed)
|
|
|
|
# Analysts (sequential for simplicity in mock)
|
|
analysts = [
|
|
("news_analyst", "gpt-4o-mini", 1.4, 480, 310),
|
|
("market_analyst", "gpt-4o-mini", 1.2, 390, 240),
|
|
("fundamentals_analyst", "gpt-4o", 2.1, 620, 430),
|
|
("social_analyst", "gpt-4o-mini", 0.9, 310, 190),
|
|
]
|
|
for node, model, latency, tok_in, tok_out in analysts:
|
|
async for evt in self._agent_with_tool(
|
|
run_id, node, ticker, model, latency, tok_in, tok_out, speed,
|
|
tool_name=f"get_{node.split('_')[0]}_data",
|
|
):
|
|
yield evt
|
|
|
|
# Research debate
|
|
for node, model, latency, tok_in, tok_out in [
|
|
("bull_researcher", "gpt-4o", 1.8, 540, 360),
|
|
("bear_researcher", "gpt-4o", 1.7, 510, 340),
|
|
("research_manager","gpt-4o", 2.3, 680, 480),
|
|
]:
|
|
async for evt in self._agent_no_tool(
|
|
run_id, node, ticker, model, latency, tok_in, tok_out, speed
|
|
):
|
|
yield evt
|
|
|
|
# Trading decision
|
|
for node, model, latency, tok_in, tok_out in [
|
|
("trader", "gpt-4o", 2.0, 600, 420),
|
|
("risk_manager", "gpt-4o", 1.5, 450, 310),
|
|
("risk_judge", "gpt-4o", 1.1, 380, 260),
|
|
]:
|
|
async for evt in self._agent_no_tool(
|
|
run_id, node, ticker, model, latency, tok_in, tok_out, speed
|
|
):
|
|
yield evt
|
|
|
|
yield self._log(f"[MOCK] Pipeline for {ticker} completed.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Scan mock
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _run_scan(
|
|
self, run_id: str, params: Dict[str, Any], speed: float
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
date = params.get("date", time.strftime("%Y-%m-%d"))
|
|
identifier = "MARKET"
|
|
|
|
yield self._log(f"[MOCK] Starting market scan for {date}")
|
|
await self._sleep(0.3, speed)
|
|
|
|
# Phase 1 — three scanners in "parallel" (interleaved)
|
|
yield self._log("[MOCK] Phase 1: Running geopolitical, market-movers, sector scanners…")
|
|
phase1 = [
|
|
("geopolitical_scanner", "gpt-4o-mini", 1.5, 420, 280),
|
|
("market_movers_scanner", "gpt-4o-mini", 1.3, 380, 250),
|
|
("sector_scanner", "gpt-4o-mini", 1.4, 400, 265),
|
|
]
|
|
# Emit thought events for all three before any complete
|
|
for node, model, _, _, _ in phase1:
|
|
yield self._thought(node, identifier, model, f"[MOCK] Scanning {node.replace('_', ' ')}…")
|
|
await self._sleep(0.1, speed)
|
|
|
|
# Then complete them in order
|
|
for node, model, latency, tok_in, tok_out in phase1:
|
|
await self._sleep(latency, speed)
|
|
yield self._result(node, identifier, model, tok_in, tok_out, round(latency * 1000),
|
|
f"[MOCK] {node.replace('_', ' ').title()} report ready.")
|
|
|
|
# Phase 2 — industry deep dive
|
|
yield self._log("[MOCK] Phase 2: Industry deep dive…")
|
|
async for evt in self._agent_no_tool(
|
|
run_id, "industry_deep_dive", identifier, "gpt-4o", 2.2, 680, 460, speed
|
|
):
|
|
yield evt
|
|
|
|
# Phase 3 — macro synthesis
|
|
yield self._log("[MOCK] Phase 3: Macro synthesis + watchlist generation…")
|
|
async for evt in self._agent_no_tool(
|
|
run_id, "macro_synthesis", identifier, "gpt-4o", 2.8, 820, 590, speed
|
|
):
|
|
yield evt
|
|
|
|
yield self._log("[MOCK] Scan completed. Top-10 watchlist generated.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Auto mock (scan → pipeline per ticker → portfolio)
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _run_auto(
|
|
self, run_id: str, params: Dict[str, Any], speed: float
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
tickers = params.get("tickers") or [params.get("ticker", "AAPL").upper()]
|
|
|
|
yield self._log(f"[MOCK] Starting auto run — scan + {len(tickers)} pipeline(s) + portfolio")
|
|
await self._sleep(0.2, speed)
|
|
|
|
# Phase 1: Scan
|
|
yield self._log("[MOCK] Phase 1/3: Market scan…")
|
|
async for evt in self._run_scan(run_id, params, speed):
|
|
yield evt
|
|
|
|
# Phase 2: Per-ticker pipeline
|
|
for ticker in tickers:
|
|
yield self._log(f"[MOCK] Phase 2/3: Pipeline for {ticker}…")
|
|
async for evt in self._run_pipeline(run_id, {**params, "ticker": ticker}, speed):
|
|
yield evt
|
|
|
|
# Phase 3: Portfolio
|
|
yield self._log("[MOCK] Phase 3/3: Portfolio manager…")
|
|
async for evt in self._agent_no_tool(
|
|
run_id, "portfolio_manager", "PORTFOLIO", "gpt-4o", 2.5, 740, 520, speed
|
|
):
|
|
yield evt
|
|
|
|
yield self._log("[MOCK] Auto run completed.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Building blocks
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _agent_with_tool(
|
|
self,
|
|
run_id: str,
|
|
node: str,
|
|
identifier: str,
|
|
model: str,
|
|
latency: float,
|
|
tok_in: int,
|
|
tok_out: int,
|
|
speed: float,
|
|
tool_name: str,
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
yield self._thought(node, identifier, model, f"[MOCK] {node} analysing {identifier}…")
|
|
await self._sleep(0.4, speed)
|
|
|
|
yield self._tool_call(node, identifier, tool_name, f'{{"ticker": "{identifier}"}}')
|
|
await self._sleep(0.6, speed)
|
|
|
|
yield self._tool_result(node, identifier, tool_name,
|
|
f"[MOCK] Retrieved {tool_name} data for {identifier}.")
|
|
await self._sleep(latency - 1.0, speed)
|
|
|
|
yield self._result(node, identifier, model, tok_in, tok_out, round(latency * 1000),
|
|
f"[MOCK] {node.replace('_', ' ').title()} analysis complete for {identifier}.")
|
|
|
|
async def _agent_no_tool(
|
|
self,
|
|
run_id: str,
|
|
node: str,
|
|
identifier: str,
|
|
model: str,
|
|
latency: float,
|
|
tok_in: int,
|
|
tok_out: int,
|
|
speed: float,
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
yield self._thought(node, identifier, model, f"[MOCK] {node} processing {identifier}…")
|
|
await self._sleep(latency, speed)
|
|
|
|
yield self._result(node, identifier, model, tok_in, tok_out, round(latency * 1000),
|
|
f"[MOCK] {node.replace('_', ' ').title()} decision for {identifier}.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Event constructors
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _ns() -> str:
|
|
return str(time.time_ns())
|
|
|
|
def _log(self, message: str) -> Dict[str, Any]:
|
|
return {
|
|
"id": f"log_{self._ns()}",
|
|
"node_id": "__system__",
|
|
"type": "log",
|
|
"agent": "SYSTEM",
|
|
"identifier": "",
|
|
"message": message,
|
|
"metrics": {},
|
|
}
|
|
|
|
def _thought(self, node: str, identifier: str, model: str, message: str) -> Dict[str, Any]:
|
|
return {
|
|
"id": f"thought_{self._ns()}",
|
|
"node_id": node,
|
|
"parent_node_id": "start",
|
|
"type": "thought",
|
|
"agent": node.upper(),
|
|
"identifier": identifier,
|
|
"message": message,
|
|
"prompt": f"[MOCK PROMPT] Analyse {identifier} using available data.",
|
|
"metrics": {"model": model},
|
|
}
|
|
|
|
def _tool_call(self, node: str, identifier: str, tool: str, inp: str) -> Dict[str, Any]:
|
|
return {
|
|
"id": f"tool_{self._ns()}",
|
|
"node_id": f"tool_{tool}",
|
|
"parent_node_id": node,
|
|
"type": "tool",
|
|
"agent": node.upper(),
|
|
"identifier": identifier,
|
|
"message": f"▶ Tool: {tool} | {inp}",
|
|
"prompt": inp,
|
|
"metrics": {},
|
|
}
|
|
|
|
def _tool_result(self, node: str, identifier: str, tool: str, output: str) -> Dict[str, Any]:
|
|
return {
|
|
"id": f"tool_res_{self._ns()}",
|
|
"node_id": f"tool_{tool}",
|
|
"parent_node_id": node,
|
|
"type": "tool_result",
|
|
"agent": node.upper(),
|
|
"identifier": identifier,
|
|
"message": f"✓ Tool result: {tool} | {output}",
|
|
"response": output,
|
|
"metrics": {},
|
|
}
|
|
|
|
def _result(
|
|
self,
|
|
node: str,
|
|
identifier: str,
|
|
model: str,
|
|
tok_in: int,
|
|
tok_out: int,
|
|
latency_ms: int,
|
|
message: str,
|
|
) -> Dict[str, Any]:
|
|
return {
|
|
"id": f"result_{self._ns()}",
|
|
"node_id": node,
|
|
"type": "result",
|
|
"agent": node.upper(),
|
|
"identifier": identifier,
|
|
"message": message,
|
|
"response": f"[MOCK RESPONSE] {message}",
|
|
"metrics": {
|
|
"model": model,
|
|
"tokens_in": tok_in,
|
|
"tokens_out": tok_out,
|
|
"latency_ms": latency_ms,
|
|
},
|
|
}
|
|
|
|
@staticmethod
|
|
async def _sleep(seconds: float, speed: float) -> None:
|
|
delay = max(seconds / speed, 0.01)
|
|
await asyncio.sleep(delay)
|