689 lines
25 KiB
Python
689 lines
25 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import re
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
from openai import AsyncOpenAI
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
SQL_DIR = SCRIPT_DIR.parent / "sql"
|
|
|
|
CompleteFn = Callable[[str], Awaitable[dict[str, Any]]]
|
|
|
|
|
|
def env(name: str, default: str | None = None) -> str | None:
|
|
val = os.environ.get(name, default)
|
|
return val if val not in (None, "") else default
|
|
|
|
|
|
def sanitize_url(url: str | None) -> str | None:
|
|
if not url:
|
|
return None
|
|
try:
|
|
parts = urllib.parse.urlsplit(url)
|
|
except ValueError:
|
|
return None
|
|
host = parts.hostname
|
|
if not host:
|
|
return None
|
|
out = f"{parts.scheme}://{host}"
|
|
if parts.port:
|
|
out += f":{parts.port}"
|
|
return out
|
|
|
|
|
|
@dataclass
|
|
class SwarmConfig:
|
|
agents: int = 8
|
|
rounds: int = 2
|
|
peers: int = 2
|
|
concurrency: int = 4
|
|
max_tokens: int = 512
|
|
seed: int = 1337
|
|
limit: int = 50
|
|
max_retries: int = 3
|
|
base_delay: float = 0.5
|
|
reasoning_effort: str = "none"
|
|
opinion_words: int = 80
|
|
include_hosts: bool = False
|
|
dry_run: bool = False
|
|
output: str = "stock_swarm_result.json"
|
|
llm_base_url: str | None = None
|
|
llm_api_key: str | None = None
|
|
llm_model_name: str | None = None
|
|
clickhouse_url: str | None = None
|
|
clickhouse_user: str | None = None
|
|
clickhouse_password: str | None = None
|
|
clickhouse_database: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class CallResult:
|
|
ok: bool
|
|
latency_ms: float
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
total_tokens: int = 0
|
|
content: str | None = None
|
|
error: str | None = None
|
|
|
|
|
|
def validate_config(cfg: SwarmConfig) -> None:
|
|
positive = {
|
|
"agents": cfg.agents,
|
|
"rounds": cfg.rounds,
|
|
"concurrency": cfg.concurrency,
|
|
"max_tokens": cfg.max_tokens,
|
|
"limit": cfg.limit,
|
|
"opinion_words": cfg.opinion_words,
|
|
}
|
|
invalid = [name for name, value in positive.items() if value <= 0]
|
|
if cfg.peers < 0:
|
|
invalid.append("peers")
|
|
if cfg.max_retries < 0:
|
|
invalid.append("max_retries")
|
|
if cfg.base_delay < 0:
|
|
invalid.append("base_delay")
|
|
if invalid:
|
|
raise ValueError(f"Invalid non-positive configuration: {', '.join(invalid)}")
|
|
|
|
|
|
class RetryableHTTP(Exception):
|
|
def __init__(self, status_code: int, message: str = ""):
|
|
super().__init__(message or f"HTTP {status_code}")
|
|
self.status_code = status_code
|
|
|
|
|
|
def is_retryable(exc: Exception) -> bool:
|
|
if isinstance(exc, (asyncio.TimeoutError,)):
|
|
return True
|
|
if isinstance(exc, RetryableHTTP):
|
|
return True
|
|
sc = getattr(exc, "status_code", None)
|
|
if sc in (429, 500, 502, 503, 504):
|
|
return True
|
|
name = type(exc).__name__
|
|
if name in ("APITimeoutError", "ConnectError", "ConnectTimeout",
|
|
"ReadTimeout", "RemoteProtocolError"):
|
|
return True
|
|
return False
|
|
|
|
|
|
def percentile(sorted_xs: list[float], q: float) -> float | None:
|
|
if not sorted_xs:
|
|
return None
|
|
k = max(0, math.ceil(q / 100 * len(sorted_xs)) - 1)
|
|
return round(sorted_xs[min(k, len(sorted_xs) - 1)], 3)
|
|
|
|
|
|
def is_valid_opinion(opinion: dict | None) -> bool:
|
|
return bool(opinion and "parse_error" not in opinion and opinion.get("view"))
|
|
|
|
|
|
def summarize(results: list[CallResult], wall_time_s: float,
|
|
opinions: list[dict | None] | None = None) -> dict[str, Any]:
|
|
lat = sorted(r.latency_ms for r in results)
|
|
requests = len(results)
|
|
transport_successes = sum(1 for r in results if r.ok)
|
|
transport_failures = requests - transport_successes
|
|
valid_opinions = (
|
|
sum(1 for opinion in opinions if is_valid_opinion(opinion))
|
|
if opinions is not None else transport_successes
|
|
)
|
|
invalid_opinions = requests - valid_opinions
|
|
prompt_tokens = sum(r.prompt_tokens for r in results)
|
|
completion_tokens = sum(r.completion_tokens for r in results)
|
|
total_tokens = sum(r.total_tokens for r in results)
|
|
wall = wall_time_s if wall_time_s and wall_time_s > 0 else 0.0
|
|
return {
|
|
"requests": requests,
|
|
"successes": valid_opinions,
|
|
"failures": invalid_opinions,
|
|
"transport_successes": transport_successes,
|
|
"transport_failures": transport_failures,
|
|
"valid_opinions": valid_opinions,
|
|
"invalid_opinions": invalid_opinions,
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": total_tokens,
|
|
"wall_time_s": round(wall_time_s, 4),
|
|
"input_tokens_per_sec": round(prompt_tokens / wall, 2) if wall else 0.0,
|
|
"output_tokens_per_sec": round(completion_tokens / wall, 2) if wall else 0.0,
|
|
"latency_ms": {
|
|
"p50": percentile(lat, 50),
|
|
"p95": percentile(lat, 95),
|
|
"p99": percentile(lat, 99),
|
|
},
|
|
}
|
|
|
|
|
|
def build_topology(n: int, peers: int, seed: int) -> dict[int, list[int]]:
|
|
if n <= 1:
|
|
return {i: [] for i in range(n)}
|
|
k = min(peers, n - 1)
|
|
rng = random.Random(seed)
|
|
topo: dict[int, list[int]] = {}
|
|
for i in range(n):
|
|
others = [j for j in range(n) if j != i]
|
|
rng.shuffle(others)
|
|
topo[i] = sorted(others[:k])
|
|
return topo
|
|
|
|
|
|
def round1_prompt(ticker: str, dossier_json: str, opinion_words: int) -> str:
|
|
return (
|
|
"You are a stock analyst agent. Ticker: "
|
|
f"{ticker}.\nDossier (JSON):\n{dossier_json}\n"
|
|
"Form an independent opinion. Respond ONLY with valid JSON, no markdown. "
|
|
f"Make analysis approximately {opinion_words} words:\n"
|
|
'{"ticker": "<ticker>", "view": "bullish|bearish|neutral", '
|
|
'"score": 0.0-1.0, "confidence": 0.0-1.0, "thesis": "<one sentence>", '
|
|
'"drivers": ["<driver>", "<driver>", "<driver>"], '
|
|
'"risks": ["<risk>", "<risk>"], "analysis": "<synthesis>"}'
|
|
)
|
|
|
|
|
|
def roundN_prompt(ticker: str, dossier_json: str, own_opinion: dict | None,
|
|
peers: list[dict], opinion_words: int) -> str:
|
|
own_block = json.dumps(own_opinion, separators=(",", ":"))
|
|
peer_block = json.dumps(peers, separators=(",", ":"))
|
|
return (
|
|
"You are a stock analyst agent. Ticker: "
|
|
f"{ticker}.\nDossier (JSON):\n{dossier_json}\n"
|
|
"Your previous opinion:\n"
|
|
f"{own_block}\n"
|
|
"Your peers' opinions from the previous round:\n"
|
|
f"{peer_block}\n"
|
|
"Revise your opinion considering peers. Respond ONLY with valid JSON, no markdown. "
|
|
f"Make analysis approximately {opinion_words} words:\n"
|
|
'{"ticker": "<ticker>", "view": "bullish|bearish|neutral", '
|
|
'"score": 0.0-1.0, "confidence": 0.0-1.0, "thesis": "<one sentence>", '
|
|
'"drivers": ["<driver>", "<driver>", "<driver>"], '
|
|
'"risks": ["<risk>", "<risk>"], "analysis": "<synthesis>", '
|
|
'"changed": true|false, "revision_note": "<short>"}'
|
|
)
|
|
|
|
|
|
def parse_opinion(content: str | None) -> dict | None:
|
|
if not content:
|
|
return None
|
|
s = content.strip()
|
|
s = re.sub(r"^```(?:json)?\s*", "", s, flags=re.IGNORECASE)
|
|
s = re.sub(r"\s*```$", "", s)
|
|
i, j = s.find("{"), s.rfind("}")
|
|
if i == -1 or j == -1 or j < i:
|
|
return {"parse_error": "no json object", "raw": s[:200]}
|
|
try:
|
|
return json.loads(s[i:j + 1])
|
|
except json.JSONDecodeError as e:
|
|
return {"parse_error": str(e), "raw": s[i:j + 1][:200]}
|
|
|
|
|
|
def load_sql(name: str) -> str:
|
|
return (SQL_DIR / name).read_text(encoding="utf-8")
|
|
|
|
|
|
def ch_query(sql: str, params: dict[str, Any],
|
|
cfg: SwarmConfig | None = None) -> list[dict]:
|
|
configured_url = cfg.clickhouse_url if cfg else env("CLICKHOUSE_URL")
|
|
if not configured_url:
|
|
raise RuntimeError("CLICKHOUSE_URL is required")
|
|
url = configured_url.rstrip("/") + "/"
|
|
database = (
|
|
cfg.clickhouse_database if cfg and cfg.clickhouse_database
|
|
else env("CLICKHOUSE_DATABASE", "default")
|
|
)
|
|
qs: dict[str, str] = {
|
|
"database": database or "default",
|
|
"default_format": "JSON",
|
|
}
|
|
for k, v in params.items():
|
|
qs[f"param_{k}"] = str(v)
|
|
full = url + "?" + urllib.parse.urlencode(qs)
|
|
user = (
|
|
cfg.clickhouse_user if cfg and cfg.clickhouse_user
|
|
else env("CLICKHOUSE_USER", "default")
|
|
) or "default"
|
|
password = (
|
|
cfg.clickhouse_password if cfg and cfg.clickhouse_password is not None
|
|
else env("CLICKHOUSE_PASSWORD", "")
|
|
) or ""
|
|
token = base64.b64encode(f"{user}:{password}".encode()).decode()
|
|
req = urllib.request.Request(
|
|
full, data=sql.encode("utf-8"), method="POST",
|
|
headers={"Authorization": f"Basic {token}"},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=30) as resp:
|
|
body = resp.read().decode("utf-8")
|
|
except urllib.error.HTTPError as e:
|
|
raise RuntimeError(f"ClickHouse HTTP {e.code}: {e.read().decode('utf-8', 'replace')[:300]}") from None
|
|
except urllib.error.URLError as e:
|
|
raise RuntimeError(f"ClickHouse unreachable: {e}") from e
|
|
return json.loads(body).get("data", [])
|
|
|
|
|
|
def fetch_ticker_universe(limit: int, cfg: SwarmConfig | None = None) -> list[dict]:
|
|
rows = ch_query(load_sql("ticker_universe.sql"), {"limit": limit}, cfg)
|
|
return [
|
|
{
|
|
"ticker": r["ticker"],
|
|
"name": r.get("name"),
|
|
"sector": r.get("sector"),
|
|
"scale_marketcap": r.get("scale_marketcap"),
|
|
}
|
|
for r in rows
|
|
]
|
|
|
|
|
|
def fetch_holding_quarter(cfg: SwarmConfig | None = None) -> str:
|
|
rows = ch_query(load_sql("holding_quarter.sql"), {}, cfg)
|
|
if not rows:
|
|
raise RuntimeError("ClickHouse has no fully filed 13F quarter")
|
|
return rows[0]["calendardate"]
|
|
|
|
|
|
def fetch_dossier(ticker: str, holding_quarter: str,
|
|
cfg: SwarmConfig | None = None) -> dict:
|
|
rows = ch_query(
|
|
load_sql("dossier.sql"),
|
|
{"ticker": ticker, "holding_quarter": holding_quarter},
|
|
cfg,
|
|
)
|
|
if not rows or not rows[0].get("latest_annual_date"):
|
|
raise RuntimeError(f"No fundamentals dossier for {ticker}")
|
|
return rows[0]
|
|
|
|
|
|
async def fetch_live_universe(cfg: SwarmConfig) -> tuple[list[dict], dict[str, Any]]:
|
|
started = time.perf_counter()
|
|
universe_limit = max(cfg.limit, cfg.agents)
|
|
universe_raw, holding_quarter = await asyncio.gather(
|
|
asyncio.to_thread(fetch_ticker_universe, universe_limit, cfg),
|
|
asyncio.to_thread(fetch_holding_quarter, cfg),
|
|
)
|
|
if len(universe_raw) < cfg.agents:
|
|
raise RuntimeError(
|
|
f"Requested {cfg.agents} agents but ClickHouse returned {len(universe_raw)} tickers"
|
|
)
|
|
selected = universe_raw[:cfg.agents]
|
|
ch_sem = asyncio.Semaphore(min(cfg.concurrency, 16))
|
|
|
|
async def load(row: dict) -> dict:
|
|
async with ch_sem:
|
|
dossier = await asyncio.to_thread(
|
|
fetch_dossier, row["ticker"], holding_quarter, cfg
|
|
)
|
|
return {**row, "dossier": dossier}
|
|
|
|
universe = await asyncio.gather(*(load(row) for row in selected))
|
|
return universe, {
|
|
"wall_time_s": round(time.perf_counter() - started, 4),
|
|
"holding_quarter": holding_quarter,
|
|
"dossiers": len(universe),
|
|
}
|
|
|
|
|
|
def synthetic_universe(n: int, seed: int) -> list[dict]:
|
|
rng = random.Random(seed)
|
|
sectors = ["Tech", "Health", "Energy", "Finance", "Consumer"]
|
|
out = []
|
|
for i in range(n):
|
|
ticker = f"SYNTH{seed % 100:02d}{i:02d}"
|
|
dossier = {
|
|
"ticker": ticker,
|
|
"revenue_annual": round(rng.uniform(1e9, 5e10), 0),
|
|
"netinc_annual": round(rng.uniform(-1e9, 8e9), 0),
|
|
"equity_annual": round(rng.uniform(2e9, 3e10), 0),
|
|
"roe": round(rng.uniform(-0.1, 0.4), 4),
|
|
"roe_percentile": round(rng.random(), 4),
|
|
"holder_count": rng.randint(50, 1200),
|
|
"synthetic": True,
|
|
}
|
|
out.append({"ticker": ticker, "name": f"Synthetic Co {i}",
|
|
"sector": rng.choice(sectors), "dossier": dossier})
|
|
return out
|
|
|
|
|
|
class FakeLLM:
|
|
|
|
def __init__(self, seed: int, fail_first: int = 0, delay: float = 0.0):
|
|
self.seed = seed
|
|
self.fail_first = fail_first
|
|
self.delay = delay
|
|
self.calls = 0
|
|
|
|
async def complete(self, prompt: str) -> dict[str, Any]:
|
|
self.calls += 1
|
|
if self.calls <= self.fail_first:
|
|
raise RetryableHTTP(429, "dry-run forced 429")
|
|
if self.delay:
|
|
await asyncio.sleep(self.delay)
|
|
h = hashlib.sha256(f"{self.seed}:{prompt}".encode()).hexdigest()
|
|
views = ["bullish", "bearish", "neutral"]
|
|
view = views[int(h[:2], 16) % 3]
|
|
score = (int(h[2:4], 16) % 100) / 100
|
|
conf = (int(h[4:6], 16) % 100) / 100
|
|
changed = "Your previous opinion:" in prompt
|
|
obj = {
|
|
"view": view, "score": round(score, 3),
|
|
"confidence": round(conf, 3), "thesis": "synthetic thesis",
|
|
"changed": changed,
|
|
"revision_note": "peer-adjusted" if changed else "initial",
|
|
}
|
|
prompt_tokens = 40 + len(prompt) % 60
|
|
completion_tokens = 60 + int(h[:2], 16) % 40
|
|
return {
|
|
"content": json.dumps(obj, separators=(",", ":")),
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
}
|
|
|
|
|
|
def build_complete_fn(cfg: SwarmConfig) -> CompleteFn:
|
|
if cfg.dry_run:
|
|
return FakeLLM(cfg.seed).complete
|
|
if not (cfg.llm_api_key and cfg.llm_base_url and cfg.llm_model_name):
|
|
raise SystemExit("LLM_BASE_URL, LLM_API_KEY, LLM_MODEL_NAME required when not --dry-run")
|
|
client = AsyncOpenAI(
|
|
api_key=cfg.llm_api_key,
|
|
base_url=cfg.llm_base_url,
|
|
max_retries=0,
|
|
timeout=120.0,
|
|
)
|
|
model = cfg.llm_model_name
|
|
max_tokens = cfg.max_tokens
|
|
effort = cfg.reasoning_effort.strip() if cfg.reasoning_effort else ""
|
|
|
|
async def real_complete(prompt: str) -> dict[str, Any]:
|
|
kwargs: dict[str, Any] = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": max_tokens,
|
|
"response_format": {"type": "json_object"},
|
|
}
|
|
if effort:
|
|
kwargs["reasoning_effort"] = effort
|
|
resp = await client.chat.completions.create(**kwargs)
|
|
choice = resp.choices[0] if resp.choices else None
|
|
message = choice.message if choice else None
|
|
content = (message.content if message else None) or ""
|
|
if not content and message:
|
|
content = (message.model_extra or {}).get("reasoning_content", "")
|
|
usage = resp.usage
|
|
return {
|
|
"content": content,
|
|
"prompt_tokens": getattr(usage, "prompt_tokens", 0) or 0,
|
|
"completion_tokens": getattr(usage, "completion_tokens", 0) or 0,
|
|
"total_tokens": getattr(usage, "total_tokens", 0) or 0,
|
|
}
|
|
|
|
return real_complete
|
|
|
|
|
|
async def invoke(complete_fn: CompleteFn, prompt: str, sem: asyncio.Semaphore,
|
|
max_retries: int, base_delay: float) -> CallResult:
|
|
start = time.perf_counter()
|
|
async with sem:
|
|
last_err: str | None = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
res = await complete_fn(prompt)
|
|
latency = (time.perf_counter() - start) * 1000
|
|
return CallResult(
|
|
ok=True, latency_ms=latency,
|
|
prompt_tokens=int(res.get("prompt_tokens", 0)),
|
|
completion_tokens=int(res.get("completion_tokens", 0)),
|
|
total_tokens=int(res.get("total_tokens", 0)),
|
|
content=res.get("content"),
|
|
)
|
|
except Exception as e:
|
|
last_err = f"{type(e).__name__}: {e}"
|
|
if not is_retryable(e) or attempt == max_retries:
|
|
latency = (time.perf_counter() - start) * 1000
|
|
return CallResult(ok=False, latency_ms=latency, error=last_err)
|
|
delay = base_delay * (2 ** attempt)
|
|
if delay:
|
|
await asyncio.sleep(delay + random.uniform(0, max(delay * 0.25, 0.01)))
|
|
latency = (time.perf_counter() - start) * 1000
|
|
return CallResult(ok=False, latency_ms=latency, error=last_err or "unknown")
|
|
|
|
|
|
async def run_agent(idx: int, ticker_info: dict, rnd: int,
|
|
opinions_prev: dict[int, dict], complete_fn: CompleteFn,
|
|
topology: dict[int, list[int]], ticker_info_by_idx: list[dict],
|
|
cfg: SwarmConfig,
|
|
sem: asyncio.Semaphore) -> tuple[int, CallResult, dict | None]:
|
|
ticker = ticker_info["ticker"]
|
|
dossier = ticker_info.get("dossier", {})
|
|
if rnd == 1:
|
|
prompt = round1_prompt(
|
|
ticker, json.dumps(dossier, separators=(",", ":")), cfg.opinion_words
|
|
)
|
|
else:
|
|
peer_idxs = topology[idx]
|
|
peers = [
|
|
{
|
|
"peer_ticker": ticker_info_by_idx[j]["ticker"],
|
|
"opinion": opinions_prev[j],
|
|
}
|
|
for j in peer_idxs
|
|
if j in opinions_prev
|
|
]
|
|
prompt = roundN_prompt(
|
|
ticker,
|
|
json.dumps(dossier, separators=(",", ":")),
|
|
opinions_prev.get(idx),
|
|
peers,
|
|
cfg.opinion_words,
|
|
)
|
|
cr = await invoke(complete_fn, prompt, sem, cfg.max_retries, cfg.base_delay)
|
|
opinion = parse_opinion(cr.content) if cr.ok else None
|
|
return idx, cr, opinion
|
|
|
|
|
|
async def run_swarm(cfg: SwarmConfig) -> dict[str, Any]:
|
|
return await run_swarm_with_complete(cfg, build_complete_fn(cfg))
|
|
|
|
|
|
async def run_swarm_with_complete(cfg: SwarmConfig, complete_fn: CompleteFn) -> dict[str, Any]:
|
|
validate_config(cfg)
|
|
if cfg.dry_run:
|
|
setup_started = time.perf_counter()
|
|
universe = synthetic_universe(cfg.agents, cfg.seed)
|
|
setup = {
|
|
"wall_time_s": round(time.perf_counter() - setup_started, 4),
|
|
"holding_quarter": None,
|
|
"dossiers": len(universe),
|
|
}
|
|
else:
|
|
universe, setup = await fetch_live_universe(cfg)
|
|
tickers = universe[:cfg.agents]
|
|
n = len(tickers)
|
|
topology = build_topology(n, cfg.peers, cfg.seed)
|
|
sem = asyncio.Semaphore(cfg.concurrency)
|
|
|
|
rounds_out: list[dict] = []
|
|
round_metrics: list[dict] = []
|
|
all_results: list[CallResult] = []
|
|
all_opinions: list[dict | None] = []
|
|
errors: list[dict] = []
|
|
opinions_prev: dict[int, dict] = {}
|
|
overall_start = time.perf_counter()
|
|
|
|
for rnd in range(1, cfg.rounds + 1):
|
|
opinions_input = opinions_prev.copy()
|
|
rstart = time.perf_counter()
|
|
tasks = [
|
|
asyncio.create_task(run_agent(i, tickers[i], rnd, opinions_input,
|
|
complete_fn, topology, tickers, cfg, sem))
|
|
for i in range(n)
|
|
]
|
|
results = await asyncio.gather(*tasks)
|
|
wall = time.perf_counter() - rstart
|
|
round_results = [cr for _, cr, _ in results]
|
|
round_opinions = [opinion for _, _, opinion in results]
|
|
all_results.extend(round_results)
|
|
all_opinions.extend(round_opinions)
|
|
round_metrics.append(summarize(round_results, wall, round_opinions))
|
|
agents_out = []
|
|
for idx, cr, opinion in results:
|
|
opinion_ok = cr.ok and is_valid_opinion(opinion)
|
|
peer_opinions_received = sum(
|
|
1 for peer_idx in topology[idx] if peer_idx in opinions_input
|
|
) if rnd > 1 else 0
|
|
entry = {
|
|
"agent_id": idx,
|
|
"ticker": tickers[idx]["ticker"],
|
|
"ok": opinion_ok,
|
|
"request_ok": cr.ok,
|
|
"opinion_ok": opinion_ok,
|
|
"peer_opinions_received": peer_opinions_received,
|
|
"latency_ms": round(cr.latency_ms, 3),
|
|
"prompt_tokens": cr.prompt_tokens,
|
|
"completion_tokens": cr.completion_tokens,
|
|
"total_tokens": cr.total_tokens,
|
|
"opinion": opinion,
|
|
"error": cr.error,
|
|
}
|
|
agents_out.append(entry)
|
|
if opinion_ok:
|
|
opinions_prev[idx] = opinion
|
|
if not cr.ok:
|
|
errors.append({
|
|
"round": rnd,
|
|
"agent_id": idx,
|
|
"ticker": tickers[idx]["ticker"],
|
|
"kind": "request",
|
|
"error": cr.error,
|
|
})
|
|
elif not opinion_ok:
|
|
errors.append({
|
|
"round": rnd,
|
|
"agent_id": idx,
|
|
"ticker": tickers[idx]["ticker"],
|
|
"kind": "opinion_parse",
|
|
"error": (opinion or {}).get("parse_error", "missing view"),
|
|
})
|
|
rounds_out.append({"round": rnd, "wall_time_s": round(wall, 4), "agents": agents_out})
|
|
|
|
total_wall = time.perf_counter() - overall_start
|
|
return {
|
|
"config": {
|
|
"mode": "dry-run" if cfg.dry_run else "live",
|
|
"model_name": cfg.llm_model_name,
|
|
"reasoning_effort": cfg.reasoning_effort,
|
|
"agents": n,
|
|
"rounds": cfg.rounds,
|
|
"peers": cfg.peers,
|
|
"concurrency": cfg.concurrency,
|
|
"max_tokens": cfg.max_tokens,
|
|
"opinion_words": cfg.opinion_words,
|
|
"seed": cfg.seed,
|
|
"limit": cfg.limit,
|
|
"max_retries": cfg.max_retries,
|
|
"llm_base_url": sanitize_url(cfg.llm_base_url) if cfg.include_hosts else None,
|
|
"clickhouse_url": sanitize_url(cfg.clickhouse_url) if cfg.include_hosts else None,
|
|
},
|
|
"topology": [
|
|
{
|
|
"agent_id": idx,
|
|
"ticker": tickers[idx]["ticker"],
|
|
"peers": [tickers[j]["ticker"] for j in topology[idx]],
|
|
}
|
|
for idx in range(n)
|
|
],
|
|
"metrics": {
|
|
"setup": setup,
|
|
"total": summarize(all_results, total_wall, all_opinions),
|
|
"rounds": round_metrics,
|
|
},
|
|
"rounds": rounds_out,
|
|
"errors": errors,
|
|
}
|
|
|
|
|
|
def build_config_from_env(args: argparse.Namespace) -> SwarmConfig:
|
|
cfg = SwarmConfig(
|
|
agents=args.agents, rounds=args.rounds, peers=args.peers,
|
|
concurrency=args.concurrency, max_tokens=args.max_tokens,
|
|
seed=args.seed, limit=args.limit, max_retries=args.max_retries,
|
|
base_delay=args.base_delay, reasoning_effort=args.reasoning_effort,
|
|
opinion_words=args.opinion_words, include_hosts=args.include_hosts,
|
|
dry_run=args.dry_run, output=args.output,
|
|
llm_base_url=env("LLM_BASE_URL"), llm_api_key=env("LLM_API_KEY"),
|
|
llm_model_name=env("LLM_MODEL_NAME"),
|
|
clickhouse_url=env("CLICKHOUSE_URL"), clickhouse_user=env("CLICKHOUSE_USER"),
|
|
clickhouse_password=env("CLICKHOUSE_PASSWORD"),
|
|
clickhouse_database=env("CLICKHOUSE_DATABASE"),
|
|
)
|
|
if not cfg.dry_run:
|
|
missing = [k for k, v in {
|
|
"CLICKHOUSE_URL": cfg.clickhouse_url,
|
|
"LLM_BASE_URL": cfg.llm_base_url,
|
|
"LLM_API_KEY": cfg.llm_api_key,
|
|
"LLM_MODEL_NAME": cfg.llm_model_name,
|
|
}.items() if not v]
|
|
if missing:
|
|
raise SystemExit(f"Missing required env (not --dry-run): {', '.join(missing)}")
|
|
return cfg
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
p = argparse.ArgumentParser(description="Stock-agent P2P throughput PoC")
|
|
p.add_argument("--agents", type=int, default=8)
|
|
p.add_argument("--rounds", type=int, default=2)
|
|
p.add_argument("--peers", type=int, default=2)
|
|
p.add_argument("--concurrency", type=int, default=4)
|
|
p.add_argument("--max-tokens", type=int, default=512)
|
|
p.add_argument("--seed", type=int, default=1337)
|
|
p.add_argument("--limit", type=int, default=50, help="universe fetch cap")
|
|
p.add_argument("--max-retries", type=int, default=3)
|
|
p.add_argument("--base-delay", type=float, default=0.5)
|
|
p.add_argument("--reasoning-effort", default="none",
|
|
help="none|low|medium|high; empty string to omit the param")
|
|
p.add_argument("--opinion-words", type=int, default=80)
|
|
p.add_argument("--include-hosts", action="store_true")
|
|
p.add_argument("--output", default="stock_swarm_result.json")
|
|
p.add_argument("--dry-run", action="store_true")
|
|
args = p.parse_args(argv)
|
|
cfg = build_config_from_env(args)
|
|
try:
|
|
validate_config(cfg)
|
|
except ValueError as exc:
|
|
p.error(str(exc))
|
|
output = asyncio.run(run_swarm(cfg))
|
|
output_path = Path(cfg.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
output_path.write_text(json.dumps(output, indent=2), encoding="utf-8")
|
|
total = output["metrics"]["total"]
|
|
print(f"agents={output['config']['agents']} rounds={cfg.rounds} "
|
|
f"requests={total['requests']} ok={total['successes']} "
|
|
f"fail={total['failures']} total_tokens={total['total_tokens']} "
|
|
f"in_tok/s={total['input_tokens_per_sec']} "
|
|
f"out_tok/s={total['output_tokens_per_sec']} "
|
|
f"p50={total['latency_ms']['p50']}ms p95={total['latency_ms']['p95']}ms "
|
|
f"-> {cfg.output}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|