Update main.py

main.py

@@ -4,10 +4,7 @@ import uvicorn
 from pydantic import BaseModel
 import requests
 from bs4 import BeautifulSoup as bs
-import mysql.connector
 import os
-import google.genai as genai
-import json
 from util.keywordExtract import *
 from typing import Optional,List, Dict, Any, Union
 import pandas as pd
@@ -26,6 +23,11 @@ from datetime import datetime, timedelta
 from starlette.concurrency import run_in_threadpool
 import FinanceDataReader as fdr
 from groq import Groq
+import asyncio
+import json
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+import ssl
+
 
 app = FastAPI()
 
@@ -44,6 +46,77 @@ else:
     groq_client = Groq(api_key=API_KEY)
     logger.info(":white_check_mark: Groq API 설정 완료 (환경 변수 사용)")
 
+KAFKA_BOOTSTRAP = os.getenv(
+    "KAFKA_BOOTSTRAP",
+    "newsnake-kafka-lsm71103186-f353.i.aivencloud.com:11897"
+)
+
+KAFKA_TOPIC = os.getenv("KAFKA_TOPIC", "news-analyze")
+KAFKA_GROUP_ID = os.getenv("KAFKA_GROUP_ID", "ai-analyzer-group")
+KAFKA_PROGRESS_TOPIC = os.getenv("KAFKA_PROGRESS_TOPIC", "analysis-progress")
+KAFKA_DONE_TOPIC = os.getenv("KAFKA_DONE_TOPIC", "analysis-done")
+
+KAFKA_CA_FILE = os.getenv("KAFKA_CA_FILE", "ca.pem")
+KAFKA_CERT_FILE = os.getenv("KAFKA_CERT_FILE", "service.cert")
+KAFKA_KEY_FILE = os.getenv("KAFKA_KEY_FILE", "service.key")
+
+producer = None
+consumer = None
+consumer_task = None
+
+def build_ssl_context():
+    ctx = ssl.create_default_context(cafile=KAFKA_CA_FILE)
+    ctx.load_cert_chain(certfile=KAFKA_CERT_FILE, keyfile=KAFKA_KEY_FILE)
+    return ctx
+
+SSL_CONTEXT = build_ssl_context()
+
+@app.on_event("startup")
+async def start_kafka():
+    global producer, consumer, consumer_task
+
+    producer = AIOKafkaProducer(
+        bootstrap_servers=KAFKA_BOOTSTRAP,
+        security_protocol="SSL",
+        ssl_context=SSL_CONTEXT,
+    )
+    await producer.start()
+    logger.info("[KAFKA] producer started (SSL)")
+
+    consumer = AIOKafkaConsumer(
+        KAFKA_TOPIC,
+        bootstrap_servers=KAFKA_BOOTSTRAP,
+        group_id=KAFKA_GROUP_ID,
+        enable_auto_commit=True,
+        auto_offset_reset="latest",
+        security_protocol="SSL",
+        ssl_context=SSL_CONTEXT,
+    )
+    await consumer.start()
+    logger.info("[KAFKA] consumer started (SSL)")
+
+    consumer_task = asyncio.create_task(consume_loop())
+
+@app.on_event("shutdown")
+async def stop_kafka():
+    global producer, consumer, consumer_task
+    if consumer_task:
+        consumer_task.cancel()
+        try:
+            await consumer_task
+        except asyncio.CancelledError:
+            pass
+
+    if consumer:
+        await consumer.stop()
+        logger.info("[KAFKA] consumer stopped")
+
+    if producer:
+        await producer.stop()
+        logger.info("[KAFKA] producer stopped")
+
+
+
 # ---------------------------------------
 # 입력/출력 모델
 # ---------------------------------------
@@ -154,7 +227,16 @@ def groq_use(text_content: Any) -> str:
 
     # 프롬프트 구성 (불필요한 특수문자 제거 및 슬라이싱)
     clean_text = text_for_ai[:500].replace('\n', ' ')
-    prompt = f"상장사 이름을 '회사명' 형식으로 하나만 답해줘: {clean_text}"
+    prompt = f'''제공되는 뉴스 본문을 읽고, 뉴스와 가장 연관성이 높은 기업 
+                 현재 주식 시장(KOSPI, KOSDAQ 등)에 상장된 기업의 이름 하나만로 출력해줘.
+                [제약 사항]
+                뉴스 본문과 가장 연관이 된 회사일 것
+                꼭 하나의 회사를 추출할 것
+                없음이라고 표시하지 말 것
+                상장되지 않은 일반 단체, 정부 기관, 비상장사는 제외할 것.
+                FinanceDataReader 이 라이브러리에 존재하는 회사만 추출할 것.
+                설명 없이 회사 이름만 나열할 것.
+                뉴스에 언급된 맥락상 '기업'임이 확실한 것만 포함할 것 : {clean_text}'''
 
     try:
         chat_completion = groq_client.chat.completions.create(
@@ -281,104 +363,189 @@ def step_predict(inp: PredictInput):
 
     return {"prediction": pred_label, "prob": prob}
 
+
+
 # ---------------------------------------
 # 호환용: 기존 parse-news (한방 요청) - 유지
 # ---------------------------------------
+def analyze_news_sync(
+    url: str,
+    user_id: str | None = None,
+    progress_cb=None,  # ✅ 추가
+) -> Dict[str, Any]:
+
+    def emit(percent: int, stage: str, message: str):
+        if progress_cb:
+            try:
+                progress_cb(percent, stage, message)
+            except Exception:
+                pass
+
+    emit(0, "START", "분석 시작")
+
+    # 1) 기사 파싱
+    emit(5, "PARSING", "뉴스 파싱 중...")
+    meta = parse_article_all(url)
+    emit(15, "PARSING", "뉴스 파싱 완료")
+
+    # 2) 요약/키워드(1차) (네가 원래 하던 resultKeyword)
+    emit(25, "SUMMARY", "요약/키워드 생성 중...")
+    rk = resultKeyword(meta["content"])
+    emit(35, "SUMMARY", "요약/키워드 생성 완료")
+
+    # 3) 회사 추론
+    emit(45, "COMPANY", "관련 회사 분석 중...")
+    targetCompany = groq_use(rk)
+    emit(55, "COMPANY", "관련 회사 분석 완료")
+
+    # 4) 감성 분석
+    emit(65, "SENTIMENT", "감정 분석 중...")
+    s = analyze_sentiment(meta["content"])
+    emit(75, "SENTIMENT", "감정 분석 완료")
+
+    # (원래 감성 후처리 로직 그대로)
+    pos, neg, neu = s["positive"], s["negative"], s["neutral"]
+
+    reduced_net = neu / 2
+    remaining = neu - reduced_net
+    total_non_neu = neg + pos
+    if total_non_neu > 0:
+        neg += remaining * (neg / total_non_neu)
+        pos += remaining * (pos / total_non_neu)
+    else:
+        neg += remaining / 2
+        pos += remaining / 2
+    neu = reduced_net  # ✅ 원래 코드에 있었던 거 유지해야 함
+
+    max_label = max([("부정", neg), ("중립", neu), ("긍정", pos)], key=lambda x: x[1])[0]
+    if max_label == "긍정":
+        if pos >= 0.9:
+            sentiment_label = f"매우 긍정 ({pos*100:.1f}%)"
+        elif pos >= 0.6:
+            sentiment_label = f"긍정 ({pos*100:.1f}%)"
+        else:
+            sentiment_label = f"약한 긍정 ({pos*100:.1f}%)"
+    elif max_label == "부정":
+        if neg >= 0.9:
+            sentiment_label = f"매우 부정 ({neg*100:.1f}%)"
+        elif neg >= 0.6:
+            sentiment_label = f"부정 ({neg*100:.1f}%)"
+        else:
+            sentiment_label = f"약한 부정 ({neg*100:.1f}%)"
+    else:
+        sentiment_label = f"중립 ({neu*100:.1f}%)"
+
+    # 5) (네 원래 코드 유지) summary_text / keywords_2nd / clean_keywords
+    emit(82, "KEYWORDS", "키워드 추출(2차) 중...")
+    summary_text = rk.get("summary") or summarize(meta["content"])
+    _, keywords_2nd = extract_keywords(summary_text)
+    clean_keywords = [kw for kw, _ in keywords_2nd]
+    emit(88, "KEYWORDS", "키워드 추출 완료")
+
+    # 6) 임베딩 + 예측
+    emit(92, "PREDICT", "주가 예측 중...")
+    keyword_vec = embed_keywords(clean_keywords)
+    input_vec = torch.tensor(keyword_vec, dtype=torch.float32).unsqueeze(0)
+
+    model = SimpleClassifier(input_vec.shape[1])
+    model.load_state_dict(torch.load("news_model.pt", map_location="cpu"))
+    model.eval()
+    with torch.no_grad():
+        prob = model(input_vec).item()
+        prediction_label = "📈 상승 (1)" if prob >= 0.5 else "📉 하락 (0)"
+    emit(98, "PREDICT", "주가 예측 완료")
+
+    emit(100, "DONE", "분석 완료")
+
+    # ✅ 리턴 키는 “원래 네 함수랑 최대한 동일하게”
+    return {
+        **meta,
+        "message": "뉴스 파싱 및 저장 완료",
+        "summary": rk.get("summary"),      # 원래: rk["summary"]
+        "keyword": rk.get("keyword"),      # 원래: rk["keyword"]
+        "company": targetCompany,
+        "sentiment": sentiment_label,
+        "sentiment_value": sentiment_label,
+        "prediction": prediction_label,
+        "prob": prob,
+        "userId": user_id,
+    }
+
+
+
 @app.post("/ai/parse-news")
 def parse_news(req: NewsRequest):
-    url = req.url.strip()
     try:
-        meta = parse_article_all(url)
-
-        # 키워드/요약(기존 resultKeyword 사용)
-        rk = resultKeyword(meta["content"])
-        targetCompany = groq_use(rk)  # 텍스트 변환은 f-string 내부에서 처리됨
-
-        # 감정(기존 로직)
-        s = analyze_sentiment(meta["content"])
-        pos, neg, neu = s["positive"], s["negative"], s["neutral"]
-        print("부정:", neg)
-        print("중립:", neu)
-        print("긍정:", pos)
-
-        reduced_net = neu / 2
-        remaining = neu - reduced_net
-        total_non_neu = neg + pos
-        if total_non_neu > 0:
-            neg += remaining * (neg / total_non_neu)
-            pos += remaining * (pos / total_non_neu)
-        else:
-            neg += remaining / 2
-            pos += remaining / 2
-        neu = reduced_net
-
-        max_label = max([("부정", neg), ("중립", neu), ("긍정", pos)], key=lambda x: x[1])[0]
-        if max_label == "긍정":
-            if pos >= 0.9: sentiment_label = f"매우 긍정 ({pos*100:.1f}%)"
-            elif pos >= 0.6: sentiment_label = f"긍정 ({pos*100:.1f}%)"
-            else: sentiment_label = f"약한 긍정 ({pos*100:.1f}%)"
-        elif max_label == "부정":
-            if neg >= 0.9: sentiment_label = f"매우 부정 ({neg*100:.1f}%)"
-            elif neg >= 0.6: sentiment_label = f"부정 ({neg*100:.1f}%)"
-            else: sentiment_label = f"약한 부정 ({neg*100:.1f}%)"
-        else:
-            sentiment_label = f"중립 ({neu*100:.1f}%)"
-
-        # 예측
-        summary_text = rk.get("summary") or summarize(meta["content"])
-        _, keywords_2nd = extract_keywords(summary_text)
-        clean_keywords = [kw for kw, _ in keywords_2nd]
-        keyword_vec = embed_keywords(clean_keywords)
-        input_vec = torch.tensor(keyword_vec, dtype=torch.float32).unsqueeze(0)
-        model = SimpleClassifier(input_vec.shape[1])
-        model.load_state_dict(torch.load("news_model.pt", map_location="cpu"))
-        model.eval()
-        with torch.no_grad():
-            prob = model(input_vec).item()
-            prediction_label = '📈 상승 (1)' if prob >= 0.5 else '📉 하락 (0)'
-
-        return {
-            **meta,
-            "message": "뉴스 파싱 및 저장 완료",
-            "summary": rk["summary"],
-            "keyword": rk["keyword"],
-            "company": targetCompany,
-            "sentiment": sentiment_label,
-            "sentiment_value": sentiment_label,
-            "prediction": prediction_label,
-            "prob": prob,
-        }
-
-    except requests.exceptions.RequestException as e:
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"서버 오류: {e}")
+        return analyze_news_sync(req.url.strip(), user_id=req.id)
     except Exception as e:
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"서버 오류: {e}")
 
+
 # ---------------------------------------
 # 주가 데이터 (기존 유지)
 # ---------------------------------------
 krx_listings: pd.DataFrame = None
 us_listings: pd.DataFrame = None
 
-
 @app.on_event("startup")
 async def load_initial_data():
     global krx_listings, us_listings
-    logger.info("✅ 서버 시작: 초기 데이터 로딩을 시작합니다...")
-    try:
-        krx_listings = await run_in_threadpool(fdr.StockListing, 'KRX')
-        logger.info("📊 한국 상장 기업 목록 로딩 완료.")
-        nasdaq = await run_in_threadpool(fdr.StockListing, 'NASDAQ')
-        nyse = await run_in_threadpool(fdr.StockListing, 'NYSE')
-        amex = await run_in_threadpool(fdr.StockListing, 'AMEX')
-        us_listings = pd.concat([nasdaq, nyse, amex], ignore_index=True)
-        logger.info("📊 미국 상장 기업 목록 로딩 완료.")
-        logger.info("🌐 번역기 초기화 완료.")
-        logger.info("✅ 초기 데이터 로딩 성공.")
-    except Exception as e:
-        logger.error(f"🚨 초기 데이터 로딩 오류: {e}", exc_info=True)
+    file_path_kr = "krx_listings.csv"
+    file_path_ns = "nas_listings.csv"
+    
+    # --- 1. 한국 시장 로딩 ---
+    if os.path.exists(file_path_kr):
+        # dtype={'Code': str} 를 설정해야 '005930'이 '5930'이 되지 않습니다.
+        krx_listings = pd.read_csv(file_path_kr, dtype={'Code': str})
+        logger.info("💾 로컬 파일에서 KRX 목록을 불러왔습니다.")
+    else:
+        try:
+            krx_listings = await run_in_threadpool(fdr.StockListing, 'KRX')
+            # 한글 깨짐 방지를 위해 utf-8-sig 권장
+            krx_listings.to_csv(file_path_kr, index=False, encoding='utf-8-sig')
+            logger.info("📊 KRX 데이터를 새로 받아 저장했습니다.")
+        except Exception as e:
+            logger.error(f"🚨 KRX 데이터 로딩 실패: {e}")
+            krx_listings = pd.DataFrame(columns=['Code', 'Name']) # 빈 데이터프레임 할당
+
+    # --- 2. 미국 시장 로딩 ---
+    if os.path.exists(file_path_ns):
+        us_listings = pd.read_csv(file_path_ns, dtype={'Symbol': str})
+        logger.info("💾 로컬 파일에서 US 목록을 불러왔습니다.")
+    else:
+        try:
+            # 여러 시장 데이터를 합칠 때 에러 방지
+            nasdaq = await run_in_threadpool(fdr.StockListing, 'NASDAQ')
+            nyse = await run_in_threadpool(fdr.StockListing, 'NYSE')
+            amex = await run_in_threadpool(fdr.StockListing, 'AMEX')
+            
+            us_listings = pd.concat([nasdaq, nyse, amex], ignore_index=True)
+            us_listings.to_csv(file_path_ns, index=False, encoding='utf-8-sig')
+            logger.info("📊 미국 상장 기업 목록 로딩 완료.")
+        except Exception as e:
+            logger.error(f"🚨 미국 상장사 로딩 실패: {e}")
+            us_listings = pd.DataFrame(columns=['Symbol', 'Name'])
+
+
+async def produce_progress(analysis_id: str, user_id: str | None, percent: int, stage: str, message: str):
+    if not producer:
+        return
+
+    payload = {
+        "analysisId": analysis_id,
+        "userId": user_id,
+        "percent": percent,
+        "stage": stage,
+        "message": message,
+    }
+    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
+
+    await producer.send_and_wait(
+        KAFKA_PROGRESS_TOPIC,
+        key=analysis_id.encode("utf-8"),
+        value=data
+    )
 
 def get_stock_info(company_name: str) -> Dict[str, str] | None:
     try:
@@ -436,6 +603,88 @@ async def get_stock_data_by_name(company_name: str = Query(..., description="조
     prices_df['Date'] = prices_df['Date'].dt.strftime('%Y-%m-%d')
     return prices_df.to_dict(orient='records')
 
+
+# ---------------------------------------
+# Kafka Consumer 루프 추가
+# ---------------------------------------
+
+async def consume_loop():
+    global consumer
+    logger.info(f"[KAFKA] started topic={KAFKA_TOPIC} group={KAFKA_GROUP_ID} bootstrap={KAFKA_BOOTSTRAP}")
+
+    try:
+        # ��� 현재 이벤트 루프를 미리 잡아둠 (threadpool 콜백에서 필요)
+        loop = asyncio.get_running_loop()
+
+        async for msg in consumer:
+            key = msg.key.decode() if msg.key else None
+            raw = msg.value.decode() if msg.value else ""
+
+            try:
+                payload = json.loads(raw)
+            except Exception:
+                logger.warning(f"[KAFKA] invalid json: {raw}")
+                continue
+
+            analysis_id = payload.get("analysisId")
+            url = (payload.get("url") or "").strip()
+            user_id = payload.get("userId")
+
+            if not analysis_id or not url:
+                logger.warning(f"[KAFKA] missing analysisId/url payload={payload}")
+                continue
+
+            logger.info(f"[KAFKA] consume key={key} analysisId={analysis_id} url={url}")
+
+            # ✅ (선택) 여기서 0%를 한 번 보내도 되지만,
+            # analyze_news_sync 내부에서도 0%를 emit 하게 만들었으면 중복될 수 있음.
+            # await produce_progress(analysis_id, user_id, 0, "START", "분석 시작")
+
+            # ✅ threadpool에서 호출될 progress 콜백
+            def progress_cb(percent: int, stage: str, message: str):
+                # threadpool(다른 스레드) -> 현재 이벤트루프에서 코루틴 실행
+                fut = asyncio.run_coroutine_threadsafe(
+                    produce_progress(analysis_id, user_id, percent, stage, message),
+                    loop
+                )
+                # (선택) progress 전송 실패 로그 보고 싶으면:
+                try:
+                    fut.result(timeout=2)
+                except Exception as e:
+                    logger.warning(f"[KAFKA] progress send failed: {e}")
+
+            try:
+                # ✅ 무거운 작업 → threadpool (progress_cb 전달!)
+                result = await run_in_threadpool(analyze_news_sync, url, user_id, progress_cb)
+
+                logger.info(f"[KAFKA] done analysisId={analysis_id} title={result.get('title')}")
+
+                # ✅ analyze_news_sync가 100% DONE까지 emit 한다면 여기서 100% 또 보낼 필요 없음
+                # await produce_progress(analysis_id, user_id, 100, "DONE", "분석 완료")
+
+                # ✅ DONE 이벤트는 analysis-done 토픽으로 (이건 그대로 유지)
+                if producer:
+                    done_payload = json.dumps({
+                        "analysisId": analysis_id,
+                        "userId": user_id,
+                        "result": result
+                    }, ensure_ascii=False).encode("utf-8")
+
+                    await producer.send_and_wait(
+                        KAFKA_DONE_TOPIC,
+                        key=analysis_id.encode("utf-8"),
+                        value=done_payload
+                    )
+                    logger.info(f"[KAFKA] produced done analysisId={analysis_id}")
+
+            except Exception as e:
+                logger.error(f"[KAFKA] analysis failed analysisId={analysis_id}: {repr(e)}", exc_info=True)
+                await produce_progress(analysis_id, user_id, 100, "ERROR", f"오류: {type(e).__name__}")
+
+    finally:
+        logger.info("[KAFKA] stopped")
+
+
 # ---------------------------------------
 # 실행
 # ---------------------------------------