Update main.py

main.py

@@ -45,12 +45,33 @@ else:
 
 
 
+# ---------------------------------------
+# 입력/출력 모델
+# ---------------------------------------
 class NewsRequest(BaseModel):
     url: str
     id: Optional[str] = None
 
+class SummaryInput(BaseModel):
+    url: str
+
+class KeywordsInput(BaseModel):
+    summary: str
+
+class CompanyInput(BaseModel):
+    summary: Optional[str] = None
+    keywords: Optional[List[str]] = None
+
+class SentimentInput(BaseModel):
+    content: str
+
+class PredictInput(BaseModel):
+    keywords: List[Union[str, Dict[str, Any]]]
 
-# 🧠 학습 모델 구조 정의
+
+# ---------------------------------------
+# 간단한 분류기 (기존과 동일)
+# ---------------------------------------
 class SimpleClassifier(torch.nn.Module):
     def __init__(self, input_dim):
         super().__init__()
@@ -60,22 +81,21 @@ class SimpleClassifier(torch.nn.Module):
             torch.nn.Linear(64, 1),
             torch.nn.Sigmoid()
         )
-
     def forward(self, x):
         return self.net(x)
 
-
-
-def fetch_html(url):
+# ---------------------------------------
+# 공통 유틸: HTML, 파서, 썸네일
+# ---------------------------------------
+def fetch_html(url: str) -> bs:
     headers = {"User-Agent": "Mozilla/5.0"}
-    response = requests.get(url, headers=headers, timeout=5)
-    response.raise_for_status()
-    return bs(response.text, "html.parser")
+    resp = requests.get(url, headers=headers, timeout=7)
+    resp.raise_for_status()
+    return bs(resp.text, "html.parser")
 
-def parse_naver(soup):
+def parse_naver(soup: bs):
     title = soup.select_one("h2.media_end_head_headline") or soup.title
     title_text = title.get_text(strip=True) if title else "제목 없음"
-
     time_tag = soup.select_one("span.media_end_head_info_datestamp_time")
     time_text = time_tag.get_text(strip=True) if time_tag else "시간 없음"
 
@@ -85,13 +105,11 @@ def parse_naver(soup):
         content = '\n'.join([p.get_text(strip=True) for p in paragraphs]) if paragraphs else content_area.get_text(strip=True)
     else:
         content = "본문 없음"
-
     return title_text, time_text, content
 
-def parse_daum(soup):
+def parse_daum(soup: bs):
     title = soup.select_one("h3.tit_view") or soup.title
     title_text = title.get_text(strip=True) if title else "제목 없음"
-
     time_tag = soup.select_one("span.num_date")
     time_text = time_tag.get_text(strip=True) if time_tag else "시간 없음"
 
@@ -101,278 +119,295 @@ def parse_daum(soup):
         content = '\n'.join([p.get_text(strip=True) for p in paragraphs]) if paragraphs else content_area.get_text(strip=True)
     else:
         content = "본문 없음"
-
     return title_text, time_text, content
 
-def extract_thumbnail(soup):
+def extract_thumbnail(soup: bs) -> Optional[str]:
     tag = soup.find("meta", property="og:image")
     return tag["content"] if tag and "content" in tag.attrs else None
 
-
-def gemini_use(resultK): 
-    generation_config = genai.GenerationConfig(
-        temperature=1,
-        response_mime_type=None  # 그냥 문자열로 응답받기
-    )
+def parse_article_all(url: str) -> Dict[str, Any]:
+    soup = fetch_html(url)
+    if "naver.com" in url:
+        title, time_str, content = parse_naver(soup)
+    elif "daum.net" in url:
+        title, time_str, content = parse_daum(soup)
+    else:
+        raise HTTPException(status_code=400, detail="지원하지 않는 뉴스 사이트입니다.")
+    thumbnail = extract_thumbnail(soup)
+    return {
+        "title": title,
+        "time": time_str,
+        "content": content,
+        "thumbnail_url": thumbnail,
+        "url": url,
+    }
+
+# ---------------------------------------
+# 회사명 추론 (Gemini)
+# ---------------------------------------
+def gemini_use(text_for_company: str) -> str:
+    generation_config = genai.GenerationConfig(temperature=1)
     model = genai.GenerativeModel('gemini-2.0-flash', generation_config=generation_config)
-
     prompt = f"""
     아래 내용을 참고해서 가장 연관성이 높은 주식 상장 회사 이름 하나만 말해줘.
     다른 설명 없이 회사 이름만 대답해.
 
-    "{resultK}"
+    "{text_for_company}"
     """
-
     response = model.generate_content(prompt)
     try:
-        result_text = response.text.strip()
+        return response.text.strip()
     except AttributeError:
-        result_text = response.candidates[0].content.parts[0].text.strip()
-
-    return result_text
-
-
-
-
+        return response.candidates[0].content.parts[0].text.strip()
+
+# ---------------------------------------
+# 1) 요약 단계
+# ---------------------------------------
+@app.post("/ai/summary")
+def step_summary(inp: SummaryInput):
+    meta = parse_article_all(inp.url)
+    # 너가 기존 resultKeyword를 먼저 쓰고 싶다면 이 한 줄로 대체 가능:
+    # rk = resultKeyword(meta["content"]); return {**meta, "summary": rk["summary"]}
+    summary_text = summarize(meta["content"])
+    return {**meta, "summary": summary_text}
+
+# 2) 키워드 단계
+@app.post("/ai/keywords")
+def step_keywords(inp: KeywordsInput):
+    print("키워드는 옴")
+    try:
+        rk = resultKeyword(inp.summary)
+        return {"keywords": rk["keyword"]}
+    except Exception as e:
+        print("❌ 키워드 추출 오류:", e)
+        return {"keywords": []}
+
+# 3) 관련 상장사 단계
+@app.post("/ai/company")
+def step_company(inp: CompanyInput):
+    if inp.summary:
+        text = inp.summary
+    elif inp.keywords:
+        text = ", ".join(inp.keywords)
+    else:
+        raise HTTPException(status_code=400, detail="summary 또는 keywords 중 하나가 필요합니다.")
+    company = gemini_use(text)
+    return {"company": company}
+
+# 4) 감정 단계
+@app.post("/ai/sentiment")
+def step_sentiment(inp: SentimentInput):
+    s = analyze_sentiment(inp.content)
+    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: label = f"매우 긍정 ({pos*100:.1f}%)"
+        elif pos >= 0.6: label = f"긍정 ({pos*100:.1f}%)"
+        else: label = f"약한 긍정 ({pos*100:.1f}%)"
+    elif max_label == "부정":
+        if neg >= 0.9: label = f"매우 부정 ({neg*100:.1f}%)"
+        elif neg >= 0.6: label = f"부정 ({neg*100:.1f}%)"
+        else: label = f"약한 부정 ({neg*100:.1f}%)"
+    else:
+        label = f"중립 ({neu*100:.1f}%)"
+
+    return {
+        "raw": {"positive": s["positive"], "negative": s["negative"], "neutral": s["neutral"]},
+        "adjusted": {"positive": pos, "negative": neg, "neutral": neu},
+        "sentiment": label
+    }
+
+# 5) 주가 예측 단계
+@app.post("/ai/predict")
+def step_predict(inp: PredictInput):
+    # 🔹 문자열 리스트로 정제 (딕셔너리인 경우 "word" 키 사용)
+    clean_keywords = []
+    for kw in inp.keywords:
+        if isinstance(kw, str):
+            clean_keywords.append(kw)
+        elif isinstance(kw, dict) and "word" in kw:
+            clean_keywords.append(kw["word"])
+
+    if not clean_keywords:
+        raise HTTPException(status_code=400, detail="keywords 리스트가 비어 있습니다.")
+
+    # 🔹 이하 기존 로직 동일
+    keyword_vec = embed_keywords(clean_keywords)
+    input_vec = torch.tensor(keyword_vec, dtype=torch.float32).unsqueeze(0)
+    input_dim = input_vec.shape[1]
+
+    model = SimpleClassifier(input_dim)
+    model.load_state_dict(torch.load("news_model.pt", map_location="cpu"))
+    model.eval()
+
+    with torch.no_grad():
+        prob = model(input_vec).item()
+        pred_label = '📈 상승 (1)' if prob >= 0.5 else '📉 하락 (0)'
+
+    return {"prediction": pred_label, "prob": prob}
+
+# ---------------------------------------
+# 호환용: 기존 parse-news (한방 요청) - 유지
+# ---------------------------------------
 @app.post("/ai/parse-news")
 def parse_news(req: NewsRequest):
     url = req.url.strip()
-    username = req.id.strip() if req.id else None
     try:
-        soup = fetch_html(url)
-
-        if "naver.com" in url:
-            title, time, content = parse_naver(soup)
-        elif "daum.net" in url:
-            title, time, content = parse_daum(soup)
-        else:
-            raise HTTPException(status_code=400, detail="지원하지 않는 뉴스 사이트입니다.")
-
-        thumbnail_url = extract_thumbnail(soup)
-
- 
-        resultK = resultKeyword(content)
-        targetCompany = gemini_use(resultK)
-
-
-        sentiment = analyze_sentiment(content)
-        pos_score = sentiment["positive"]
-        neg_score = sentiment["negative"]
-        net_score = sentiment["neutral"]
-        print("부정:", sentiment["negative"])
-        print("중립:", sentiment["neutral"])
-        print("긍정:", sentiment["positive"])
-
-        # 중립 점수 절반으로 줄이고 나머지를 부정/긍정에 재분배
-        reduced_net = net_score / 2
-        remaining = net_score - reduced_net
-        total_non_neu = neg_score + pos_score
-
+        meta = parse_article_all(url)
+
+        # 키워드/요약(기존 resultKeyword 사용)
+        rk = resultKeyword(meta["content"])
+        targetCompany = gemini_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_score += remaining * (neg_score / total_non_neu)
-            pos_score += remaining * (pos_score / total_non_neu)
+            neg += remaining * (neg / total_non_neu)
+            pos += remaining * (pos / total_non_neu)
         else:
-            neg_score += remaining / 2
-            pos_score += remaining / 2
-
-        net_score = reduced_net
-
-        max_label = max(
-            [("부정", neg_score), ("중립", net_score), ("긍정", pos_score)],
-            key=lambda x: x[1]
-        )[0]
+            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_score >= 0.9:
-                sentiment_label = f"매우 긍정 ({pos_score*100:.1f}%)"
-            elif pos_score >= 0.6:
-                sentiment_label = f"긍정 ({pos_score*100:.1f}%)"
-            else:
-                sentiment_label = f"약한 긍정 ({pos_score*100:.1f}%)"
+            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_score >= 0.9:
-                sentiment_label = f"매우 부정 ({neg_score*100:.1f}%)"
-            elif neg_score >= 0.6:
-                sentiment_label = f"부정 ({neg_score*100:.1f}%)"
-            else:
-                sentiment_label = f"약한 부정 ({neg_score*100:.1f}%)"
+            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"중립 ({net_score*100:.1f}%)"
-                
-#밑에 부분은 모델로 주가 예측 한거임
-    
+            sentiment_label = f"중립 ({neu*100:.1f}%)"
 
-        summary = summarize(content)
-        print(summary)
-
-        _, keywords_2nd = extract_keywords(summary)
+        # 예측
+        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)  # (1, D)
-
-        input_dim = input_vec.shape[1]
-        model = SimpleClassifier(input_dim)
-
+        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 = int(prob >= 0.5)
-
-        
-        prediction = '📈 상승 (1)' if prediction == 1 else '📉 하락 (0)'
-        print(type(prob))
-        print(type(prediction))
+            prediction_label = '📈 상승 (1)' if prob >= 0.5 else '📉 하락 (0)'
 
-        
-        
         return {
+            **meta,
             "message": "뉴스 파싱 및 저장 완료",
-            "title": title,
-            "time": time,
-            "content": content,
-            "thumbnail_url": thumbnail_url,
-            "url": url,
-            "summary": resultK["summary"],
-            "keyword": resultK["keyword"],
+            "summary": rk["summary"],
+            "keyword": rk["keyword"],
             "company": targetCompany,
             "sentiment": sentiment_label,
             "sentiment_value": sentiment_label,
-            "prediction": prediction,
+            "prediction": prediction_label,
             "prob": prob,
         }
 
     except requests.exceptions.RequestException as e:
-        traceback.print_exc()  # 전체 스택트레이스 콘솔에 출력
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"서버 오류: {e}")
     except Exception as e:
-        traceback.print_exc()  # 전체 스택트레이스 콘솔에 출력
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"서버 오류: {e}")
-    
-from fastapi.concurrency import run_in_threadpool # 동기 함수를 비동기처럼 실행하기 위해
-from typing import List, Dict, Any # 반환 타입 명시를 위해 (선택 사항)
-# --- 전역 변수 (서버 시작 시 초기화) ---
+
+# ---------------------------------------
+# 주가 데이터 (기존 유지)
+# ---------------------------------------
 krx_listings: pd.DataFrame = None
 us_listings: pd.DataFrame = None
 translator: Translator = None
 
-# --- 서버 시작 시 실행될 로직 ---
 @app.on_event("startup")
 async def load_initial_data():
-    """
-    서버가 시작될 때 주식 목록과 번역기를 미리 로드하여
-    API 요청마다 반복적으로 로드하는 것을 방지합니다.
-    """
     global krx_listings, us_listings, translator
-    
     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("📊 미국 상장 기업 목록 로딩 완료.")
-        
         translator = Translator()
         logger.info("🌐 번역기 초기화 완료.")
-        
-        logger.info("✅ 모든 초기 데이터 로딩이 성공적으로 완료되었습니다.")
-        
+        logger.info("✅ 초기 데이터 로딩 성공.")
     except Exception as e:
-        logger.error(f"🚨 초기 데이터 로딩 중 심각한 오류 발생: {e}", exc_info=True)
-        # 필요하다면 여기서 서버 실행을 중단시킬 수도 있습니다.
-        # raise RuntimeError("Failed to load initial stock listings.") from e
+        logger.error(f"🚨 초기 데이터 로딩 오류: {e}", exc_info=True)
 
-# --- 핵심 로직 함수 ---
 def get_stock_info(company_name: str) -> Dict[str, str] | None:
-    """
-    회사명을 받아 한국 또는 미국 시장에서 종목 정보를 찾아 반환합니다.
-    (정상 동작하는 스크립트의 로직을 그대로 적용)
-    """
-    # 1. 한국 주식에서 먼저 검색
     kr_match = krx_listings[krx_listings['Name'].str.contains(company_name, case=False, na=False)]
     if not kr_match.empty:
-        stock = kr_match.iloc[0]
-        logger.info(f"KRX에서 '{company_name}' 발견: {stock['Name']} ({stock['Code']})")
-        return {"market": "KRX", "symbol": stock['Code'], "name": stock['Name']}
-
-    # 2. 한국에 없으면 미국 주식에서 검색 (번역기 사용)
+        s = kr_match.iloc[0]
+        return {"market": "KRX", "symbol": s['Code'], "name": s['Name']}
     try:
-        # 번역은 I/O 작업이므로 스레드풀에서 실행하는 것이 더 안전할 수 있으나,
-        # googletrans의 내부 구현에 따라 여기서 직접 호출해도 큰 문제가 없을 수 있습니다.
         company_name_eng = translator.translate(company_name, src='ko', dest='en').text
-        logger.info(f"'{company_name}' -> 영어로 번역: '{company_name_eng}'")
-        
-        # 이름 또는 심볼에서 검색
         us_match = us_listings[
             us_listings['Name'].str.contains(company_name_eng, case=False, na=False) |
             us_listings['Symbol'].str.fullmatch(company_name_eng, case=False)
         ]
-
         if not us_match.empty:
-            stock = us_match.iloc[0]
-            logger.info(f"US에서 '{company_name}' 발견: {stock['Name']} ({stock['Symbol']})")
-            return {"market": "US", "symbol": stock['Symbol'], "name": stock['Name']}
-            
+            s = us_match.iloc[0]
+            return {"market": "US", "symbol": s['Symbol'], "name": s['Name']}
     except Exception as e:
-        logger.error(f"'{company_name}' 번역 또는 미국 주식 검색 중 오류: {e}")
-
-    # 3. 최종적으로 찾지 못한 경우
-    logger.warning(f"'{company_name}'에 해당하는 종목을 찾지 못했습니다.")
+        logger.error(f"번역/미국 주식 검색 오류: {e}")
     return None
 
-def fetch_stock_prices_sync(symbol: str, days: int = 365) -> pd.DataFrame:
-    """
-    지정된 기간 동안의 주가 데이터를 가져옵니다 (동기 함수).
-    """
+def fetch_stock_prices_sync(symbol: str, days: int = 365) -> Optional[pd.DataFrame]:
     end_date = datetime.today()
     start_date = end_date - timedelta(days=days)
-    
-    logger.info(f"FinanceDataReader로 '{symbol}'의 주가 데이터 조회를 시작합니다 ({start_date.date()} ~ {end_date.date()}).")
     try:
         df = fdr.DataReader(symbol, start=start_date, end=end_date)
         if df.empty:
-            logger.warning(f"'{symbol}'에 대한 데이터가 없습니다.")
             return None
         return df
     except Exception as e:
-        logger.error(f"'{symbol}' 데이터 조회 중 오류 발생: {e}", exc_info=True)
+        logger.error(f"'{symbol}' 데이터 조회 오류: {e}", exc_info=True)
         return None
 
-# --- API 엔드포인트 ---
-
 @app.get("/ai/stock-data/by-name",
          summary="회사명으로 최근 1년 주가 데이터 조회 (JSON)",
-         description="회사명(예: 삼성전자, 애플)을 입력받아 최근 1년간의 일별 주가 데이터를 JSON 형식으로 반환합니다.")
-async def get_stock_data_by_name(
-    company_name: str = Query(..., description="조회할 회사명")
-) -> List[Dict[str, Any]]:
-    
+         description="회사명(예: 삼성전자, 애플)을 입력받아 최근 1년간의 일별 주가 데이터를 JSON 형식으로 반환")
+async def get_stock_data_by_name(company_name: str = Query(..., description="조회할 회사명")) -> List[Dict[str, Any]]:
     if not company_name or not company_name.strip():
         raise HTTPException(status_code=400, detail="회사명을 입력해주세요.")
-
     stock_info = await run_in_threadpool(get_stock_info, company_name.strip())
-    
     if not stock_info:
         raise HTTPException(status_code=404, detail=f"'{company_name}'에 해당하는 종목을 찾을 수 없습니다.")
-
     prices_df = await run_in_threadpool(fetch_stock_prices_sync, stock_info['symbol'], 365)
-    
     if prices_df is None or prices_df.empty:
         raise HTTPException(status_code=404, detail=f"'{stock_info['name']}'의 시세 데이터를 찾을 수 없습니다.")
 
-    prices_df.index.name = 'Date'  # 👈 이 줄을 추가하여 인덱스 이름을 명시적으로 설정
+    prices_df.index.name = 'Date'
     prices_df.reset_index(inplace=True)
     prices_df['Date'] = prices_df['Date'].dt.strftime('%Y-%m-%d')
-    
     return prices_df.to_dict(orient='records')
 
-
-
+# ---------------------------------------
+# 실행
+# ---------------------------------------
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)