diff --git a/backend/services/sms_service.py b/backend/services/sms_service.py index dbec524..b1922c5 100644 --- a/backend/services/sms_service.py +++ b/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