Add complete Ouroboros trading system with TDD test suite
Some checks failed
CI / test (push) Has been cancelled
Some checks failed
CI / test (push) Has been cancelled
Implement the full autonomous trading agent architecture: - KIS broker with async API, token refresh, leaky bucket rate limiter, and hash key signing - Gemini-powered decision engine with JSON parsing and confidence threshold enforcement - Risk manager with circuit breaker (-3% P&L) and fat finger protection (30% cap) - Evolution engine for self-improving strategy generation via failure analysis - 35 passing tests written TDD-first covering risk, broker, and brain modules - CI/CD pipeline, Docker multi-stage build, and AI agent context docs Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
152
src/brain/gemini_client.py
Normal file
152
src/brain/gemini_client.py
Normal file
@@ -0,0 +1,152 @@
|
||||
"""Decision engine powered by Google Gemini.
|
||||
|
||||
Constructs prompts from market data, calls Gemini, and parses structured
|
||||
JSON responses into validated TradeDecision objects.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import google.generativeai as genai
|
||||
|
||||
from src.config import Settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
VALID_ACTIONS = {"BUY", "SELL", "HOLD"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TradeDecision:
|
||||
"""Validated decision from the AI brain."""
|
||||
|
||||
action: str # "BUY" | "SELL" | "HOLD"
|
||||
confidence: int # 0-100
|
||||
rationale: str
|
||||
|
||||
|
||||
class GeminiClient:
|
||||
"""Wraps the Gemini API for trade decision-making."""
|
||||
|
||||
def __init__(self, settings: Settings) -> None:
|
||||
self._settings = settings
|
||||
self._confidence_threshold = settings.CONFIDENCE_THRESHOLD
|
||||
genai.configure(api_key=settings.GEMINI_API_KEY)
|
||||
self._model = genai.GenerativeModel(settings.GEMINI_MODEL)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Prompt Construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def build_prompt(self, market_data: dict[str, Any]) -> str:
|
||||
"""Build a structured prompt from market data.
|
||||
|
||||
The prompt instructs Gemini to return valid JSON with action,
|
||||
confidence, and rationale fields.
|
||||
"""
|
||||
return (
|
||||
"You are a professional Korean stock market trading analyst.\n"
|
||||
"Analyze the following market data and decide whether to BUY, SELL, or HOLD.\n\n"
|
||||
f"Stock Code: {market_data['stock_code']}\n"
|
||||
f"Current Price: {market_data['current_price']}\n"
|
||||
f"Orderbook: {json.dumps(market_data['orderbook'], ensure_ascii=False)}\n"
|
||||
f"Foreigner Net Buy/Sell: {market_data['foreigner_net']}\n\n"
|
||||
"You MUST respond with ONLY valid JSON in the following format:\n"
|
||||
'{"action": "BUY"|"SELL"|"HOLD", "confidence": <int 0-100>, "rationale": "<string>"}\n\n'
|
||||
"Rules:\n"
|
||||
"- action must be exactly one of: BUY, SELL, HOLD\n"
|
||||
"- confidence must be an integer from 0 to 100\n"
|
||||
"- rationale must explain your reasoning concisely\n"
|
||||
"- Do NOT wrap the JSON in markdown code blocks\n"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Response Parsing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def parse_response(self, raw: str) -> TradeDecision:
|
||||
"""Parse a raw Gemini response into a TradeDecision.
|
||||
|
||||
Handles: valid JSON, JSON wrapped in markdown code blocks,
|
||||
malformed JSON, missing fields, and invalid action values.
|
||||
|
||||
On any failure, returns a safe HOLD with confidence 0.
|
||||
"""
|
||||
if not raw or not raw.strip():
|
||||
logger.warning("Empty response from Gemini — defaulting to HOLD")
|
||||
return TradeDecision(action="HOLD", confidence=0, rationale="Empty response")
|
||||
|
||||
# Strip markdown code fences if present
|
||||
cleaned = raw.strip()
|
||||
match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", cleaned, re.DOTALL)
|
||||
if match:
|
||||
cleaned = match.group(1).strip()
|
||||
|
||||
try:
|
||||
data = json.loads(cleaned)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Malformed JSON from Gemini — defaulting to HOLD")
|
||||
return TradeDecision(
|
||||
action="HOLD", confidence=0, rationale="Malformed JSON response"
|
||||
)
|
||||
|
||||
# Validate required fields
|
||||
if not all(k in data for k in ("action", "confidence", "rationale")):
|
||||
logger.warning("Missing fields in Gemini response — defaulting to HOLD")
|
||||
return TradeDecision(
|
||||
action="HOLD", confidence=0, rationale="Missing required fields"
|
||||
)
|
||||
|
||||
action = str(data["action"]).upper()
|
||||
if action not in VALID_ACTIONS:
|
||||
logger.warning("Invalid action '%s' from Gemini — defaulting to HOLD", action)
|
||||
return TradeDecision(
|
||||
action="HOLD", confidence=0, rationale=f"Invalid action: {action}"
|
||||
)
|
||||
|
||||
confidence = int(data["confidence"])
|
||||
rationale = str(data["rationale"])
|
||||
|
||||
# Enforce confidence threshold
|
||||
if confidence < self._confidence_threshold:
|
||||
logger.info(
|
||||
"Confidence %d < threshold %d — forcing HOLD",
|
||||
confidence,
|
||||
self._confidence_threshold,
|
||||
)
|
||||
action = "HOLD"
|
||||
|
||||
return TradeDecision(action=action, confidence=confidence, rationale=rationale)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# API Call
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def decide(self, market_data: dict[str, Any]) -> TradeDecision:
|
||||
"""Build prompt, call Gemini, and return a parsed decision."""
|
||||
prompt = self.build_prompt(market_data)
|
||||
logger.info("Requesting trade decision from Gemini")
|
||||
|
||||
try:
|
||||
response = await self._model.generate_content_async(prompt)
|
||||
raw = response.text
|
||||
except Exception as exc:
|
||||
logger.error("Gemini API error: %s", exc)
|
||||
return TradeDecision(
|
||||
action="HOLD", confidence=0, rationale=f"API error: {exc}"
|
||||
)
|
||||
|
||||
decision = self.parse_response(raw)
|
||||
logger.info(
|
||||
"Gemini decision",
|
||||
extra={
|
||||
"action": decision.action,
|
||||
"confidence": decision.confidence,
|
||||
},
|
||||
)
|
||||
return decision
|
||||
Reference in New Issue
Block a user