🐛 fix a bug.

backend/services/sms_service.py

@@ -2,11 +2,25 @@ import re
 from typing import Optional, Tuple
 from datetime import datetime, timedelta
 from sqlmodel import select, Session
-
+from openai import OpenAI
+from backend.config.settings import settings
 from backend.models.sms import SMSRecord, engine
 
+
 def extract_code_with_context(sms_content: str) -> Optional[str]:
-    """从短信内容中提取验证码"""
+    # 判断AI相关的三个环境变量是否都有值
+    # 如果有值，调用AI接口
+    # 如果没有值，调用正则表达式
+    if settings.ai_base_url != "" and settings.ai_api_key != "" and settings.ai_model != "":
+        return _extract_code_with_ai(sms_content)
+    else:
+        return _extract_code_with_reg(sms_content)
+
+
+def _extract_code_with_reg(sms_content: str) -> Optional[str]:
+    """
+    从短信内容中提取验证码
+    """
     pattern = r'(?:验证码|auth|code)[^0-9]{0,20}(\d{4,8})'
     match = re.search(pattern, sms_content, re.IGNORECASE)
     if match:
@@ -14,6 +28,29 @@ def extract_code_with_context(sms_content: str) -> Optional[str]:
     fallback = re.findall(r'(\d{4,8})', sms_content)
     return fallback[0] if fallback else None
 
+
+def _extract_code_with_ai(sms_content: str) -> Optional[str]:
+    """
+    通过AI提取短信中的验证码
+    """
+    client = OpenAI(api_key=settings.ai_api_key, base_url=settings.ai_base_url)
+    response = client.chat.completions.create(
+        messages=[
+            {
+                "role": "system",
+                "content": "你是一个验证码提取助手，我会告诉你短信内容，请将其中的验证码提取给我，只需要返回验证码即可，如果没有验证码，请返回null，不需要返回任何其他无关的内容。",
+            },
+            {
+                "role": "user",
+                "content": sms_content,
+            }
+        ],
+        model=settings.ai_model,
+    )
+    code = response.choices[0].message.content
+    return code if code != "null" else None
+
+
 def extract_phone_from_sim_slot(sim_slot: Optional[str]) -> Optional[str]:
     """从SIM卡信息中提取手机号"""
     if not sim_slot:
@@ -24,31 +61,31 @@ def extract_phone_from_sim_slot(sim_slot: Optional[str]) -> Optional[str]:
         return match.group(0)
     return None
 
+
 def query_latest_code(phone_number: str, platform_keyword: Optional[str]) -> Tuple[Optional[str], Optional[SMSRecord]]:
     """查询最新的验证码"""
     ten_minutes_ago = datetime.utcnow() - timedelta(minutes=10)
-    
+
     with Session(engine) as session:
         query = select(SMSRecord).where(
             SMSRecord.receive_time >= ten_minutes_ago
         )
-        
 
         primary_query = query.where(SMSRecord.phone_number == phone_number)
         if platform_keyword:
             primary_query = primary_query.where(SMSRecord.sms.contains(platform_keyword))
-        
+
         primary_records = session.exec(primary_query.order_by(SMSRecord.receive_time.desc()).limit(5)).all()
-        
+
         if not primary_records:
             fallback_query = query.where(SMSRecord.sim_slot.contains(phone_number))
             if platform_keyword:
                 fallback_query = fallback_query.where(SMSRecord.sms.contains(platform_keyword))
-            
+
             primary_records = session.exec(fallback_query.order_by(SMSRecord.receive_time.desc()).limit(5)).all()
-        
+
         for record in primary_records:
             if record.extracted_code:
                 return record.extracted_code, record
-        
+
         return None, None