50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
from typing import Optional
|
||
import re
|
||
|
||
from openai import OpenAI
|
||
|
||
from backend.config.settings import settings
|
||
|
||
|
||
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:
|
||
return match.group(1)
|
||
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]:
|
||
client = OpenAI(api_key=settings.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]:
|
||
if not sim_slot:
|
||
return None
|
||
match = re.search(r'1\d{10}', sim_slot)
|
||
if match:
|
||
return match.group(0)
|
||
return None |