improve: re+llm

modules/info_extractor.py

@@ -1,7 +1,7 @@
 import json
 import re
 from utils.logger import log
-from .ai_model import AIModel 
+from .ai_model import AIModel
 
 class InfoExtractor:
     def __init__(self, ai_model: AIModel):
@@ -15,29 +15,139 @@ class InfoExtractor:
 
     def extract(self, user_message: str) -> dict:
         """从用户消息中提取结构化信息，确保使用确定性解码。"""
+        
+        # 输入验证
+        if not user_message or not isinstance(user_message, str):
+            log.warning("⚠️ 收到无效的用户消息")
+            return {}
+        
+        # 构建prompt
         prompt = self._build_prompt_template(user_message)
         
         # --- 核心修复：强制使用确定性解码以杜绝幻觉 ---
-        # 确保调用AI模型时，使用类似 do_sample=False 或 temperature=0 的参数
-        # 这里我们模拟这个调用，并强调其重要性
         log.info("🧠 使用LLM开始提取信息 (模式: 确定性)")
+        
+        # 注意：ai_model.generate() 方法不支持 do_sample 和 temperature 参数
+        # 需要通过其他方式确保确定性输出
         raw_response = self.ai_model.run_inference(
-            input_type='text',
+            input_type="text",
             formatted_input=None,
             prompt=prompt,
-            temperature=0.0
+            temperature=0.0  # 使用最低温度确保确定性
         )
         
         try:
-            extracted_json = json.loads(raw_response)
+            # 清理响应，提取纯JSON部分
+            cleaned_response = self._clean_json_response(raw_response)
+            extracted_json = json.loads(cleaned_response)
             log.info(f"✅ LLM成功提取并解析JSON: {extracted_json}")
+            
             # 使用新的验证方法
             validated_data = self._validate_and_normalize(extracted_json)
             log.info(f"📊 LLM最终提取结果 (安全处理后): {validated_data}")
             return validated_data
+            
         except (json.JSONDecodeError, TypeError) as e:
             log.error(f"❌ 解析或验证LLM提取的JSON失败: {e}", exc_info=True)
-            return {} # 返回一个空字典而不是列表
+            log.debug(f"🔍 原始响应: {raw_response}")
+            # 尝试备用提取方法
+            return self._fallback_extraction(user_message)
+
+    def _clean_json_response(self, response: str) -> str:
+        """清理LLM响应，提取纯JSON部分"""
+        if not response:
+            return "{}"
+        
+        # 移除可能的markdown代码块标记
+        response = re.sub(r'```json\s*', '', response)
+        response = re.sub(r'```\s*', '', response)
+        
+        # 移除可能的前导文字
+        response = re.sub(r'^[^{]*', '', response)
+        
+        # 查找第一个{和最后一个}
+        start_idx = response.find('{')
+        end_idx = response.rfind('}')
+        
+        if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
+            return response[start_idx:end_idx+1].strip()
+        
+        # 如果找不到有效的JSON结构，返回空对象
+        return "{}"
+
+    def _fallback_extraction(self, user_message: str) -> dict:
+        """基于规则的备用信息提取"""
+        log.info("🔄 使用基于规则的备用提取方法")
+        
+        result = {}
+        message_lower = user_message.lower()
+        
+        # 目的地提取 - 更全面的模式
+        city_patterns = [
+            r'去(\w+)', r'到(\w+)', r'想去(\w+)', r'前往(\w+)',
+            r'旅行(\w+)', r'游(\w+)', r'玩(\w+)', r'访问(\w+)',
+            r'目的地[\s是：:]*(\w+)', r'地方[\s是：:]*(\w+)'
+        ]
+        
+        for pattern in city_patterns:
+            match = re.search(pattern, user_message)
+            if match:
+                city_name = match.group(1)
+                if len(city_name) >= 2 and not city_name.isdigit():
+                    result["destination"] = {"name": city_name}
+                    break
+        
+        # 天数提取 - 更全面的模式
+        day_patterns = [
+            r'(\d+)天', r'(\d+)日', r'玩(\d+)天', r'住(\d+)天',
+            r'(\d+)个天', r'呆(\d+)天', r'待(\d+)天', r'(\d+)天行程'
+        ]
+        
+        for pattern in day_patterns:
+            match = re.search(pattern, user_message)
+            if match:
+                days = int(match.group(1))
+                if 1 <= days <= 365:  # 合理范围检查
+                    result["duration"] = {"days": days}
+                    break
+        
+        # 预算提取 - 更全面的模式
+        budget_patterns = [
+            r'(\d+)元', r'(\d+)块', r'预算(\d+)', r'(\d+)rmb',
+            r'(\d+)人民币', r'花(\d+)', r'费用(\d+)', r'(\d+)万'
+        ]
+        
+        for pattern in budget_patterns:
+            match = re.search(pattern, user_message)
+            if match:
+                amount = int(match.group(1))
+                # 处理"万"的情况
+                if '万' in pattern:
+                    amount *= 10000
+                result["budget"] = {
+                    "type": None,
+                    "amount": amount,
+                    "currency": "RMB"
+                }
+                break
+        
+        # 预算类型识别
+        budget_type_keywords = {
+            'economy': ['经济', '便宜', '省钱', '实惠', '节省'],
+            'comfortable': ['舒适', '中等', '适中', '一般'],
+            'luxury': ['豪华', '奢华', '高端', '贵一点', '不差钱']
+        }
+        
+        for budget_type, keywords in budget_type_keywords.items():
+            if any(keyword in message_lower for keyword in keywords):
+                if "budget" not in result:
+                    result["budget"] = {"type": budget_type, "amount": None, "currency": None}
+                else:
+                    result["budget"]["type"] = budget_type
+                break
+        
+        log.info(f"🛠️ 备用提取结果: {result}")
+        return result
 
     def _validate_and_normalize(self, data: dict) -> dict:
         """
@@ -48,40 +158,108 @@ class InfoExtractor:
             return {}
             
         validated_output = {}
-        for key, schema in self.extraction_schema.items():
-            if key in data and isinstance(data[key], schema["type"]):
-                # 这里可以添加更深层次的字段类型验证
-                # 为简化，我们暂时只验证第一层
-                validated_output[key] = data[key]
+        
+        # 验证destination
+        if "destination" in data:
+            dest_data = data["destination"]
+            if isinstance(dest_data, dict):
+                validated_dest = {}
+                if "name" in dest_data and isinstance(dest_data["name"], str):
+                    name = dest_data["name"].strip()
+                    if name:
+                        validated_dest["name"] = name
+                if "country" in dest_data and isinstance(dest_data["country"], str):
+                    country = dest_data["country"].strip()
+                    if country:
+                        validated_dest["country"] = country
+                if validated_dest:
+                    validated_output["destination"] = validated_dest
+        
+        # 验证duration
+        if "duration" in data:
+            duration_data = data["duration"]
+            if isinstance(duration_data, dict):
+                validated_duration = {}
+                if "days" in duration_data:
+                    days = duration_data["days"]
+                    if isinstance(days, (int, float)) and 1 <= days <= 365:
+                        validated_duration["days"] = int(days)
+                if "description" in duration_data and isinstance(duration_data["description"], str):
+                    desc = duration_data["description"].strip()
+                    if desc:
+                        validated_duration["description"] = desc
+                if validated_duration:
+                    validated_output["duration"] = validated_duration
+        
+        # 验证budget
+        if "budget" in data:
+            budget_data = data["budget"]
+            if isinstance(budget_data, dict):
+                validated_budget = {}
+                
+                # 验证type
+                if "type" in budget_data:
+                    budget_type = budget_data["type"]
+                    if budget_type in ["economy", "comfortable", "luxury"]:
+                        validated_budget["type"] = budget_type
+                
+                # 验证amount
+                if "amount" in budget_data:
+                    amount = budget_data["amount"]
+                    if isinstance(amount, (int, float)) and amount > 0:
+                        validated_budget["amount"] = int(amount)
+                
+                # 验证currency
+                if "currency" in budget_data and isinstance(budget_data["currency"], str):
+                    currency = budget_data["currency"].strip()
+                    if currency:
+                        validated_budget["currency"] = currency
+                
+                # 验证description
+                if "description" in budget_data and isinstance(budget_data["description"], str):
+                    desc = budget_data["description"].strip()
+                    if desc:
+                        validated_budget["description"] = desc
+                
+                if validated_budget:
+                    validated_output["budget"] = validated_budget
         
         if not validated_output:
             log.warning(f"⚠️ 提取的数据 {data} 未通过验证，未发现任何有效字段。")
         
         return validated_output
 
-    def _build_prompt_template(self,user_message: str) -> str:
-        # --- 重点更新：使用更严格的指令和结构化示例 ---
+    def _build_prompt_template(self, user_message: str) -> str:
+        """构建包含多个示例的提取prompt"""
+        
+        # 输��长度控制
+        if len(user_message) > 300:
+            user_message = user_message[:300] + "..."
+        
         return f"""你的任务是且仅是作为文本解析器。
 严格分析用户输入，并以一个纯净、无注释的JSON对象格式返回。
 
 **核心规则:**
-1.  **绝对禁止** 在JSON之外添加任何文本、注释、解释或Markdown标记。你的输出必须从 `{{` 开始，到 `}}` 结束。
-2.  **必须严格遵守** 下方定义的嵌套JSON结构。不要创造新的键，也不要改变层级。
-3.  如果信息未提供，对应的键值必须为 `null`，而不是省略该键。
-4.  如果用户输入与旅行无关（如 "你好"），必须返回一个空的JSON对象: `{{}}`。
+1. **绝对禁止** 在JSON之外添加任何文本、注释、解释或Markdown标记。你的输出必须从 `{{` 开始，到 `}}` 结束。
+2. **必须严格遵守** 下方定义的嵌套JSON结构。不要创造新的键，也不要改变层级。
+3. 如果信息未提供，对应的键值必须为 `null`，而不是省略该键。
+4. 如果用户输入与旅行无关(如 "你好")，必须返回一个空的JSON对象: `{{}}`。
 
 **强制JSON输出结构:**
 {{
   "destination": {{
-    "name": "string or null"
+    "name": "string or null",
+    "country": "string or null"
   }},
   "duration": {{
-    "days": "integer or null"
+    "days": "integer or null",
+    "description": "string or null"
   }},
   "budget": {{
     "type": "string ('economy', 'comfortable', 'luxury') or null",
     "amount": "number or null",
-    "currency": "string or null"
+    "currency": "string or null",
+    "description": "string or null"
   }}
 }}
 
@@ -106,12 +284,12 @@ class InfoExtractor:
 }}
 
 **示例2:**
-用户输入: "计划去西班牙巴塞罗那旅行一周，预算2万元"
+用户输入: "计划去日本东京旅行一周，预算2万元"
 你的输出:
 {{
   "destination": {{
-    "name": "巴塞罗那",
-    "country": "西班牙"
+    "name": "东京",
+    "country": "日本"
   }},
   "duration": {{
     "days": 7,
@@ -125,11 +303,54 @@ class InfoExtractor:
   }}
 }}
 
+**示例3:**
+用户输入: "想要一个经济实惠的巴黎5天行程"
+你的输出:
+{{
+  "destination": {{
+    "name": "巴黎",
+    "country": null
+  }},
+  "duration": {{
+    "days": 5,
+    "description": null
+  }},
+  "budget": {{
+    "type": "economy",
+    "amount": null,
+    "currency": null,
+    "description": "经济实惠"
+  }}
+}}
+
+**示例4:**
+用户输入: "你好"
+你的输出:
+{{}}
+
+**示例5:**
+用户输入: "想去泰国普吉岛度蜜月，10天左右，豪华一点不差钱"
+你的输出:
+{{
+  "destination": {{
+    "name": "普吉岛",
+    "country": "泰国"
+  }},
+  "duration": {{
+    "days": 10,
+    "description": "10天左右"
+  }},
+  "budget": {{
+    "type": "luxury",
+    "amount": null,
+    "currency": null,
+    "description": "豪华一点不差钱"
+  }}
+}}
+
 ---
 **用户输入:**
-`{user_message}`
+{user_message}
 
 **你的输出 (必须是纯JSON):**
-"""
-
-    
+"""