""" Async vLLM inference client. Wraps the OpenAI-compatible endpoint served by vLLM on AMD MI300X. All agent calls go through here, batched via asyncio.gather(). """ import asyncio import json import time from dataclasses import dataclass import openai @dataclass class LLMResponse: """Parsed response from the LLM.""" action: str # "buy", "sell", "hold", "cancel" price: float quantity: int raw_text: str latency_ms: float success: bool orders: list[dict] = None # Added for multiple orders # Default hold response for when LLM returns garbage HOLD_RESPONSE = LLMResponse( action="hold", price=0.0, quantity=0, raw_text="fallback_hold", latency_ms=0.0, success=False, orders=[] ) def parse_llm_output(raw: str) -> dict | None: """ Parse the LLM's JSON output. Returns dict or None on failure. Handles common LLM failure modes: markdown wrapping, trailing text. """ text = raw.strip() # Strip markdown code fences if present if text.startswith("```"): lines = text.split("\n") # Remove first and last lines (```json and ```) lines = [l for l in lines if not l.strip().startswith("```")] text = "\n".join(lines).strip() try: data = json.loads(text) except json.JSONDecodeError: # Try to find JSON object in the text start = text.find("{") end = text.rfind("}") if start != -1 and end != -1 and end > start: try: data = json.loads(text[start:end + 1]) except json.JSONDecodeError: return None else: return None if "orders" in data and isinstance(data["orders"], list): parsed_orders = [] for o in data["orders"]: action = o.get("action", "").lower() if action not in ("buy", "sell", "hold", "cancel"): continue if action in ("hold", "cancel"): parsed_orders.append({"action": action, "price": 0.0, "quantity": 0}) else: try: price = float(o.get("price", 0)) quantity = int(o.get("quantity", 0)) if price > 0 and quantity > 0: parsed_orders.append({"action": action, "price": round(price, 2), "quantity": min(quantity, 10)}) except (ValueError, TypeError): continue return {"orders": parsed_orders} # Validate required fields action = data.get("action", "").lower() if action not in ("buy", "sell", "hold", "cancel"): return None if action in ("hold", "cancel"): return {"action": action, "price": 0.0, "quantity": 0} try: price = float(data.get("price", 0)) quantity = int(data.get("quantity", 0)) except (ValueError, TypeError): return None if price <= 0 or quantity <= 0: return None # Clamp quantity to spec max quantity = min(quantity, 10) return {"action": action, "price": round(price, 2), "quantity": quantity} class VLLMClient: """ Async client for vLLM's OpenAI-compatible API. Usage: client = VLLMClient(base_url="http://localhost:8000/v1") responses = await client.batch_infer(requests) """ def __init__( self, base_url: str = "http://localhost:8000/v1", api_key: str = "EMPTY", model: str = "Qwen/Qwen2.5-7B-Instruct", max_tokens: int = 64, temperature: float = 0.8, ): # Auto-fix common URL mistakes (like missing /v1 for Ollama/vLLM) base_url = base_url.strip() if base_url and not base_url.endswith("/v1") and not base_url.endswith("/v1/"): if base_url.endswith("/"): base_url += "v1" else: base_url += "/v1" print(f"INFO: Initializing VLLMClient - Model: {model} | Base URL: {base_url}") self.client = openai.AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=10.0) self.model = model self.max_tokens = max_tokens self.temperature = temperature self.error_count = 0 # Track consecutive errors async def infer(self, system_prompt: str, user_message: str) -> LLMResponse: """Single inference call. Returns parsed LLMResponse.""" t0 = time.perf_counter() max_retries = 3 for attempt in range(max_retries): try: response = await self.client.chat.completions.create( model=self.model, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_message}, ], response_format={"type": "json_object"}, max_tokens=self.max_tokens, temperature=self.temperature, ) raw_text = response.choices[0].message.content or "" latency_ms = (time.perf_counter() - t0) * 1000 self.error_count = 0 # Reset on success parsed = parse_llm_output(raw_text) if parsed is None: return LLMResponse( action="hold", price=0.0, quantity=0, raw_text=raw_text, latency_ms=latency_ms, success=False, orders=[] ) if "orders" in parsed: return LLMResponse( action="orders", price=0.0, quantity=0, raw_text=raw_text, latency_ms=latency_ms, success=True, orders=parsed["orders"] ) return LLMResponse( action=parsed["action"], price=parsed["price"], quantity=parsed["quantity"], raw_text=raw_text, latency_ms=latency_ms, success=True, orders=[] ) except Exception as e: error_str = str(e).lower() if "429" in error_str or "rate limit" in error_str or "503" in error_str: if attempt < max_retries - 1: sleep_time = 2 ** attempt print(f"API Rate Limited. Retrying in {sleep_time}s...") await asyncio.sleep(sleep_time) continue latency_ms = (time.perf_counter() - t0) * 1000 print(f"LLM API Error for {self.model}: {e}") self.error_count += 1 return LLMResponse( action="hold", price=0.0, quantity=0, raw_text=f"ERROR: {e}", latency_ms=latency_ms, success=False, orders=[] ) async def batch_infer( self, requests: list[tuple[str, str, str]] ) -> dict[str, LLMResponse]: """ Batch inference for multiple agents concurrently. Args: requests: list of (agent_id, system_prompt, user_message) Returns: dict mapping agent_id → LLMResponse """ async def _call(agent_id: str, sys_prompt: str, user_msg: str): resp = await self.infer(sys_prompt, user_msg) return agent_id, resp tasks = [_call(aid, sp, um) for aid, sp, um in requests] results = await asyncio.gather(*tasks, return_exceptions=True) output: dict[str, LLMResponse] = {} for result in results: if isinstance(result, Exception): # Shouldn't happen since infer() catches exceptions, but be safe continue agent_id, response = result output[agent_id] = response return output