1"""LLM client for the agent laboratory.23Design rules (CLAUDE.md, "LLM agent laboratory"):4* local OpenAI-compatible inference (mlx_lm.server on MacLustr nodes) is the default; optional BYOK providers are5 configured only through environment variables (never embedded credentials);6* every call is recorded: provider, model, system fingerprint, prompt and context hashes, full messages, temperature,7 seed, response, latency, token usage, estimated cost, cache hit;8* responses are cached content-addressed (sha256 of provider+model+messages+temperature+seed+schema) so a run is9 reproducible bit-for-bit from the cache;10* JSON actions are validated against a JSON schema; invalid output is retried once with a repair message and then11 rejected (the caller treats rejection as "no action" and logs it). LLM output is never accounting truth.12* `MockLLM` is a deterministic *policy surrogate* (rule-based) exposing the same interface; it is labelled as such in13 every log line and must never be reported as LLM inference.14"""15from __future__ import annotations1617import hashlib18import json19import os20import re21import time22from dataclasses import dataclass, field23from pathlib import Path2425import httpx26import jsonschema2728_JSON_RE = re.compile(r"\{.*\}", re.S)293031def _hash(obj) -> str:32 return hashlib.sha256(json.dumps(obj, sort_keys=True, default=str).encode()).hexdigest()333435def extract_json(text: str):36 """Return the first JSON object found in `text` (handles ```json fences), or None."""37 if text is None:38 return None39 t = text.strip()40 t = re.sub(r"^```(?:json)?\s*|\s*```$", "", t, flags=re.S)41 try:42 return json.loads(t)43 except Exception:44 pass45 m = _JSON_RE.search(t)46 if not m:47 return None48 try:49 return json.loads(m.group(0))50 except Exception:51 return None525354@dataclass55class ProviderConfig:56 name: str # 'local-qwen', 'anthropic', ...57 kind: str # 'openai_compat' | 'anthropic' | 'mock'58 model: str59 base_url: str = ''60 api_key_env: str = '' # name of the env var holding the key (optional providers only)61 price_in_per_mtok: float = 0.0 # USD per million input tokens (0 for local)62 price_out_per_mtok: float = 0.063 timeout: float = 120.064 merge_system: bool = False # chat templates that reject a system role (Mistral): prepend system text to the first user turn6566 @property67 def api_key(self):68 return os.environ.get(self.api_key_env, '') if self.api_key_env else ''697071LOCAL_PROVIDERS = {72 'qwen7b': ProviderConfig('qwen7b', 'openai_compat', 'mlx-community/Qwen2.5-7B-Instruct-4bit', 'http://localhost:8081/v1'),73 'mistral7b': ProviderConfig('mistral7b', 'openai_compat', 'mlx-community/Mistral-7B-Instruct-v0.3-4bit', 'http://localhost:8082/v1', merge_system=True),74 'llama3b': ProviderConfig('llama3b', 'openai_compat', 'mlx-community/Llama-3.2-3B-Instruct-4bit', 'http://localhost:8083/v1'),75 'mock': ProviderConfig('mock', 'mock', 'rule-surrogate-v1'),76}777879@dataclass80class LLMResult:81 parsed: dict | None82 raw: str83 valid: bool84 cached: bool85 latency: float86 usage: dict87 cost_usd: float88 model: str89 fingerprint: str90 attempts: int91 error: str = ''929394class LLMClient:95 def __init__(self, provider: ProviderConfig, cache_dir: str | Path = 'results/llm_cache',96 audit_path: str | Path | None = None, run_id: str = 'run'):97 self.p = provider98 self.cache_dir = Path(cache_dir)99 self.cache_dir.mkdir(parents=True, exist_ok=True)100 self.audit_path = Path(audit_path) if audit_path else Path('results/llm_audit') / f'{run_id}.jsonl'101 self.audit_path.parent.mkdir(parents=True, exist_ok=True)102 self.run_id = run_id103 self.n_calls = 0104 self.n_cached = 0105 self.n_invalid = 0106 self.total_tokens = 0107 self.total_cost = 0.0108 self._http = httpx.Client(timeout=provider.timeout)109110 # ------------------------------------------------------------------ raw completion111 def _complete_raw(self, messages, temperature, seed, max_tokens):112 p = self.p113 if p.kind == 'openai_compat':114 if p.merge_system and messages and messages[0]['role'] == 'system':115 sys_txt = messages[0]['content']116 rest = [dict(m) for m in messages[1:]]117 if rest and rest[0]['role'] == 'user':118 rest[0]['content'] = sys_txt + '\n\n' + rest[0]['content']119 messages = rest120 headers = {'Content-Type': 'application/json'}121 if p.api_key:122 headers['Authorization'] = f'Bearer {p.api_key}'123 body = dict(model=p.model, messages=messages, temperature=temperature, max_tokens=max_tokens, seed=seed)124 r = self._http.post(f'{p.base_url}/chat/completions', json=body, headers=headers)125 r.raise_for_status()126 j = r.json()127 text = j['choices'][0]['message']['content']128 usage = j.get('usage', {})129 fp = j.get('system_fingerprint', '')130 return text, usage, fp131 if p.kind == 'anthropic':132 if not p.api_key:133 raise RuntimeError(f'missing credential in env {p.api_key_env}')134 headers = {'x-api-key': p.api_key, 'anthropic-version': '2023-06-01', 'content-type': 'application/json'}135 system = '\n'.join(m['content'] for m in messages if m['role'] == 'system')136 msgs = [m for m in messages if m['role'] != 'system']137 body = dict(model=p.model, max_tokens=max_tokens, temperature=temperature, messages=msgs)138 if system:139 body['system'] = system140 r = self._http.post(f'{p.base_url or "https://api.anthropic.com"}/v1/messages', json=body, headers=headers)141 r.raise_for_status()142 j = r.json()143 text = ''.join(b.get('text', '') for b in j.get('content', []))144 u = j.get('usage', {})145 usage = dict(prompt_tokens=u.get('input_tokens', 0), completion_tokens=u.get('output_tokens', 0))146 return text, usage, j.get('model', '')147 raise ValueError(p.kind)148149 # ------------------------------------------------------------------ structured call150 def complete_json(self, messages: list[dict], schema: dict, temperature: float = 0.0, seed: int = 0,151 max_tokens: int = 300, tags: dict | None = None) -> LLMResult:152 key = _hash(dict(provider=self.p.name, model=self.p.model, messages=messages, temperature=temperature,153 seed=seed, schema=schema))154 cpath = self.cache_dir / f'{key}.json'155 cached = cpath.exists()156 t0 = time.time()157 attempts = 0158 error = ''159 if cached:160 rec = json.loads(cpath.read_text())161 text, usage, fp, attempts = rec['raw'], rec['usage'], rec['fingerprint'], rec['attempts']162 parsed = extract_json(text)163 else:164 msgs = list(messages)165 parsed, text, usage, fp = None, '', {}, ''166 for attempt in range(2):167 attempts += 1168 try:169 text, usage, fp = self._complete_raw(msgs, temperature, seed, max_tokens)170 except Exception as e: # network / provider error171 error = f'{type(e).__name__}: {e}'172 text = ''173 break174 parsed = extract_json(text)175 if parsed is not None:176 try:177 jsonschema.validate(parsed, schema)178 break179 except jsonschema.ValidationError as e:180 error = f'schema: {e.message}'181 parsed = None182 msgs = msgs + [{'role': 'assistant', 'content': text},183 {'role': 'user', 'content': 'Your reply was not a valid JSON object matching the schema. '184 'Reply with ONLY the JSON object, no prose.'}]185 if not error.startswith(('ConnectError', 'ReadTimeout', 'HTTPStatusError', 'RemoteProtocolError', 'ConnectTimeout')):186 cpath.write_text(json.dumps(dict(raw=text, usage=usage, fingerprint=fp, attempts=attempts, error=error)))187 valid = parsed is not None188 if valid:189 try:190 jsonschema.validate(parsed, schema)191 except jsonschema.ValidationError as e:192 valid, error = False, f'schema: {e.message}'193 latency = time.time() - t0194 ptok = int(usage.get('prompt_tokens', 0) or 0)195 ctok = int(usage.get('completion_tokens', 0) or 0)196 cost = (ptok * self.p.price_in_per_mtok + ctok * self.p.price_out_per_mtok) / 1e6197 self.n_calls += 1198 self.n_cached += int(cached)199 self.n_invalid += int(not valid)200 self.total_tokens += ptok + ctok201 self.total_cost += cost202 res = LLMResult(parsed if valid else None, text, valid, cached, latency, usage, cost, self.p.model, fp, attempts, error)203 self._audit(messages, schema, temperature, seed, res, tags or {})204 return res205206 def _audit(self, messages, schema, temperature, seed, res: LLMResult, tags):207 rec = dict(ts=time.time(), run_id=self.run_id, provider=self.p.name, kind=self.p.kind, model=self.p.model,208 fingerprint=res.fingerprint, surrogate=(self.p.kind == 'mock'),209 prompt_hash=_hash(messages[0]['content'] if messages else ''), context_hash=_hash(messages),210 schema_hash=_hash(schema), messages=messages, temperature=temperature, seed=seed,211 response=res.raw, parsed=res.parsed, valid=res.valid, error=res.error, cached=res.cached,212 latency=res.latency, usage=res.usage, cost_usd=res.cost_usd, attempts=res.attempts, tags=tags)213 with self.audit_path.open('a') as f:214 f.write(json.dumps(rec, default=str) + '\n')215216 def stats(self):217 return dict(provider=self.p.name, model=self.p.model, calls=self.n_calls, cached=self.n_cached,218 invalid=self.n_invalid, tokens=self.total_tokens, cost_usd=self.total_cost)219220221class MockLLM(LLMClient):222 """Deterministic rule-based *surrogate* with the LLM interface. Produces a mean-variance-like order from the223 observation JSON embedded in the last user message. Labelled surrogate=True in every audit record."""224225 def __init__(self, **kw):226 super().__init__(LOCAL_PROVIDERS['mock'], **kw)227228 def _complete_raw(self, messages, temperature, seed, max_tokens):229 obs = extract_json(messages[-1]['content']) or {}230 price = float(obs.get('price', 1.0))231 val = float(obs.get('own_valuation', price))232 h = float(obs.get('holdings', 0.0))233 cash = float(obs.get('cash', 0.0))234 target = float(obs.get('target_holdings', h))235 q = target - h236 side = 'buy' if q > 0 else ('sell' if q < 0 else 'hold')237 qty = abs(round(q, 2))238 limit = round(val * (0.995 if side == 'buy' else 1.005), 4)239 act = dict(side=side, quantity=qty, order_type='limit', limit_price=limit, borrow=0.0,240 rationale='surrogate: move toward target holdings')241 return json.dumps(act), dict(prompt_tokens=0, completion_tokens=0), 'mock-surrogate'242