322 lines
13 KiB
Python
322 lines
13 KiB
Python
"""Gemini AI 클라이언트 모듈 (gemini_client.py)
|
|
|
|
Google Gemini API를 사용하여 주식 분석, 리스크 평가, 매매일지 생성 등을 수행.
|
|
Claude 클라이언트와 동일한 인터페이스를 제공하여 AI_PROVIDER 설정으로 전환 가능.
|
|
|
|
주요 기능:
|
|
- analyze(): 범용 Gemini 분석
|
|
- validate_trading_rule(): 트레이딩 규칙 검증
|
|
- analyze_market(): 시장 분석
|
|
- calculate_fair_price(): 적정가 산출
|
|
- assess_risk(): 리스크 평가
|
|
- optimize_trading_rules(): 전략 최적화
|
|
- generate_trade_journal(): 매매일지 생성
|
|
|
|
공통 헬퍼:
|
|
- _parse_json_response(): AI 응답에서 JSON 추출 (Claude 클라이언트에도 동일 로직 존재)
|
|
"""
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
from config import settings
|
|
from utils import _dumps
|
|
from typing import Dict, Any, List
|
|
import json
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GeminiClient:
|
|
"""Gemini AI 클라이언트 (최신 google-genai SDK 버전)"""
|
|
|
|
def __init__(self):
|
|
if not settings.GEMINI_API_KEY:
|
|
raise ValueError("GEMINI_API_KEY is not set")
|
|
|
|
# 최신 SDK에서는 Client 객체를 생성하여 사용합니다.
|
|
self.client = genai.Client(api_key=settings.GEMINI_API_KEY)
|
|
self.model_id = settings.GEMINI_MODEL # 예: 'gemini-2.0-flash'
|
|
|
|
# GenerationConfig 설정
|
|
self.config = types.GenerateContentConfig(
|
|
max_output_tokens=settings.MAX_TOKENS,
|
|
temperature=settings.TEMPERATURE,
|
|
)
|
|
|
|
async def analyze(
|
|
self,
|
|
system_prompt: str,
|
|
user_message: str,
|
|
tools: List[Dict[str, Any]] = None,
|
|
max_tokens: int = None,
|
|
) -> Dict[str, Any]:
|
|
"""Gemini를 통한 분석 실행"""
|
|
try:
|
|
config_with_system = types.GenerateContentConfig(
|
|
system_instruction=system_prompt,
|
|
max_output_tokens=max_tokens or self.config.max_output_tokens,
|
|
temperature=self.config.temperature,
|
|
)
|
|
|
|
# 비동기 호출
|
|
response = await self.client.models.generate_content(
|
|
model=self.model_id,
|
|
contents=user_message,
|
|
config=config_with_system
|
|
)
|
|
|
|
text = response.text if response.text else ""
|
|
|
|
return {
|
|
"content": [{"type": "text", "text": text}],
|
|
"stop_reason": "end_turn",
|
|
"usage": {
|
|
"input_tokens": response.usage_metadata.prompt_token_count if response.usage_metadata else 0,
|
|
"output_tokens": response.usage_metadata.candidates_token_count if response.usage_metadata else 0,
|
|
}
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Gemini analysis error: {e}")
|
|
raise
|
|
|
|
def _parse_json_response(self, response: Dict[str, Any], fallback: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""응답 텍스트에서 JSON 파싱 (공통 헬퍼)"""
|
|
for content in response["content"]:
|
|
if content["type"] == "text":
|
|
try:
|
|
text = content["text"]
|
|
json_start = text.find('{')
|
|
json_end = text.rfind('}') + 1
|
|
if json_start != -1 and json_end > json_start:
|
|
return json.loads(text[json_start:json_end])
|
|
except (json.JSONDecodeError, Exception):
|
|
return {**fallback, "analysis": content["text"]}
|
|
return fallback
|
|
|
|
async def validate_trading_rule(
|
|
self,
|
|
rule: Dict[str, Any],
|
|
market_data: Dict[str, Any],
|
|
news_data: List[Dict[str, Any]] = None
|
|
) -> Dict[str, Any]:
|
|
"""트레이딩 규칙 검증"""
|
|
system_prompt = """주식 트레이딩 규칙 검증 전문가. 규칙/시장/뉴스를 분석하여 JSON으로만 응답.
|
|
필수 필드: decision(BUY/SELL/HOLD), confidence(0-1), risk_level, recommendation.
|
|
악재(유상증자,감자,소송,실적악화) 시 confidence 낮추고 경고 포함. JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"규칙:{_dumps(rule)}\n시장:{_dumps(market_data)}"
|
|
if news_data:
|
|
user_message += f"\n뉴스:{_dumps(news_data)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"decision": "HOLD", "confidence": 0.0, "analysis": "분석 실패"}
|
|
)
|
|
|
|
async def analyze_market(
|
|
self,
|
|
stock_code: str,
|
|
market_data: Dict[str, Any],
|
|
news: List[Dict[str, Any]] = None
|
|
) -> Dict[str, Any]:
|
|
"""시장 분석"""
|
|
system_prompt = """주식 시장 분석 전문가. JSON으로만 응답.
|
|
필수 필드: trend(UP/DOWN/SIDEWAYS), sentiment(POSITIVE/NEGATIVE/NEUTRAL), recommendation(BUY/SELL/HOLD), confidence(0-1), analysis(간략설명). JSON 외 텍스트 금지."""
|
|
|
|
news_limit = settings.NEWS_LIMIT if settings.TOKEN_SAVE_MODE else 5
|
|
news_text = ""
|
|
if news:
|
|
news_text = "\n뉴스:" + ",".join([item.get('title', '') for item in news[:news_limit]])
|
|
|
|
user_message = f"종목:{stock_code}\n시장:{_dumps(market_data)}{news_text}"
|
|
|
|
response = await self.analyze(system_prompt, user_message)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"trend": "UNKNOWN", "sentiment": "NEUTRAL", "recommendation": "HOLD", "confidence": 0.0}
|
|
)
|
|
|
|
async def calculate_fair_price(
|
|
self,
|
|
stock_code: str,
|
|
market_data: Dict[str, Any],
|
|
financial_data: Dict[str, Any] = None
|
|
) -> Dict[str, Any]:
|
|
"""AI 기반 적정 매수/매도 가격 산출"""
|
|
system_prompt = """주식 가격 전략 전문가. JSON으로만 응답.
|
|
필수 필드: buy_prices(aggressive,neutral,conservative), sell_prices(target1,target2,target3), stop_loss, strategy, confidence(0-1). JSON 외 텍스트 금지."""
|
|
|
|
user_message = (
|
|
f"종목:{stock_code} 현재가:{market_data.get('current_price')} "
|
|
f"고가:{market_data.get('high_price')} 저가:{market_data.get('low_price')} "
|
|
f"거래량:{market_data.get('volume')} 지표:{_dumps(market_data.get('indicators', {}))}"
|
|
)
|
|
|
|
response = await self.analyze(system_prompt, user_message)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"error": "분석 실패", "confidence": 0.0}
|
|
)
|
|
|
|
async def assess_risk(
|
|
self,
|
|
portfolio: Dict[str, Any],
|
|
trade_plan: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
"""리스크 평가 (토큰 최소화)"""
|
|
system_prompt = """리스크 평가. JSON만 출력.
|
|
필수: risk_level(LOW/MEDIUM/HIGH), approval(APPROVED/REJECTED), reason(한줄). JSON 외 금지."""
|
|
|
|
slim_portfolio = {k: portfolio[k] for k in ("strategy", "position_count", "mdd") if k in portfolio}
|
|
user_message = f"포트:{_dumps(slim_portfolio)} 계획:{_dumps(trade_plan)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message, max_tokens=settings.RISK_MAX_TOKENS)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"error": "Failed to assess risk"}
|
|
)
|
|
|
|
async def optimize_trading_rules(
|
|
self,
|
|
stock_code: str,
|
|
market_data: Dict[str, Any],
|
|
news_data: List[Dict[str, Any]] = None
|
|
) -> Dict[str, Any]:
|
|
"""시장 상황에 맞는 트레이딩 규칙 최적화"""
|
|
system_prompt = """알고리즘 트레이딩 전략가. JSON으로만 응답.
|
|
필수 필드: strategy_name, market_condition, parameters(buy_conditions,sell_conditions(take_profit_rate,stop_loss_rate,trailing_stop)), risk_assessment, confidence_score(0-100).
|
|
악재시 보수적, 호재시 적극적 설정. JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"종목:{stock_code}\n시장:{_dumps(market_data)}"
|
|
if news_data:
|
|
user_message += f"\n뉴스:{_dumps(news_data)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"error": "Failed to generate optimized rules"}
|
|
)
|
|
|
|
async def generate_trade_journal(
|
|
self,
|
|
trade_info: Dict[str, Any],
|
|
market_data: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
"""매매일지 자동 작성"""
|
|
system_prompt = """매매일지 작성 전문가. JSON으로만 응답.
|
|
필수 필드: title, content(간략 Markdown), mood(평정심/흥분/공포/아쉬움/만족), tags(배열), score(0-100).
|
|
매매근거, 잘한점/아쉬운점, 개선사항 포함. JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"거래:{_dumps(trade_info)}\n시장:{_dumps(market_data)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"error": "Failed to generate journal"}
|
|
)
|
|
|
|
async def analyze_batch_stocks(
|
|
self,
|
|
stocks: List[Dict[str, Any]],
|
|
) -> Dict[str, Any]:
|
|
"""장 전 배치 분석 — 10종목을 한 번에 받아 각각 매수적합/부적합 판정.
|
|
|
|
max_tokens=2048 (10종목 응답이므로 넉넉하게).
|
|
입력 stocks: [{stock_code, stock_name, current_price, rsi, bb_position,
|
|
volume_ratio, score, news_titles}, ...]
|
|
"""
|
|
system_prompt = """주식 스윙 트레이딩 분석가. 종목 리스트를 분석하여 JSON 배열로 응답.
|
|
각 종목: stock_code, suitable(true/false), confidence(0-1), reason(한줄), risk_factors(배열).
|
|
JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"종목리스트:{_dumps(stocks)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message, max_tokens=2048)
|
|
return self._parse_batch_response(response)
|
|
|
|
def _parse_batch_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""배치 분석 응답 파싱 — JSON 배열 또는 객체 모두 처리.
|
|
|
|
1) 응답 텍스트에서 JSON 배열([...]) 파싱 시도
|
|
2) JSON 객체({...}) 파싱 시도 (AI가 {stocks: [...]} 형태로 응답할 수 있음)
|
|
3) 모두 실패하면 빈 결과 반환
|
|
"""
|
|
for content in response["content"]:
|
|
if content["type"] == "text":
|
|
text = content["text"]
|
|
# 배열 형태 시도
|
|
try:
|
|
arr_start = text.find('[')
|
|
arr_end = text.rfind(']') + 1
|
|
if arr_start != -1 and arr_end > arr_start:
|
|
parsed = json.loads(text[arr_start:arr_end])
|
|
if isinstance(parsed, list):
|
|
return {"stocks": parsed}
|
|
except (json.JSONDecodeError, Exception):
|
|
pass
|
|
# 객체 형태 시도
|
|
try:
|
|
obj_start = text.find('{')
|
|
obj_end = text.rfind('}') + 1
|
|
if obj_start != -1 and obj_end > obj_start:
|
|
return json.loads(text[obj_start:obj_end])
|
|
except (json.JSONDecodeError, Exception):
|
|
pass
|
|
return {"stocks": [], "analysis": text}
|
|
return {"stocks": []}
|
|
|
|
async def check_news_sentiment(
|
|
self,
|
|
stock_code: str,
|
|
stock_name: str,
|
|
news_titles: List[str],
|
|
) -> Dict[str, Any]:
|
|
"""장중 뉴스 감성 분석. max_tokens=256 (짧은 응답)."""
|
|
system_prompt = """주식 뉴스 감성 분석. JSON만 출력.
|
|
필수: sentiment(POSITIVE/NEGATIVE/NEUTRAL), severity(LOW/MEDIUM/HIGH), action(HOLD/SELL/WATCH), reason(한줄).
|
|
JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"종목:{stock_code}({stock_name})\n뉴스제목:{_dumps(news_titles)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message, max_tokens=256)
|
|
return self._parse_json_response(
|
|
response,
|
|
{"sentiment": "NEUTRAL", "severity": "LOW", "action": "HOLD", "reason": "분석 실패"}
|
|
)
|
|
|
|
async def review_daily_trades(
|
|
self,
|
|
trades_summary: Dict[str, Any],
|
|
market_summary: Dict[str, Any],
|
|
) -> Dict[str, Any]:
|
|
"""장 후 당일 매매 복기.
|
|
|
|
max_tokens=1024.
|
|
trades_summary: {total_trades, wins, losses, pnl, strategies_used, notable_trades}
|
|
market_summary: {kospi_change, regime, top_movers}
|
|
"""
|
|
system_prompt = """주식 매매 복기 전문가. 당일 매매를 분석하여 JSON 응답.
|
|
필수: overall_grade(A/B/C/D/F), good_points(배열), bad_points(배열), improvements(배열), strategy_drift(bool), drift_detail(있으면), tomorrow_plan(한줄).
|
|
JSON 외 텍스트 금지."""
|
|
|
|
user_message = f"매매요약:{_dumps(trades_summary)}\n시장요약:{_dumps(market_summary)}"
|
|
|
|
response = await self.analyze(system_prompt, user_message, max_tokens=1024)
|
|
return self._parse_json_response(
|
|
response,
|
|
{
|
|
"overall_grade": "C",
|
|
"good_points": [],
|
|
"bad_points": [],
|
|
"improvements": [],
|
|
"strategy_drift": False,
|
|
"tomorrow_plan": "분석 실패",
|
|
}
|
|
)
|
|
|
|
|
|
gemini_client = GeminiClient()
|