Update main.py

main.py

@@ -25,6 +25,7 @@ from datetime import datetime, timedelta
 from googletrans import Translator
 from starlette.concurrency import run_in_threadpool
 import FinanceDataReader as fdr
+from groq import Groq
 
 app = FastAPI()
 
@@ -33,15 +34,16 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-API_KEY = os.getenv("GEMINI_API_KEY") 
+API_KEY = os.getenv("GEMINI_API_KEY")
+
 
 if not API_KEY:
     # API 키가 없으면 에러를 발생시키거나 경고
-    print("❌ GEMINI_API_KEY 환경 변수가 설정되지 않았습니다.")
+    print("❌ Groq_API_KEY 환경 변수가 설정되지 않았습니다.")
     
 else:
-    genai.configure(api_key=API_KEY)
-    logger.info("✅ Gemini API 설정 완료 (환경 변수 사용)")
+    groq_client = Groq(api_key=API_KEY)
+    logger.info("✅ Groq API 설정 완료 (환경 변수 사용)")
 
 
 
@@ -146,20 +148,44 @@ def parse_article_all(url: str) -> Dict[str, Any]:
 # 회사명 추론 (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"""
-    아래 내용을 참고해서 가장 연관성이 높은 주식 상장 회사 이름 하나만 말해줘.
-    다른 설명 없이 회사 이름만 대답해.
+# 3. 함수 이름 및 내용 변경 (gemini_use -> groq_use)
+def groq_use(text_content: Any) -> str:
+    # 텍스트 추출 및 정제
+    if isinstance(text_content, dict):
+        text_for_ai = text_content.get('summary', '')
+    else:
+        text_for_ai = str(text_content)
+
+    # 프롬프트 구성 (불필요한 특수문자 제거 및 슬라이싱)
+    clean_text = text_for_ai[:500].replace('\n', ' ')
+    prompt = f"상장사 이름을 '회사명' 형식으로 하나만 답해줘: {clean_text}"
 
-    "{text_for_company}"
-    """
-    response = model.generate_content(prompt)
     try:
-        return response.text.strip()
-    except AttributeError:
-        return response.candidates[0].content.parts[0].text.strip()
+        chat_completion = groq_client.chat.completions.create(
+            messages=[{
+                "role": "user",
+                "content": prompt
+            }],
+            model="llama-3.3-70b-versatile",
+        )
+        return chat_completion.choices[0].message.content.strip()
+
+    except Exception as e:
+        # 에러 객체를 직접 출력하지 말고 repr()을 사용해 ASCII 충돌 방지
+        logger.error(f"🚨 Groq 호출 실패: {repr(e)}")
+        return "추출 실패"
+    
+def get_stock_info(company_info: str) -> Dict[str, str] | None:
+    try:
+        if krx_listings is not None:
+            # 포함 관계 확인
+            for _, row in krx_listings.iterrows():
+                if row['Name'] in company_info:
+                    return {"market": "KRX", "symbol": row['Code'], "name": row['Name']}
+    except Exception as e:
+        logger.error(f"주식 검색 에러: {e}")
+    return None
+
 
 # ---------------------------------------
 # 1) 요약 단계
@@ -192,7 +218,7 @@ def step_company(inp: CompanyInput):
         text = ", ".join(inp.keywords)
     else:
         raise HTTPException(status_code=400, detail="summary 또는 keywords 중 하나가 필요합니다.")
-    company = gemini_use(text)
+    company = groq_use(text)
     return {"company": company}
 
 # 4) 감정 단계
@@ -270,7 +296,7 @@ def parse_news(req: NewsRequest):
 
         # 키워드/요약(기존 resultKeyword 사용)
         rk = resultKeyword(meta["content"])
-        targetCompany = gemini_use(rk)  # 텍스트 변환은 f-string 내부에서 처리됨
+        targetCompany = groq_use(rk)  # 텍스트 변환은 f-string 내부에서 처리됨
 
         # 감정(기존 로직)
         s = analyze_sentiment(meta["content"])
@@ -360,22 +386,30 @@ async def load_initial_data():
         logger.error(f"🚨 초기 데이터 로딩 오류: {e}", exc_info=True)
 
 def get_stock_info(company_name: str) -> Dict[str, str] | None:
-    kr_match = krx_listings[krx_listings['Name'].str.contains(company_name, case=False, na=False)]
-    if not kr_match.empty:
-        s = kr_match.iloc[0]
-        return {"market": "KRX", "symbol": s['Code'], "name": s['Name']}
     try:
-        company_name_eng = translator.translate(company_name, src='ko', dest='en').text
-        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:
-            s = us_match.iloc[0]
-            return {"market": "US", "symbol": s['Symbol'], "name": s['Name']}
+        # 1. 한국 주식(KRX) 검색
+        if krx_listings is not None:
+            kr_match = krx_listings[krx_listings['Name'].str.contains(company_name, case=False, na=False)]
+            if not kr_match.empty:
+                s = kr_match.iloc[0]
+                return {"market": "KRX", "symbol": s['Code'], "name": s['Name']}
+
+        # 2. 미국 주식 검색 (번역기 없이 이름으로 직접 검색)
+        if us_listings is not None:
+            # 영어 이름이 포함되어 있을 수 있으므로 대소문자 무시하고 검색
+            us_match = us_listings[
+                us_listings['Name'].str.contains(company_name, case=False, na=False) |
+                us_listings['Symbol'].str.fullmatch(company_name, case=False)
+            ]
+            if not us_match.empty:
+                s = us_match.iloc[0]
+                return {"market": "US", "symbol": s['Symbol'], "name": s['Name']}
+
     except Exception as e:
-        logger.error(f"번역/미국 주식 검색 오류: {e}")
+        logger.error(f"주식 종목 검색 중 오류 발생: {e}")
+    
     return None
+    
 
 def fetch_stock_prices_sync(symbol: str, days: int = 365) -> Optional[pd.DataFrame]:
     end_date = datetime.today()