Update main.py

main.py

@@ -538,27 +538,35 @@ async def produce_progress(analysis_id: str, user_id: str | None, percent: int,
 
 def get_stock_info(company_name: str) -> Dict[str, str] | None:
     try:
-        # 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']}
+        name = (company_name or "").strip()
+        if not name:
+            return None
+
+        # 1) KRX
+        if krx_listings is not None and not krx_listings.empty:
+            # code 컬럼 자동 선택
+            code_col = 'Code' if 'Code' in krx_listings.columns else ('Symbol' if 'Symbol' in krx_listings.columns else None)
+            if code_col and 'Name' in krx_listings.columns:
+                kr_match = krx_listings[krx_listings['Name'].str.contains(name, case=False, na=False)]
+                if not kr_match.empty:
+                    s = kr_match.iloc[0]
+                    code = str(s[code_col]).zfill(6)
+                    return {"market": "KRX", "symbol": code, "name": s['Name']}
+
+        # 2) US
+        if us_listings is not None and not us_listings.empty:
+            if 'Name' in us_listings.columns and 'Symbol' in us_listings.columns:
+                us_match = us_listings[
+                    us_listings['Name'].str.contains(name, case=False, na=False) |
+                    us_listings['Symbol'].str.fullmatch(name, case=False)
+                ]
+                if not us_match.empty:
+                    s = us_match.iloc[0]
+                    return {"market": "US", "symbol": str(s['Symbol']), "name": s['Name']}
 
     except Exception as e:
-        logger.error(f"주식 종목 검색 중 오류 발생: {e}")
-    
+        logger.error(f"주식 종목 검색 중 오류 발생: {repr(e)}", exc_info=True)
+
     return None