diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 36727dc..ae306dd 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -13,12 +13,12 @@ use arbos_core::constants::{ }; use serde::Deserialize; -use std::str::FromStr; #[derive(Debug, Deserialize)] pub struct BrainStatePayload { pub implications: Vec, pub partitions: Vec, + pub contradictions: Vec, pub asset_end_timestamps: HashMap, } @@ -37,6 +37,13 @@ pub struct PartitionMapping { pub confidence: f64, } +#[derive(Debug, Deserialize)] +pub struct ContradictionMapping { + pub asset_a: String, + pub asset_b: String, + pub confidence: f64, +} + pub struct EngineActor { // State Caches pub orderbook_cache: HashMap, @@ -45,11 +52,17 @@ pub struct EngineActor { // Map of relation edges: Child -> List of Parents that imply it pub implication_reverse_edges: HashMap>, - // Map of Market (Condition) -> Set of Asset IDs (Outcomes) - pub market_outcomes: HashMap>, + // Map of undirected contradicting mutually-exclusive peers (A implies NOT B, etc.) + pub contradiction_edges: HashMap>, // Peer -> (Peer ID, Default Weight 1.0) + + // Map of Partition ID -> Set of Asset IDs (Outcomes) + pub partition_outcomes: HashMap>, - // Map of Market (Condition) -> Required Number of Outcomes for Exhaustiveness - pub market_expected_outcome_counts: HashMap, + // Map of Asset ID -> Partition ID (for O(1) partition lookups) + pub asset_to_partition: HashMap, + + // Map of Partition ID -> Required Number of Outcomes for Exhaustiveness + pub partition_expected_outcome_counts: HashMap, // Map of Asset ID -> Resolution Deadline (Timestamp in seconds) pub market_end_timestamps: HashMap, @@ -64,6 +77,12 @@ pub struct EngineActor { pub api_client: reqwest::Client, } +pub type PartitionParseResult = ( + HashMap>, + HashMap, + HashMap, +); + impl Default for EngineActor { fn default() -> Self { Self::new() @@ -81,8 +100,10 @@ impl EngineActor { orderbook_cache: HashMap::new(), implication_edges: HashMap::new(), implication_reverse_edges: HashMap::new(), - market_outcomes: HashMap::new(), - market_expected_outcome_counts: HashMap::new(), + contradiction_edges: HashMap::new(), + partition_outcomes: HashMap::new(), + asset_to_partition: HashMap::new(), + partition_expected_outcome_counts: HashMap::new(), market_end_timestamps: HashMap::new(), tracked_assets: HashSet::new(), current_gas_usdc_per_leg: (ESTIMATED_GAS_UNITS_PER_LEG * ESTIMATED_GWEI_PRICE) @@ -176,6 +197,8 @@ impl EngineActor { .await; self.evaluate_partitions(&book, bot_tx, fee_rate, gas_per_leg_usdc) .await; + self.evaluate_contradictions(&book, bot_tx, fee_rate, gas_per_leg_usdc) + .await; } async fn evaluate_implications( @@ -479,19 +502,24 @@ impl EngineActor { fee_rate: Decimal, gas_per_leg_usdc: Decimal, ) { - let condition_id = book.market; + let asset_id = book.asset_id; - let bids = match self.get_partition_bids(&condition_id) { + let partition_id = match self.asset_to_partition.get(&asset_id) { + Some(pid) => pid.clone(), + None => return, + }; + + let bids = match self.get_partition_bids(&partition_id) { Some(b) => b, None => return, }; - if !self.is_partition_exhaustively_tradable(&condition_id, &bids) { + if !self.is_partition_exhaustively_tradable(&partition_id, &bids) { return; } self.check_and_send_partition( - condition_id, + partition_id, bids, book.timestamp, bot_tx, @@ -501,11 +529,8 @@ impl EngineActor { .await; } - fn get_partition_bids( - &self, - condition_id: &polymarket_client_sdk::types::B256, - ) -> Option> { - let outcome_assets = self.market_outcomes.get(condition_id)?; + fn get_partition_bids(&self, partition_id: &str) -> Option> { + let outcome_assets = self.partition_outcomes.get(partition_id)?; let mut bids = Vec::new(); for asset in outcome_assets { @@ -525,13 +550,13 @@ impl EngineActor { fn is_partition_exhaustively_tradable( &self, - condition_id: &polymarket_client_sdk::types::B256, + partition_id: &str, bids: &[(U256, Decimal)], ) -> bool { // Enforce strict exhaustiveness: Are we receiving the EXACT number of outcomes for this partition? let required_outcome_count = self - .market_expected_outcome_counts - .get(condition_id) + .partition_expected_outcome_counts + .get(partition_id) .copied() .unwrap_or(0); @@ -545,7 +570,7 @@ impl EngineActor { async fn check_and_send_partition( &mut self, - condition_id: polymarket_client_sdk::types::B256, + partition_id: String, bids: Vec<(U256, Decimal)>, timestamp: i64, bot_tx: &Sender, @@ -578,7 +603,7 @@ impl EngineActor { timestamp, }; - info!(market = ?condition_id, leg_count = bids.len(), profit = %profit, "Routing Partition ArbSignal to Bot"); + info!(partition_id = %partition_id, leg_count = bids.len(), profit = %profit, "Routing Partition ArbSignal to Bot"); if let Err(e) = bot_tx.send(signal).await { error!("Failed to route Partition ArbSignal to Bot channel: {}", e); } @@ -590,6 +615,108 @@ impl EngineActor { } } + async fn evaluate_contradictions( + &mut self, + book: &NormalizedOrderbook, + bot_tx: &Sender, + fee_rate: Decimal, + gas_per_leg_usdc: Decimal, + ) { + let asset_id = book.asset_id; + + let peers_len = match self.contradiction_edges.get(&asset_id) { + Some(p) => p.len(), + None => return, + }; + + let mut executed_peer = None; + + for i in 0..peers_len { + let (peer_id, weight) = { + let peers = self.contradiction_edges.get(&asset_id).unwrap(); + peers[i] + }; + + let peer_book = match self.orderbook_cache.get(&peer_id) { + Some(b) => b, + None => continue, + }; + + if book.bid_untradeable || peer_book.bid_untradeable { + continue; + } + + let bid_a = book.vwap_bid; + let bid_b = peer_book.vwap_bid; + + if bid_a <= Decimal::ZERO || bid_b <= Decimal::ZERO { + continue; + } + + // A Contradiction strategy also executes exactly 2 legs (Sell A, Sell B) + let total_est_gas_usdc = gas_per_leg_usdc * Decimal::from(2); + + // Weight logic: for binary markets, usually 1.0 vs 1.0. The tier 2 brain assigns 1.0 for mutually exclusive currently. + let weight_a = Decimal::ONE; + let weight_b = weight; + + if let Some(profit) = crate::strategies::contradiction::ContradictionStrategy::check( + asset_id, + bid_a, + weight_a, + peer_id, + bid_b, + weight_b, + fee_rate, + total_est_gas_usdc, + ) { + // Sizing follows logic from strategy: TARGET_LIQUIDITY / sum_bids + let sum_bids = bid_a + bid_b; + let consistent_size_shares = arbos_core::constants::TARGET_LIQUIDITY / sum_bids; + + let signal = ArbSignal { + legs: vec![ + TradeAction::Sell { + asset_id, + size: consistent_size_shares, + }, + TradeAction::Sell { + asset_id: peer_id, + size: consistent_size_shares, + }, + ], + expected_profit_usdc: profit, + timestamp: book.timestamp, + }; + + info!( + asset_a = %asset_id, + asset_b = %peer_id, + profit = %profit, + "Routing Contradiction ArbSignal to Bot" + ); + if let Err(e) = bot_tx.send(signal).await { + error!( + "Failed to route Contradiction ArbSignal to Bot channel: {}", + e + ); + } + + // Defer ejection + executed_peer = Some(peer_id); + + // Break after executing one contradiction for this asset to avoid double-spend + // on the same loop tick if it contradicts multiple peers + break; + } + } + + if let Some(peer_id) = executed_peer { + self.orderbook_cache.remove(&asset_id); + self.orderbook_cache.remove(&peer_id); + } + } + /// The Max_Duration filter. Markets resolving $>30$ days out trap capital if the Oracle takes weeks. fn violates_max_duration(&self, asset_id: U256) -> bool { if let Some(end_ts) = self.market_end_timestamps.get(&asset_id) { @@ -619,13 +746,18 @@ impl EngineActor { // 1. Parse Graph Edges let (new_imp_edges, new_imp_rev) = Self::parse_implications(&payload.implications); - let (new_market_outcomes, new_market_expected) = + let (new_partition_outcomes, new_asset_to_partition, new_partition_expected) = Self::parse_partitions(&payload.partitions); + let new_contradictions = Self::parse_contradictions(&payload.contradictions); let new_timestamps = Self::parse_timestamps(&payload.asset_end_timestamps); // 2. Compute Tracked Assets (with Expiry GC and MAX_TRACKED limits) - let new_tracked_assets = - Self::compute_tracked_assets(&new_imp_edges, &new_market_outcomes, &new_timestamps); + let new_tracked_assets = Self::compute_tracked_assets( + &new_imp_edges, + &new_contradictions, + &new_partition_outcomes, + &new_timestamps, + ); // 3. Diff & Command Ingestor self.apply_asset_diffs(&new_tracked_assets, cmd_tx).await; @@ -634,8 +766,10 @@ impl EngineActor { self.tracked_assets = new_tracked_assets; self.implication_edges = new_imp_edges; self.implication_reverse_edges = new_imp_rev; - self.market_outcomes = new_market_outcomes; - self.market_expected_outcome_counts = new_market_expected; + self.contradiction_edges = new_contradictions; + self.partition_outcomes = new_partition_outcomes; + self.asset_to_partition = new_asset_to_partition; + self.partition_expected_outcome_counts = new_partition_expected; self.market_end_timestamps = new_timestamps; debug!("Successfully synchronized in-memory Graph from Tier 2 Brain API."); @@ -656,11 +790,15 @@ impl EngineActor { return Err(anyhow::anyhow!("BRAIN_API_URL must have a valid host")); } - let res = self - .api_client - .get(valid_brain_url.join("state")?) - .send() - .await?; + let mut req = self.api_client.get(valid_brain_url.join("state")?); + + if let Ok(brain_api_key) = std::env::var("ADMIN_API_KEY") + && !brain_api_key.is_empty() + { + req = req.header("X-API-Key", brain_api_key); + } + + let res = req.send().await?; if !res.status().is_success() { return Err(anyhow::anyhow!( @@ -701,48 +839,66 @@ impl EngineActor { (edges, rev_edges) } - fn parse_partitions( - partitions: &[crate::actor::PartitionMapping], - ) -> ( - HashMap>, - HashMap, - ) { + fn parse_contradictions( + contradictions: &[crate::actor::ContradictionMapping], + ) -> HashMap> { + let mut edges: HashMap> = HashMap::new(); + + for cont in contradictions { + if cont.confidence < MIN_CONFIDENCE_THRESHOLD { + continue; + } + match ( + U256::from_str_radix(&cont.asset_a, 10), + U256::from_str_radix(&cont.asset_b, 10), + ) { + (Ok(a_id), Ok(b_id)) => { + let weight = Decimal::ONE; // Using 1.0 for mutually exclusive edges by default + edges.entry(a_id).or_default().push((b_id, weight)); + edges.entry(b_id).or_default().push((a_id, weight)); + } + _ => { + warn!( + "Failed to parse Contradiction U256 IDs, silently dropping: A: {} | B: {}", + cont.asset_a, cont.asset_b + ); + } + } + } + edges + } + + fn parse_partitions(partitions: &[crate::actor::PartitionMapping]) -> PartitionParseResult { let mut outcomes = HashMap::new(); + let mut asset_to_partition = HashMap::new(); let mut expected = HashMap::new(); for part in partitions { if part.confidence < MIN_CONFIDENCE_THRESHOLD { continue; } - match polymarket_client_sdk::types::B256::from_str(&part.condition_id) { - Ok(cond_id) => { - expected.insert(cond_id, part.expected_outcomes_count); - - let mut asset_set = HashSet::new(); - for a_str in &part.assets { - match U256::from_str_radix(a_str, 10) { - Ok(a_id) => { - asset_set.insert(a_id); - } - Err(_) => { - warn!( - "Failed to parse Partition Asset U256 ID, dropping asset: {}", - a_str - ); - } - } + + let partition_id = part.condition_id.clone(); + expected.insert(partition_id.clone(), part.expected_outcomes_count); + + let mut asset_set = HashSet::new(); + for a_str in &part.assets { + match U256::from_str_radix(a_str, 10) { + Ok(a_id) => { + asset_set.insert(a_id); + asset_to_partition.insert(a_id, partition_id.clone()); + } + Err(_) => { + warn!( + "Failed to parse Partition Asset U256 ID, dropping asset: {}", + a_str + ); } - outcomes.insert(cond_id, asset_set); - } - Err(_) => { - warn!( - "Failed to parse Partition Condition B256 ID, dropping partition: {}", - part.condition_id - ); } } + outcomes.insert(partition_id, asset_set); } - (outcomes, expected) + (outcomes, asset_to_partition, expected) } fn parse_timestamps(asset_end_timestamps: &HashMap) -> HashMap { @@ -757,7 +913,8 @@ impl EngineActor { fn compute_tracked_assets( edges: &HashMap>, - outcomes: &HashMap>, + contradictions: &HashMap>, + outcomes: &HashMap>, timestamps: &HashMap, ) -> HashSet { let current_ts = std::time::SystemTime::now() @@ -768,6 +925,12 @@ impl EngineActor { let mut frequency: HashMap = HashMap::new(); Self::tally_edge_frequencies(&mut frequency, edges, timestamps, current_ts); + Self::tally_contradiction_frequencies( + &mut frequency, + contradictions, + timestamps, + current_ts, + ); Self::tally_outcome_frequencies(&mut frequency, outcomes, timestamps, current_ts); Self::sort_and_limit_tracked_assets(frequency) @@ -799,9 +962,29 @@ impl EngineActor { } } + fn tally_contradiction_frequencies( + frequency: &mut HashMap, + edges: &HashMap>, + timestamps: &HashMap, + current_ts: i64, + ) { + for (&p, children) in edges { + for &(c, _) in children { + if p < c { + if Self::is_asset_unexpired(p, timestamps, current_ts) { + *frequency.entry(p).or_insert(0) += 1; + } + if Self::is_asset_unexpired(c, timestamps, current_ts) { + *frequency.entry(c).or_insert(0) += 1; + } + } + } + } + } + fn tally_outcome_frequencies( frequency: &mut HashMap, - outcomes: &HashMap>, + outcomes: &HashMap>, timestamps: &HashMap, current_ts: i64, ) { @@ -981,4 +1164,62 @@ mod tests { assert!(signal.expected_profit_usdc > dec!(0.0)); assert_eq!(signal.legs.len(), 2); } + + #[tokio::test] + async fn test_engine_actor_triggers_contradiction_arb() { + let mut actor = EngineActor::new(); + let (tx, mut rx) = mpsc::channel(10); + + let peer_a_id = U256::from(10_u64); + let peer_b_id = U256::from(20_u64); + + actor + .contradiction_edges + .insert(peer_a_id, vec![(peer_b_id, Decimal::ONE)]); + actor + .contradiction_edges + .insert(peer_b_id, vec![(peer_a_id, Decimal::ONE)]); + + let peer_a_book = NormalizedOrderbook { + asset_id: peer_a_id, + market: B256::default(), + vwap_bid: dec!(0.60), // High bid, P = 0.60 + vwap_ask: dec!(0.65), + bid_untradeable: false, + ask_untradeable: false, + timestamp: 123456, + }; + + let peer_b_book = NormalizedOrderbook { + asset_id: peer_b_id, + market: B256::default(), + vwap_bid: dec!(0.60), // High bid, P = 0.60 (Sum P = 1.20) + vwap_ask: dec!(0.65), + bid_untradeable: false, + ask_untradeable: false, + timestamp: 123457, + }; + + let gas_per_leg_usdc = dec!(0.05); + + // Inject first book, then second book + actor + .handle_book_update(peer_a_book.clone(), &tx, dec!(0.0), gas_per_leg_usdc) + .await; + actor + .handle_book_update(peer_b_book.clone(), &tx, dec!(0.01), gas_per_leg_usdc) + .await; + + // Check cache clearing + assert!(!actor.orderbook_cache.contains_key(&peer_a_id)); + assert!(!actor.orderbook_cache.contains_key(&peer_b_id)); + + // Wait for signal + let signal = rx.recv().await.unwrap(); + + // High ask prices indicate sum > 1.0 (Contradiction). + // Profit should be strictly positive. + assert!(signal.expected_profit_usdc > dec!(0.0)); + assert_eq!(signal.legs.len(), 2); + } } diff --git a/crates/engine/src/strategies/mod.rs b/crates/engine/src/strategies/mod.rs index 33b940c..2190f4b 100644 --- a/crates/engine/src/strategies/mod.rs +++ b/crates/engine/src/strategies/mod.rs @@ -1,7 +1,5 @@ // Note: ContradictionStrategy is implemented as a Tier-3 strategy -// but is not currently wired into the Engine runtime (EngineActor). -// It will be wired once the mapping/source for mutually-exclusive pairs -// is provided by the Tier-2 Brain API. +// and is wired into the Engine runtime (EngineActor). pub mod contradiction; pub mod implication; pub mod partition; diff --git a/services/brain/main.py b/services/brain/main.py index 482b4bf..c1c428c 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -7,11 +7,12 @@ from __future__ import annotations +import hashlib import sys from contextlib import asynccontextmanager from datetime import UTC, datetime - from typing import Any + from fastapi import Depends, FastAPI, HTTPException, Query, Security from fastapi.security import APIKeyHeader from loguru import logger @@ -25,8 +26,12 @@ from models import ( AnalyzePairRequest, AnalyzePairResponse, + BrainStatePayload, GraphStats, HealthResponse, + ImplicationMapping, + PartitionMapping, + ContradictionMapping, Market, MarketRead, Relationship, @@ -478,3 +483,154 @@ async def get_relationships_for_market( async def graph_stats(): """Return high-level statistics about the in-memory relationship graph.""" return await _graph_manager.get_stats() + + +# ── Graph Sync Endpoint ──────────────────────────────────────────────────────── + +@app.get("/state", response_model=BrainStatePayload, dependencies=[Depends(_verify_admin)]) +async def get_brain_state(session: AsyncSession = Depends(get_session)): + """Return the entire synchronized graph state for the Rust Engine to poll.""" + + # 1. Fetch relationships and partition groups + implications_dicts = await _graph_manager.get_all_relationships() + partitions_lists = await _graph_manager.get_partition_groups() + + # Collect all needed condition IDs to optimize the DB query + needed_conditions = set() + for rel in implications_dicts: + needed_conditions.add(rel["parent_condition_id"]) + needed_conditions.add(rel["child_condition_id"]) + for group in partitions_lists: + for cid in group: + needed_conditions.add(cid) + + if not needed_conditions: + return BrainStatePayload( + implications=[], partitions=[], contradictions=[], asset_end_timestamps={} + ) + + # 2. Fetch metadata (tokens and timestamps) + result = await session.execute( + select(Market.condition_id, Market.end_date, Market.clob_token_ids) + .where( + Market.active.is_(True), + Market.closed.is_(False), + Market.condition_id.in_(needed_conditions) + ) + ) + + asset_end_timestamps = {} + condition_to_yes_token = {} + + for condition_id, end_date, clob_token_ids in result.all(): + if not clob_token_ids or len(clob_token_ids) == 0: + continue + + # The YES token is almost always the first element in the array for binary markets + yes_token = clob_token_ids[0] + condition_to_yes_token[condition_id] = yes_token + + if end_date: + try: + # End dates from Polymarket are usually ISO8601 strings + dt = datetime.fromisoformat(end_date.replace("Z", "+00:00")) + asset_end_timestamps[yes_token] = int(dt.timestamp()) + except Exception as e: + (logger.bind(condition_id=condition_id, end_date=end_date, error=str(e)) + .warning("failed_to_parse_end_date")) + + # 3. Implications and Contradictions + implications = [] + contradictions = [] + + for rel in implications_dicts: + parent_token = condition_to_yes_token.get(rel["parent_condition_id"]) + child_token = condition_to_yes_token.get(rel["child_condition_id"]) + + if not parent_token or not child_token: + logger.bind( + relationship_id=rel.get("id"), + parent_condition_id=rel["parent_condition_id"], + child_condition_id=rel["child_condition_id"], + logic_type=rel["logic_type"] + ).debug("skipped_relationship_missing_tokens") + continue + + if rel["logic_type"] == "IMPLIES": + implications.append( + ImplicationMapping( + parent_asset_id=parent_token, + child_asset_id=child_token, + confidence=rel["confidence"], + ) + ) + elif rel["logic_type"] == "MUTUALLY_EXCLUSIVE": + contradictions.append( + ContradictionMapping( + asset_a=parent_token, + asset_b=child_token, + confidence=rel["confidence"], + ) + ) + + # 4. Partitions + partitions = [] + + # Build a lookup for mutual exclusion confidences + me_confidences = {} + for rel in implications_dicts: + if rel["logic_type"] == "MUTUALLY_EXCLUSIVE": + c1, c2 = rel["parent_condition_id"], rel["child_condition_id"] + me_confidences[f"{c1}|{c2}"] = rel["confidence"] + me_confidences[f"{c2}|{c1}"] = rel["confidence"] + + for i, group in enumerate(partitions_lists): + if len(group) >= 2: + group_tokens = [] + for cid in group: + tok = condition_to_yes_token.get(cid) + if tok: + group_tokens.append(tok) + else: + (logger.bind(condition_id=cid, group=group) + .debug("partition_member_missing_token")) + + if len(group_tokens) >= 2: + # Calculate minimum confidence and verify clique + group_confs = [] + is_clique = True + for j in range(len(group)): + for k in range(j + 1, len(group)): + conf = me_confidences.get(f"{group[j]}|{group[k]}") + if conf is not None: + group_confs.append(conf) + else: + is_clique = False + break + if not is_clique: + break + + if not is_clique: + logger.bind(group=group).debug("skipped_partition_not_clique") + continue + + partition_confidence = min(group_confs) if group_confs else 1.0 + + # Hash the sorted assets to get a deterministic B256-like condition ID + hash_input = "v1_" + "_".join(sorted(group_tokens)) + hash_id = hashlib.sha256(hash_input.encode()).hexdigest() + partitions.append( + PartitionMapping( + condition_id=f"0x{hash_id}", + expected_outcomes_count=len(group_tokens), + assets=group_tokens, + confidence=partition_confidence + ) + ) + + return BrainStatePayload( + implications=implications, + partitions=partitions, + contradictions=contradictions, + asset_end_timestamps=asset_end_timestamps + ) diff --git a/services/brain/models.py b/services/brain/models.py index 6a05e17..f0c6939 100644 --- a/services/brain/models.py +++ b/services/brain/models.py @@ -138,3 +138,33 @@ class HealthResponse(BaseModel): status: str db_connected: bool llm_configured: bool + + +class ImplicationMapping(BaseModel): + """An implication edge between two markets.""" + parent_asset_id: str + child_asset_id: str + confidence: float + + +class PartitionMapping(BaseModel): + """A group of mutually exclusive markets representing options to a question.""" + condition_id: str + expected_outcomes_count: int + assets: list[str] + confidence: float + + +class ContradictionMapping(BaseModel): + """A pair of mutually exclusive markets.""" + asset_a: str + asset_b: str + confidence: float + + +class BrainStatePayload(BaseModel): + """The aggregate graph state polled by the Rust Engine.""" + implications: list[ImplicationMapping] + partitions: list[PartitionMapping] + contradictions: list[ContradictionMapping] + asset_end_timestamps: dict[str, int] diff --git a/services/brain/tests/test_main.py b/services/brain/tests/test_main.py index 1f8d918..b52033c 100644 --- a/services/brain/tests/test_main.py +++ b/services/brain/tests/test_main.py @@ -16,6 +16,7 @@ def _mock_env(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://postgres:password@localhost:5432/arbos") monkeypatch.setenv("WATSONX_APIKEY", "test-key") monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project") + monkeypatch.setenv("ADMIN_API_KEY", "test-admin-key") # Clear cached settings from config import get_settings @@ -82,3 +83,49 @@ async def test_returns_stats(self, client): assert "total_markets" in data assert "total_relationships" in data assert data["total_markets"] == 0 + + +class TestBrainStateEndpoint: + async def test_requires_auth_for_state(self, client): + response = await client.get("/state") + assert response.status_code == 401 + + async def test_returns_state(self, client): + from main import app, get_session + from unittest.mock import AsyncMock + + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.all.return_value = [ + ("cond_a", "2024-12-31T00:00:00Z", ["yes_a", "no_a"]), + ("cond_b", "2024-12-31T00:00:00Z", ["yes_b", "no_b"]) + ] + mock_session.execute.return_value = mock_result + + async def override_get_session_for_state(): + yield mock_session + + app.dependency_overrides[get_session] = override_get_session_for_state + + with patch("main._graph_manager") as mock_gm: + mock_gm.get_all_relationships = AsyncMock(return_value=[ + { + "parent_condition_id": "cond_a", + "child_condition_id": "cond_b", + "logic_type": "IMPLIES", + "confidence": 0.9, + } + ]) + mock_gm.get_partition_groups = AsyncMock(return_value=[]) + + response = await client.get("/state", headers={"X-API-Key": "test-admin-key"}) + assert response.status_code == 200 + data = response.json() + + assert "implications" in data + assert len(data["implications"]) == 1 + assert data["implications"][0]["parent_asset_id"] == "yes_a" + assert data["implications"][0]["child_asset_id"] == "yes_b" + assert "yes_a" in data["asset_end_timestamps"] + + app.dependency_overrides.pop(get_session, None) diff --git a/services/brain/tests/test_models.py b/services/brain/tests/test_models.py index 2e70fa6..4550007 100644 --- a/services/brain/tests/test_models.py +++ b/services/brain/tests/test_models.py @@ -13,6 +13,10 @@ MarketRead, RelationshipOutput, RelationshipRead, + ImplicationMapping, + PartitionMapping, + ContradictionMapping, + BrainStatePayload, ) @@ -144,3 +148,27 @@ def test_basic(self): condition_id_b="0xdef", ) assert req.condition_id_a == "0xabc" + + +class TestBrainStateModels: + def test_implication_mapping(self): + m = ImplicationMapping(parent_asset_id="A", child_asset_id="B", confidence=0.9) + assert m.parent_asset_id == "A" + + def test_partition_mapping(self): + p = PartitionMapping(condition_id="C", expected_outcomes_count=2, + assets=["A", "B"], confidence=1.0) + assert p.expected_outcomes_count == 2 + + def test_contradiction_mapping(self): + c = ContradictionMapping(asset_a="A", asset_b="B", confidence=0.8) + assert c.asset_b == "B" + + def test_brain_state_payload(self): + payload = BrainStatePayload( + implications=[], + partitions=[], + contradictions=[], + asset_end_timestamps={"A": 1234567890} + ) + assert payload.asset_end_timestamps["A"] == 1234567890