Travel_Assistant / modules /info_extractor.py
Eliot0110's picture
improve: 优化模型调用并对各组件升级
af60cba
raw
history blame
1.04 kB
# modules/info_extractor.py
import re
from .config_loader import ConfigLoader
class InfoExtractor:
def __init__(self, config_loader: ConfigLoader):
self.configs = config_loader
def extract(self, user_input: str) -> dict:
"""从用户输入中提取目的地、天数和旅行风格"""
extracted_info = {}
user_lower = user_input.lower()
# 提取目的地
for alias, city_info in self.configs.cities.items():
if alias in user_lower:
extracted_info["destination"] = city_info
break
# 提取天数
match = re.search(r'(\d+)\s*天', user_input)
if match:
extracted_info["duration"] = {"days": int(match.group(1))}
# 提取旅行风格 (persona)
for p_name, p_info in self.configs.personas.items():
if p_info['name'] in user_input or p_name in user_input:
extracted_info["persona"] = p_info
break
return extracted_info