目录

一、文本型疾病描述数据的获取

二、大模型工具提取疾病结构化数据及效果对比

(1)deepseek版(付费,可跳)

(2)阿里云百炼版(有免费额度)

(3)效果对比

二、接入Neo4j和TuGraph

(1)Neo4j

(2)TuGraph


一、文本型疾病描述数据的获取

通过爬虫软件在国内公开的健康网站爬取疾病的症状

推荐的爬虫软件(免费):

八爪鱼:https://www.bazhuayu.com

后裔:https://www.houyicaiji.com

医疗网站推荐:

39:https://www.39.net

复禾:https://dise.fh21.com.cn

需要对获取数据进行预处理

import pandas as pd
import os
import re
from datetime import datetime
import json
import numpy as np


def process_medical_data():
    """主处理函数:读取、清洗、分析并保存医疗数据"""

    # 1. 设置路径
    base_path = r"文件目录"
    input_file = os.path.join(base_path, "疾病数据.csv")

    # 2. 创建输出文件夹
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_folder = os.path.join(base_path, f"医疗数据处理结果_{timestamp}")

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        print(f"✅ 已创建输出文件夹: {output_folder}")
    else:
        print(f"📁 使用现有输出文件夹: {output_folder}")

    print("=" * 60)
    print("开始处理医疗数据...")
    print("=" * 60)

    # 3. 尝试读取CSV文件
    print(f"\n📥 正在读取文件: {input_file}")

    if not os.path.exists(input_file):
        print(f"❌ 文件不存在: {input_file}")
        print("请检查文件路径是否正确")
        return None

    # 尝试不同的编码方式读取CSV
    encodings = ['utf-8', 'gbk', 'gb2312', 'latin1']
    df = None

    for encoding in encodings:
        try:
            df = pd.read_csv(input_file, encoding=encoding)
            print(f"✅ 使用 {encoding} 编码成功读取文件")
            break
        except UnicodeDecodeError:
            print(f"编码 {encoding} 失败,尝试下一个...")
            continue
        except Exception as e:
            print(f"读取文件时发生错误 ({encoding}): {e}")
            continue

    if df is None:
        print("❌ 无法读取CSV文件,请检查文件格式")
        return None

    # 4. 显示数据基本信息
    print(f"\n📊 数据基本信息:")
    print(f"   - 数据形状: {df.shape}")
    print(f"   - 行数: {len(df)}")
    print(f"   - 列数: {len(df.columns)}")
    print(f"   - 列名: {list(df.columns)}")

    # 5. 数据预览
    print(f"\n🔍 数据前3行预览:")
    for i in range(min(3, len(df))):
        row_data = []
        for val in df.iloc[i].tolist():
            if isinstance(val, str) and len(val) > 50:
                row_data.append(val[:50] + "...")
            else:
                row_data.append(str(val)[:50])
        print(f"行{i + 1}: {row_data}")

    # 6. 数据清洗
    print(f"\n🧹 开始数据清洗...")

    # 清理列名
    df.columns = [str(col).strip() for col in df.columns]
    print(f"   清理后列名: {list(df.columns)}")

    # 7. 创建清洗函数
    def clean_text(text):
        """清理文本数据"""
        if pd.isna(text):
            return ""

        text_str = str(text)

        # 去除多余空格
        text_str = re.sub(r'\s+', ' ', text_str).strip()

        return text_str

    def extract_symptoms_from_row(row):
        """从一行中提取所有症状"""
        symptoms = []

        # 症状相关列
        symptom_columns = [
            'result_item_top_l',
            'result_item_content_label1',
            'result_item_content_label2',
            'result_item_content_label3',
            'result_item_content_label4'
        ]

        for col in symptom_columns:
            if col in df.columns and pd.notna(row[col]):
                symptom = clean_text(row[col])
                if symptom and len(symptom) > 1:
                    symptoms.append(symptom)

        return list(set(symptoms))  # 去重

    # 8. 提取疾病和症状信息
    print(f"\n🎯 提取疾病和症状信息...")

    cleaned_data = []

    for idx, row in df.iterrows():
        # 提取疾病名称
        disease_name = clean_text(row['标题']) if '标题' in df.columns else f"疾病_{idx + 1}"

        # 提取症状
        symptoms = extract_symptoms_from_row(row)

        # 提取疾病描述
        disease_desc = clean_text(row['result_item_content']) if 'result_item_content' in df.columns else ""

        # 提取链接
        disease_link = clean_text(row['标题链接']) if '标题链接' in df.columns else ""

        # 添加到结果
        cleaned_data.append({
            '疾病ID': idx + 1,
            '疾病名称': disease_name,
            '疾病描述': disease_desc,
            '疾病链接': disease_link,
            '主要症状': clean_text(row['result_item_top_l']) if 'result_item_top_l' in df.columns else "",
            '症状列表': ';'.join(symptoms) if symptoms else '',
            '症状数量': len(symptoms)
        })

    # 9. 创建清洗后的DataFrame
    cleaned_df = pd.DataFrame(cleaned_data)

    print(f"\n✅ 数据清洗完成:")
    print(f"   原始记录数: {len(df)}")
    print(f"   清洗后记录数: {len(cleaned_df)}")

    if len(cleaned_df) == 0:
        print("❌ 没有提取到有效数据")
        return None

    # 10. 显示数据示例
    print(f"\n🔍 清洗后数据示例 (前3行):")
    for i in range(min(3, len(cleaned_df))):
        print(f"行{i + 1}:")
        print(f"  疾病: {cleaned_df.iloc[i]['疾病名称']}")
        print(f"  症状数: {cleaned_df.iloc[i]['症状数量']}")
        print(f"  症状: {cleaned_df.iloc[i]['症状列表'][:50]}...")
        print()

    # 11. 保存各种格式的文件
    print(f"\n💾 正在保存处理结果...")

    # 保存为Excel
    excel_path = os.path.join(output_folder, "疾病数据_清洗后.xlsx")
    cleaned_df.to_excel(excel_path, index=False)
    print(f"   ✅ Excel文件: {excel_path}")

    # 保存为CSV
    csv_path = os.path.join(output_folder, "疾病数据_清洗后.csv")
    cleaned_df.to_csv(csv_path, index=False, encoding='utf-8-sig')
    print(f"   ✅ CSV文件: {csv_path}")

    # 12. 生成JSON格式(适合大模型处理)
    json_path = os.path.join(output_folder, "疾病数据_大模型格式.json")
    llm_data = []

    for _, row in cleaned_df.iterrows():
        symptoms = row['症状列表'].split(';') if row['症状列表'] else []

        llm_data.append({
            "disease_id": int(row['疾病ID']),
            "disease_name": row['疾病名称'],
            "disease_description": row['疾病描述'][:100] + "..." if row['疾病描述'] and len(row['疾病描述']) > 100 else
            row['疾病描述'],
            "main_symptom": row['主要症状'],
            "symptoms": symptoms,
            "symptom_count": len(symptoms),
            "disease_link": row['疾病链接'],
            "text_for_llm": f"疾病名称:{row['疾病名称']}。主要症状:{row['主要症状']}。相关症状:{'、'.join(symptoms)}。疾病描述:{row['疾病描述'][:50] if row['疾病描述'] else '无详细描述'}。"
        })

    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(llm_data, f, ensure_ascii=False, indent=2)

    print(f"   ✅ JSON文件: {json_path}")

    # 13. 生成50个样本(用于大模型测试)
    sample_size = min(50, len(llm_data))
    if sample_size == len(llm_data):
        samples = llm_data
    else:
        # 随机抽样
        indices = np.random.choice(len(llm_data), sample_size, replace=False)
        samples = [llm_data[i] for i in indices]

    # 为每个样本创建不同temperature的prompt
    for i, sample in enumerate(samples):
        disease_name = sample['disease_name']
        main_symptom = sample['main_symptom']
        symptoms = sample['symptoms'][:3]  # 取前3个症状
        description = sample['disease_description'][:50] if sample['disease_description'] else "暂无详细描述"

        sample[
            'temperature_0_prompt'] = f"请精确提取以下医疗信息:疾病名称:{disease_name},主要症状:{main_symptom},相关症状:{'、'.join(symptoms)}。"
        sample[
            'temperature_0.5_prompt'] = f"请分析以下疾病描述:{disease_name}是一种疾病,主要表现为{main_symptom},常伴有{'、'.join(symptoms)}等症状。{description}"
        sample[
            'temperature_1_prompt'] = f"这是一条医疗记录:患者被诊断为{disease_name},主诉{main_symptom},检查发现与{'、'.join(symptoms)}相关。请提取关键医疗信息。"

    sample_path = os.path.join(output_folder, "疾病数据_50样本.json")
    with open(sample_path, 'w', encoding='utf-8') as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)

    print(f"   ✅ {sample_size}个样本文件: {sample_path}")

    # 14. 生成统计报告
    report_path = os.path.join(output_folder, "数据统计报告.txt")

    with open(report_path, 'w', encoding='utf-8') as f:
        f.write("=" * 60 + "\n")
        f.write("        医疗数据处理统计报告\n")
        f.write("=" * 60 + "\n\n")

        f.write(f"处理时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"原始文件: {os.path.basename(input_file)}\n")
        f.write(f"输出文件夹: {os.path.basename(output_folder)}\n\n")

        f.write("1. 基础统计:\n")
        f.write("-" * 40 + "\n")
        f.write(f"   原始记录数: {len(df)}\n")
        f.write(f"   清洗后记录数: {len(cleaned_df)}\n")
        f.write(f"   数据完整率: {len(cleaned_df) / len(df) * 100:.1f}%\n\n")

        f.write("2. 症状数量统计:\n")
        f.write("-" * 40 + "\n")
        if '症状数量' in cleaned_df.columns:
            stats = cleaned_df['症状数量'].describe()
            for stat, value in stats.items():
                f.write(f"   {stat}: {value:.2f}\n")
        f.write("\n")

        f.write("3. 症状数量分布:\n")
        f.write("-" * 40 + "\n")
        symptom_dist = cleaned_df['症状数量'].value_counts().sort_index()
        for count, freq in symptom_dist.items():
            percentage = (freq / len(cleaned_df)) * 100
            f.write(f"   {count}个症状: {freq}条 ({percentage:.1f}%)\n")
        f.write("\n")

        f.write("4. 常见疾病类型 (前10个):\n")
        f.write("-" * 40 + "\n")
        disease_counts = cleaned_df['疾病名称'].value_counts().head(10)
        for disease, count in disease_counts.items():
            f.write(f"   {disease}: {count}次\n")
        f.write("\n")

        f.write("5. 常见症状 (前10个):\n")
        f.write("-" * 40 + "\n")
        # 统计所有症状
        all_symptoms = []
        for symptoms_str in cleaned_df['症状列表']:
            if symptoms_str:
                all_symptoms.extend(symptoms_str.split(';'))

        from collections import Counter
        symptom_counter = Counter(all_symptoms)
        for symptom, count in symptom_counter.most_common(10):
            f.write(f"   {symptom}: {count}次\n")
        f.write("\n")

        f.write("6. 生成的文件:\n")
        f.write("-" * 40 + "\n")
        f.write("   疾病数据_清洗后.xlsx - Excel格式\n")
        f.write("   疾病数据_清洗后.csv - CSV格式\n")
        f.write("   疾病数据_大模型格式.json - JSON格式\n")
        f.write(f"   疾病数据_{sample_size}样本.json - 测试样本\n")
        f.write("   数据统计报告.txt - 本报告\n")

    print(f"   ✅ 统计报告: {report_path}")

    # 15. 显示处理结果摘要
    print(f"\n" + "=" * 60)
    print("处理完成摘要")
    print("=" * 60)
    print(f"✅ 成功处理 {len(cleaned_df)} 条疾病记录")
    print(f"📁 所有文件已保存到: {output_folder}")

    # 显示输出文件夹内容
    print(f"\n📁 输出文件夹内容:")
    print("-" * 40)
    for file in os.listdir(output_folder):
        file_path = os.path.join(output_folder, file)
        if os.path.isfile(file_path):
            size_kb = os.path.getsize(file_path) / 1024
            print(f"  {file} ({size_kb:.1f} KB)")

    # 16. 显示样本示例
    print(f"\n🔍 样本示例 (前3个):")
    print("-" * 40)
    for i, sample in enumerate(samples[:3], 1):
        print(f"样本 {i}:")
        print(f"  疾病: {sample['disease_name']}")
        print(f"  主要症状: {sample['main_symptom']}")
        print(f"  症状数: {sample['symptom_count']}")
        print(f"  症状示例: {', '.join(sample['symptoms'][:3])}")
        print(f"  Temperature 0 Prompt: {sample['temperature_0_prompt'][:60]}...")
        print()

    return cleaned_df


def main():
    """主函数"""
    print("=" * 60)
    print("医疗数据处理程序")
    print("=" * 60)

    try:
        result_df = process_medical_data()

        if result_df is not None:
            print("\n✨ 处理完成!您可以:")
            print("1. 使用 '疾病数据_50样本.json' 进行大模型测试")
            print("2. 测试3种不同的temperature值 (0, 0.5, 1)")
            print("3. 使用 '疾病数据_清洗后.csv' 导入到数据库")
            print("4. 查看 '数据统计报告.txt' 了解详细统计信息")
        else:
            print("\n❌ 处理失败")

    except Exception as e:
        print(f"\n❌ 处理过程中发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

二、大模型工具提取疾病结构化数据及效果对比

(1)deepseek版(付费,可跳)

使用过deepseek的可以登录deepseek平台:https://platform.deepseek.com

创建API

"""
疾病结构化数据提取实验
要求测试3种不同temperature值、3种不同prompt的提取效果
使用DeepSeek模型
"""

import pandas as pd
import os
import json
import time
from tqdm import tqdm
import requests
from typing import List, Dict, Any
import re


class DeepSeekExtractor:
    """DeepSeek模型结构化数据提取器"""

    def __init__(self, api_key: str, api_base: str = "https://api.deepseek.com"):
        """
        初始化DeepSeek提取器

        Args:
            api_key: DeepSeek API密钥
            api_base: API基础URL
        """
        self.api_key = api_key
        self.api_base = api_base
        self.model = "deepseek-chat"  # DeepSeek对话模型

    def call_deepseek_api(self, prompt: str, temperature: float = 0.1) -> str:
        """
        调用DeepSeek API

        Args:
            prompt: 提示词
            temperature: 温度参数

        Returns:
            API返回的文本内容
        """
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "model": self.model,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": temperature,
            "max_tokens": 1000
        }

        try:
            response = requests.post(
                f"{self.api_base}/chat/completions",
                headers=headers,
                json=payload,
                timeout=60
            )

            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content']
            else:
                print(f"API调用失败: {response.status_code}, {response.text}")
                return ""

        except Exception as e:
            print(f"API调用异常: {e}")
            return ""

    def extract_json_from_response(self, response_text: str) -> Dict:
        """
        从响应文本中提取JSON数据

        Args:
            response_text: API返回文本

        Returns:
            解析后的字典
        """
        try:
            # 尝试直接解析JSON
            response_text = response_text.strip()

            # 提取JSON部分(可能包含markdown代码块)
            if '```json' in response_text:
                json_str = response_text.split('```json')[1].split('```')[0].strip()
            elif '```' in response_text:
                json_str = response_text.split('```')[1].split('```')[0].strip()
            else:
                json_str = response_text

            # 尝试解析JSON
            data = json.loads(json_str)
            return data

        except json.JSONDecodeError:
            # 如果JSON解析失败,尝试提取键值对
            print("JSON解析失败,尝试提取键值对...")
            return self.extract_key_value_pairs(response_text)
        except Exception as e:
            print(f"解析响应失败: {e}")
            return {}

    def extract_key_value_pairs(self, text: str) -> Dict:
        """
        从文本中提取键值对

        Args:
            text: 文本内容

        Returns:
            提取的键值对字典
        """
        result = {}

        # 定义要提取的字段
        fields = ["疾病名称", "发病部位", "就诊科室", "主要症状", "常用药物", "是否传染"]

        for field in fields:
            # 尝试不同模式匹配
            patterns = [
                rf'{field}[::]\s*([^\n]+)',
                rf'{field}[::]\s*"([^"]+)"',
                rf'{field}[::]\s*([^,。;\n]+)'
            ]

            for pattern in patterns:
                match = re.search(pattern, text)
                if match:
                    result[field] = match.group(1).strip()
                    break
            else:
                result[field] = ""

        return result

    def normalize_fields(self, data: Dict) -> Dict:
        """
        标准化字段名

        Args:
            data: 原始数据

        Returns:
            标准化后的数据
        """
        # 字段映射关系
        field_mapping = {
            "疾病名称": "disease",
            "disease": "disease",
            "疾病": "disease",
            "发病部位": "part",
            "part": "part",
            "部位": "part",
            "就诊科室": "department",
            "department": "department",
            "科室": "department",
            "主要症状": "symptom",
            "symptom": "symptom",
            "症状": "symptom",
            "常用药物": "drug",
            "drug": "drug",
            "药物": "drug",
            "是否传染": "infectious",
            "infectious": "infectious",
            "传染性": "infectious"
        }

        normalized = {}
        for key, value in data.items():
            if key in field_mapping:
                normalized_key = field_mapping[key]
                normalized[normalized_key] = value
            elif key in ["disease", "part", "department", "symptom", "drug", "infectious"]:
                normalized[key] = value

        # 确保所有字段都存在
        required_fields = ["disease", "part", "department", "symptom", "drug", "infectious"]
        for field in required_fields:
            if field not in normalized:
                normalized[field] = ""

        return normalized


class DiseaseDataProcessor:
    """疾病数据处理器"""

    def __init__(self, data_path: str):
        """
        初始化处理器

        Args:
            data_path: 数据文件路径
        """
        self.data_path = data_path
        self.df = None
        self.samples = None

    def load_data(self) -> bool:
        """
        加载数据

        Returns:
            是否成功加载
        """
        try:
            # 尝试不同格式
            if self.data_path.endswith('.csv'):
                self.df = pd.read_csv(self.data_path, encoding='utf-8')
            elif self.data_path.endswith('.xlsx'):
                self.df = pd.read_excel(self.data_path)
            elif self.data_path.endswith('.json'):
                with open(self.data_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.df = pd.DataFrame(data)
            else:
                print("不支持的文件格式")
                return False

            print(f"成功加载数据,共 {len(self.df)} 条记录")

            # 查找包含疾病描述的列
            text_columns = [col for col in self.df.columns if
                            'text' in col.lower() or '描述' in col or 'content' in col.lower()]
            if text_columns:
                print(f"找到疾病描述列: {text_columns[0]}")
                self.df['disease_text'] = self.df[text_columns[0]]
            else:
                # 尝试合并多列
                self.df['disease_text'] = self.df.apply(
                    lambda row: ' '.join([str(row[col]) for col in self.df.columns]), axis=1)

            return True

        except Exception as e:
            print(f"加载数据失败: {e}")
            return False

    def sample_data(self, n: int = 50, random_seed: int = 42) -> List[Dict]:
        """
        随机抽取样本

        Args:
            n: 样本数量
            random_seed: 随机种子

        Returns:
            样本列表
        """
        if self.df is None or len(self.df) == 0:
            print("没有可用数据")
            return []

        # 确保有足够的样本
        n = min(n, len(self.df))

        # 随机抽取
        sampled_df = self.df.sample(n=n, random_state=random_seed)

        samples = []
        for idx, row in sampled_df.iterrows():
            sample = {
                "id": idx + 1,
                "disease_text": str(row.get('disease_text', '')).strip(),
                "original_disease": str(row.get('疾病名称', row.get('disease_name', ''))).strip()
            }
            if sample['disease_text']:  
                samples.append(sample)

        self.samples = samples
        print(f"抽取了 {len(samples)} 个样本")
        return samples


class ExperimentRunner:
    """实验运行器"""

    def __init__(self, extractor: DeepSeekExtractor):
        """
        初始化实验运行器

        Args:
            extractor: DeepSeek提取器实例
        """
        self.extractor = extractor

        # 定义实验参数
        self.temperatures = [0.1, 0.5, 1.0]
        self.prompt_types = ["简单版", "专家版", "详细版"]

        # 实验结果存储
        self.results = []
        self.summary = []

    def build_prompt(self, prompt_type: str, text: str) -> str:
        """
        构建不同版本的prompt

        Args:
            prompt_type: prompt类型
            text: 疾病文本

        Returns:
            构建好的prompt
        """
        # JSON模板
        json_template = """{
  "disease": "疾病名称",
  "part": "发病部位",
  "department": "就诊科室", 
  "symptom": "主要症状",
  "drug": "常用药物",
  "infectious": "是否传染"
}"""

        if prompt_type == "简单版":
            return f"""请从以下医学文本中提取结构化信息,严格按照JSON格式输出:
{json_template}

文本内容:
{text}

要求:
1. 只输出JSON,不要有任何额外说明
2. 如果字段信息不存在,使用空字符串""
3. 保持字段顺序一致"""

        elif prompt_type == "专家版":
            return f"""你是一名专业的医学信息抽取专家,请从以下临床文本中提取关键医学信息。

任务要求:
1. 严格按照给定的JSON结构输出
2. 只提取文本中明确提到的信息,不要推理或补充
3. 字段说明:
   - disease: 疾病名称
   - part: 发病部位(如:肺部、心脏、消化道等)
   - department: 就诊科室(如:呼吸内科、心血管科、消化内科等)
   - symptom: 主要症状(多个症状用逗号分隔)
   - drug: 常用药物(多个药物用逗号分隔)
   - infectious: 是否传染(是/否/未提及)

输出格式:
{json_template}

待处理文本:
{text}"""

        elif prompt_type == "详细版":
            example_text = "患者男性,45岁,因持续咳嗽、咳痰、发热3天就诊。查体:体温38.5℃,双肺可闻及湿啰音。诊断为社区获得性肺炎,建议呼吸内科就诊。治疗药物包括阿莫西林、左氧氟沙星。该疾病具有传染性。"

            example_output = """{
  "disease": "社区获得性肺炎",
  "part": "肺部",
  "department": "呼吸内科",
  "symptom": "咳嗽,咳痰,发热",
  "drug": "阿莫西林,左氧氟沙星", 
  "infectious": "是"
}"""

            return f"""请参考以下示例,从提供的医学文本中提取结构化信息。

【示例】
输入文本:{example_text}
输出JSON:{example_output}

【抽取规则】
1. disease: 提取明确的疾病诊断名称
2. part: 提取发病器官或部位
3. department: 提取建议就诊的科室
4. symptom: 提取患者的主要症状,用逗号分隔
5. drug: 提取提到的治疗药物,用逗号分隔
6. infectious: 根据文本判断传染性,填写"是"、"否"或"未提及"

【待处理文本】
{text}

请严格按照示例格式输出JSON:"""

        else:
            raise ValueError(f"不支持的prompt类型: {prompt_type}")

    def run_experiment(self, samples: List[Dict]) -> None:
        """
        运行3×3实验

        Args:
            samples: 样本数据
        """
        print(f"\n{'=' * 60}")
        print("开始3×3实验")
        print(f"样本数: {len(samples)}")
        print(f"Temperature值: {self.temperatures}")
        print(f"Prompt类型: {self.prompt_types}")
        print(f"{'=' * 60}")

        for temp in self.temperatures:
            for prompt_type in self.prompt_types:
                print(f"\n▶ 测试组合: Temperature={temp}, Prompt={prompt_type}")

                temp_results = []
                success_count = 0
                total_score = 0

                for sample in tqdm(samples, desc=f"处理进度"):
                    text = sample["disease_text"]
                    sample_id = sample["id"]

                    if not text:
                        continue

                    # 构建prompt
                    prompt = self.build_prompt(prompt_type, text)

                    # 调用API
                    response = self.extractor.call_deepseek_api(prompt, temperature=temp)

                    # 提取和解析结果
                    if response:
                        extracted_data = self.extractor.extract_json_from_response(response)
                        normalized_data = self.extractor.normalize_fields(extracted_data)

                        # 计算得分(非空字段数量)
                        score = sum(1 for value in normalized_data.values() if value and str(value).strip())

                        # 保存结果
                        result = {
                            "sample_id": sample_id,
                            "original_disease": sample.get("original_disease", ""),
                            "temperature": temp,
                            "prompt_type": prompt_type,
                            "raw_response": response[:200] + "..." if len(response) > 200 else response,
                            "extracted_data": normalized_data,
                            "score": score,
                            "max_score": 6  # 6个字段
                        }

                        temp_results.append(result)
                        success_count += 1
                        total_score += score

                        # API限速控制
                        time.sleep(0.5)
                    else:
                        print(f"样本{sample_id}: API调用失败")

                # 计算统计信息
                avg_score = total_score / success_count if success_count > 0 else 0
                success_rate = (success_count / len(samples)) * 100 if samples else 0

                # 保存统计信息
                self.summary.append({
                    "temperature": temp,
                    "prompt_type": prompt_type,
                    "total_samples": len(samples),
                    "success_count": success_count,
                    "success_rate": round(success_rate, 2),
                    "avg_score": round(avg_score, 2),
                    "max_possible_score": 6
                })

                # 保存详细结果
                self.results.extend(temp_results)

                print(f"✓ 完成: 成功率={success_rate:.1f}%, 平均得分={avg_score:.2f}/6")

    def save_results(self, output_dir: str = "experiment_results"):
        """
        保存实验结果

        Args:
            output_dir: 输出目录
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        timestamp = time.strftime("%Y%m%d_%H%M%S")

        # 保存汇总结果
        summary_df = pd.DataFrame(self.summary)
        summary_path = os.path.join(output_dir, f"experiment_summary_{timestamp}.xlsx")
        summary_df.to_excel(summary_path, index=False)

        # 保存详细结果
        detailed_data = []
        for result in self.results:
            detailed_result = {
                "sample_id": result["sample_id"],
                "original_disease": result["original_disease"],
                "temperature": result["temperature"],
                "prompt_type": result["prompt_type"],
                "score": result["score"],
                "max_score": result["max_score"],
                "raw_response": result["raw_response"]
            }

            # 添加提取的字段
            extracted = result["extracted_data"]
            for key in ["disease", "part", "department", "symptom", "drug", "infectious"]:
                detailed_result[key] = extracted.get(key, "")

            detailed_data.append(detailed_result)

        detailed_df = pd.DataFrame(detailed_data)
        detailed_path = os.path.join(output_dir, f"experiment_detailed_{timestamp}.xlsx")
        detailed_df.to_excel(detailed_path, index=False)

        # 保存为JSON格式
        json_path = os.path.join(output_dir, f"experiment_results_{timestamp}.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "summary": self.summary,
                "results": [
                    {
                        k: v for k, v in result.items()
                        if k != "extracted_data"
                    } | {"extracted_data": result["extracted_data"]}
                    for result in self.results
                ]
            }, f, ensure_ascii=False, indent=2)

        print(f"\n{'=' * 60}")
        print("实验结果已保存:")
        print(f"1. 汇总结果: {summary_path}")
        print(f"2. 详细结果: {detailed_path}")
        print(f"3. JSON格式: {json_path}")

        # 打印汇总统计
        print(f"\n实验汇总统计:")
        print(summary_df.to_string(index=False))

        return summary_df, detailed_df

    def analyze_results(self):
        """分析实验结果"""
        if not self.summary:
            print("没有实验结果可分析")
            return

        print(f"\n{'=' * 60}")
        print("实验结果分析")
        print(f"{'=' * 60}")

        summary_df = pd.DataFrame(self.summary)

        # 按temperature分析
        print("\n1. 不同Temperature的效果对比:")
        temp_groups = summary_df.groupby('temperature')
        for temp, group in temp_groups:
            print(f"\nTemperature = {temp}:")
            for _, row in group.iterrows():
                print(f"  Prompt={row['prompt_type']}: 成功率={row['success_rate']}%, 平均得分={row['avg_score']:.2f}")

        # 按prompt类型分析
        print("\n2. 不同Prompt类型的效果对比:")
        prompt_groups = summary_df.groupby('prompt_type')
        for prompt, group in prompt_groups:
            print(f"\nPrompt类型 = {prompt}:")
            for _, row in group.iterrows():
                print(
                    f"  Temperature={row['temperature']}: 成功率={row['success_rate']}%, 平均得分={row['avg_score']:.2f}")

        # 最佳组合
        best_result = summary_df.loc[summary_df['avg_score'].idxmax()]
        print(f"\n3. 最佳组合:")
        print(f"   Temperature: {best_result['temperature']}")
        print(f"   Prompt类型: {best_result['prompt_type']}")
        print(f"   平均得分: {best_result['avg_score']:.2f}/6")
        print(f"   成功率: {best_result['success_rate']}%")


def main():
    """主函数"""

    # 1. 配置DeepSeek API
    print("=" * 60)
    print("DeepSeek疾病结构化数据提取实验")
    print("=" * 60)

    api_key = os.environ.get("DEEPSEEK_API_KEY")
    if not api_key:
        api_key = input("请输入您的DeepSeek API Key: ").strip()
        os.environ["DEEPSEEK_API_KEY"] = api_key

    if not api_key:
        print("错误: 需要提供DeepSeek API Key")
        return

    # 2. 初始化提取器
    extractor = DeepSeekExtractor(api_key=api_key)

    # 3. 加载数据
    data_path = r"预处理生成的json文件位置"

    if not os.path.exists(data_path):
        # 尝试其他可能的路径
        alternatives = [
            r"文件目录\疾病数据_清洗后.csv",
            r"文件目录\疾病数据_清洗后.xlsx"
        ]

        for alt_path in alternatives:
            if os.path.exists(alt_path):
                data_path = alt_path
                break
        else:
            print("错误: 未找到数据文件")
            print("请确保数据文件存在于指定路径")
            return

    print(f"使用数据文件: {data_path}")

    processor = DiseaseDataProcessor(data_path)
    if not processor.load_data():
        return

    # 4. 抽取样本
    samples = processor.sample_data(n=50, random_seed=42)
    if not samples:
        print("错误: 未能抽取到样本数据")
        return

    # 5. 运行实验
    experiment = ExperimentRunner(extractor)
    experiment.run_experiment(samples)

    # 6. 保存结果
    output_dir = r"输出位置"
    experiment.save_results(output_dir)

    # 7. 分析结果
    experiment.analyze_results()

    print(f"\n{'=' * 60}")
    print("实验完成!")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()

在上图位置输出API

(2)阿里云百炼版(有免费额度)

阿里云新用户有免费的100万tokens额度,一般实验完全够用

注册账号,在密钥管理处创建API

"""
疾病结构化数据提取实验
要求测试3种不同temperature值、3种不同prompt的提取效果
使用阿里云百炼(DashScope)通义千问模型
"""

import pandas as pd
import os
import json
import time
from tqdm import tqdm
import dashscope
from typing import List, Dict, Any
import re


class DashScopeExtractor:
    """阿里云百炼模型结构化数据提取器"""

    def __init__(self, api_key: str, model: str = "qwen-max"):
        """
        初始化百炼提取器

        Args:
            api_key: DashScope API密钥
            model: 模型名称,可选 qwen-max, qwen-plus, qwen-turbo
        """
        self.api_key = api_key
        self.model = model
        # 设置API密钥
        dashscope.api_key = api_key

    def call_dashscope_api(self, prompt: str, temperature: float = 0.1) -> str:
        """
        调用DashScope API

        Args:
            prompt: 提示词
            temperature: 温度参数(在DashScope中对应top_p参数)

        Returns:
            API返回的文本内容
        """
        try:
            # 创建消息
            messages = [{'role': 'user', 'content': prompt}]

            # 构建请求参数
            from dashscope import Generation

            response = Generation.call(
                model=self.model,
                messages=messages,
                result_format='message',  # 返回消息格式
                temperature=temperature,  # 温度参数
                top_p=0.8,  # 核采样参数
                seed=42,  # 随机种子
                max_tokens=1500  # 最大token数
            )

            if response.status_code == 200:
                return response.output.choices[0]['message']['content']
            else:
                print(f"API调用失败: {response.status_code}, {response.message}")
                return ""

        except Exception as e:
            print(f"API调用异常: {e}")
            return ""

    def extract_json_from_response(self, response_text: str) -> Dict:
        """
        从响应文本中提取JSON数据

        Args:
            response_text: API返回文本

        Returns:
            解析后的字典
        """
        try:
            # 尝试直接解析JSON
            response_text = response_text.strip()

            # 提取JSON部分(可能包含markdown代码块)
            if '```json' in response_text:
                json_str = response_text.split('```json')[1].split('```')[0].strip()
            elif '```' in response_text:
                json_str = response_text.split('```')[1].split('```')[0].strip()
            else:
                json_str = response_text

            # 尝试解析JSON
            data = json.loads(json_str)
            return data

        except json.JSONDecodeError:
            # 如果JSON解析失败,尝试提取键值对
            print("JSON解析失败,尝试提取键值对...")
            return self.extract_key_value_pairs(response_text)
        except Exception as e:
            print(f"解析响应失败: {e}")
            return {}

    def extract_key_value_pairs(self, text: str) -> Dict:
        """
        从文本中提取键值对

        Args:
            text: 文本内容

        Returns:
            提取的键值对字典
        """
        result = {}

        # 定义要提取的字段
        fields = ["疾病名称", "发病部位", "就诊科室", "主要症状", "常用药物", "是否传染"]

        for field in fields:
            # 尝试不同模式匹配
            patterns = [
                rf'{field}[::]\s*([^\n]+)',
                rf'{field}[::]\s*"([^"]+)"',
                rf'{field}[::]\s*([^,。;\n]+)'
            ]

            for pattern in patterns:
                match = re.search(pattern, text)
                if match:
                    result[field] = match.group(1).strip()
                    break
            else:
                result[field] = ""

        return result

    def normalize_fields(self, data: Dict) -> Dict:
        """
        标准化字段名

        Args:
            data: 原始数据

        Returns:
            标准化后的数据
        """
        # 字段映射关系
        field_mapping = {
            "疾病名称": "disease",
            "disease": "disease",
            "疾病": "disease",
            "发病部位": "part",
            "part": "part",
            "部位": "part",
            "就诊科室": "department",
            "department": "department",
            "科室": "department",
            "主要症状": "symptom",
            "symptom": "symptom",
            "症状": "symptom",
            "常用药物": "drug",
            "drug": "drug",
            "药物": "drug",
            "是否传染": "infectious",
            "infectious": "infectious",
            "传染性": "infectious"
        }

        normalized = {}
        for key, value in data.items():
            if key in field_mapping:
                normalized_key = field_mapping[key]
                normalized[normalized_key] = value
            elif key in ["disease", "part", "department", "symptom", "drug", "infectious"]:
                normalized[key] = value

        # 确保所有字段都存在
        required_fields = ["disease", "part", "department", "symptom", "drug", "infectious"]
        for field in required_fields:
            if field not in normalized:
                normalized[field] = ""

        return normalized


class DiseaseDataProcessor:
    """疾病数据处理器"""

    def __init__(self, data_path: str):
        """
        初始化处理器

        Args:
            data_path: 数据文件路径
        """
        self.data_path = data_path
        self.df = None
        self.samples = None

    def load_data(self) -> bool:
        """
        加载数据

        Returns:
            是否成功加载
        """
        try:
            # 尝试不同格式
            if self.data_path.endswith('.csv'):
                self.df = pd.read_csv(self.data_path, encoding='utf-8')
            elif self.data_path.endswith('.xlsx'):
                self.df = pd.read_excel(self.data_path)
            elif self.data_path.endswith('.json'):
                with open(self.data_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.df = pd.DataFrame(data)
            else:
                print("不支持的文件格式")
                return False

            print(f"成功加载数据,共 {len(self.df)} 条记录")

            # 查找包含疾病描述的列
            text_columns = [col for col in self.df.columns if
                            'text' in col.lower() or '描述' in col or 'content' in col.lower()]
            if text_columns:
                print(f"找到疾病描述列: {text_columns[0]}")
                self.df['disease_text'] = self.df[text_columns[0]]
            else:
                # 尝试合并多列
                self.df['disease_text'] = self.df.apply(
                    lambda row: ' '.join([str(row[col]) for col in self.df.columns]), axis=1)

            return True

        except Exception as e:
            print(f"加载数据失败: {e}")
            return False

    def sample_data(self, n: int = 50, random_seed: int = 42) -> List[Dict]:
        """
        随机抽取样本

        Args:
            n: 样本数量
            random_seed: 随机种子

        Returns:
            样本列表
        """
        if self.df is None or len(self.df) == 0:
            print("没有可用数据")
            return []

        # 确保有足够的样本
        n = min(n, len(self.df))

        # 随机抽取
        sampled_df = self.df.sample(n=n, random_state=random_seed)

        samples = []
        for idx, row in sampled_df.iterrows():
            sample = {
                "id": idx + 1,
                "disease_text": str(row.get('disease_text', '')).strip(),
                "original_disease": str(row.get('疾病名称', row.get('disease_name', ''))).strip()
            }
            if sample['disease_text']:  # 只添加非空文本
                samples.append(sample)

        self.samples = samples
        print(f"抽取了 {len(samples)} 个样本")
        return samples


class ExperimentRunner:
    """实验运行器"""

    def __init__(self, extractor: DashScopeExtractor):
        """
        初始化实验运行器

        Args:
            extractor: 百炼提取器实例
        """
        self.extractor = extractor

        # 定义实验参数
        self.temperatures = [0.1, 0.5, 1.0]
        self.prompt_types = ["简单版", "专家版", "详细版"]

        # 实验结果存储
        self.results = []
        self.summary = []

    def build_prompt(self, prompt_type: str, text: str) -> str:
        """
        构建不同版本的prompt

        Args:
            prompt_type: prompt类型
            text: 疾病文本

        Returns:
            构建好的prompt
        """
        # JSON模板
        json_template = """{
  "disease": "疾病名称",
  "part": "发病部位",
  "department": "就诊科室", 
  "symptom": "主要症状",
  "drug": "常用药物",
  "infectious": "是否传染"
}"""

        if prompt_type == "简单版":
            return f"""请从以下医学文本中提取结构化信息,严格按照JSON格式输出。
要求:
1. 只输出JSON,不要有任何额外说明
2. 如果字段信息不存在,使用空字符串""
3. 保持字段顺序一致

JSON格式要求:
{json_template}

待处理文本:
{text}"""

        elif prompt_type == "专家版":
            return f"""你是一名专业的医学信息抽取专家,请从以下临床文本中提取关键医学信息。

任务要求:
1. 严格按照给定的JSON结构输出
2. 只提取文本中明确提到的信息,不要推理或补充
3. 字段说明:
   - disease: 疾病名称
   - part: 发病部位(如:肺部、心脏、消化道等)
   - department: 就诊科室(如:呼吸内科、心血管科、消化内科等)
   - symptom: 主要症状(多个症状用逗号分隔)
   - drug: 常用药物(多个药物用逗号分隔)
   - infectious: 是否传染(是/否/未提及)

输出格式:
{json_template}

待处理文本:
{text}"""

        elif prompt_type == "详细版":
            example_text = "患者男性,45岁,因持续咳嗽、咳痰、发热3天就诊。查体:体温38.5℃,双肺可闻及湿啰音。诊断为社区获得性肺炎,建议呼吸内科就诊。治疗药物包括阿莫西林、左氧氟沙星。该疾病具有传染性。"

            example_output = """{
  "disease": "社区获得性肺炎",
  "part": "肺部",
  "department": "呼吸内科",
  "symptom": "咳嗽,咳痰,发热",
  "drug": "阿莫西林,左氧氟沙星", 
  "infectious": "是"
}"""

            return f"""请参考以下示例,从提供的医学文本中提取结构化信息。

【示例】
输入文本:{example_text}
输出JSON:{example_output}

【抽取规则】
1. disease: 提取明确的疾病诊断名称
2. part: 提取发病器官或部位
3. department: 提取建议就诊的科室
4. symptom: 提取患者的主要症状,用逗号分隔
5. drug: 提取提到的治疗药物,用逗号分隔
6. infectious: 根据文本判断传染性,填写"是"、"否"或"未提及"

【待处理文本】
{text}

请严格按照示例格式输出JSON:"""

        else:
            raise ValueError(f"不支持的prompt类型: {prompt_type}")

    def run_experiment(self, samples: List[Dict]) -> None:
        """
        运行3×3实验

        Args:
            samples: 样本数据
        """
        print(f"\n{'=' * 60}")
        print("开始3×3实验")
        print(f"样本数: {len(samples)}")
        print(f"Temperature值: {self.temperatures}")
        print(f"Prompt类型: {self.prompt_types}")
        print(f"{'=' * 60}")

        for temp in self.temperatures:
            for prompt_type in self.prompt_types:
                print(f"\n▶ 测试组合: Temperature={temp}, Prompt={prompt_type}")

                temp_results = []
                success_count = 0
                total_score = 0

                for sample in tqdm(samples, desc=f"处理进度"):
                    text = sample["disease_text"]
                    sample_id = sample["id"]

                    if not text:
                        continue

                    # 构建prompt
                    prompt = self.build_prompt(prompt_type, text)

                    # 调用API
                    response = self.extractor.call_dashscope_api(prompt, temperature=temp)

                    # 提取和解析结果
                    if response:
                        extracted_data = self.extractor.extract_json_from_response(response)
                        normalized_data = self.extractor.normalize_fields(extracted_data)

                        # 计算得分(非空字段数量)
                        score = sum(1 for value in normalized_data.values() if value and str(value).strip())

                        # 保存结果
                        result = {
                            "sample_id": sample_id,
                            "original_disease": sample.get("original_disease", ""),
                            "temperature": temp,
                            "prompt_type": prompt_type,
                            "raw_response": response[:200] + "..." if len(response) > 200 else response,
                            "extracted_data": normalized_data,
                            "score": score,
                            "max_score": 6  # 6个字段
                        }

                        temp_results.append(result)
                        success_count += 1
                        total_score += score

                        # API限速控制
                        time.sleep(0.3)  # 百炼API调用频率限制
                    else:
                        print(f"样本{sample_id}: API调用失败")

                # 计算统计信息
                avg_score = total_score / success_count if success_count > 0 else 0
                success_rate = (success_count / len(samples)) * 100 if samples else 0

                # 保存统计信息
                self.summary.append({
                    "temperature": temp,
                    "prompt_type": prompt_type,
                    "total_samples": len(samples),
                    "success_count": success_count,
                    "success_rate": round(success_rate, 2),
                    "avg_score": round(avg_score, 2),
                    "max_possible_score": 6
                })

                # 保存详细结果
                self.results.extend(temp_results)

                print(f"✓ 完成: 成功率={success_rate:.1f}%, 平均得分={avg_score:.2f}/6")

    def save_results(self, output_dir: str = "experiment_results_dashscope"):
        """
        保存实验结果

        Args:
            output_dir: 输出目录
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        timestamp = time.strftime("%Y%m%d_%H%M%S")

        # 保存汇总结果
        summary_df = pd.DataFrame(self.summary)
        summary_path = os.path.join(output_dir, f"experiment_summary_{timestamp}.xlsx")
        summary_df.to_excel(summary_path, index=False)

        # 保存详细结果
        detailed_data = []
        for result in self.results:
            detailed_result = {
                "sample_id": result["sample_id"],
                "original_disease": result["original_disease"],
                "temperature": result["temperature"],
                "prompt_type": result["prompt_type"],
                "score": result["score"],
                "max_score": result["max_score"],
                "raw_response": result["raw_response"]
            }

            # 添加提取的字段
            extracted = result["extracted_data"]
            for key in ["disease", "part", "department", "symptom", "drug", "infectious"]:
                detailed_result[key] = extracted.get(key, "")

            detailed_data.append(detailed_result)

        detailed_df = pd.DataFrame(detailed_data)
        detailed_path = os.path.join(output_dir, f"experiment_detailed_{timestamp}.xlsx")
        detailed_df.to_excel(detailed_path, index=False)

        # 保存为JSON格式
        json_path = os.path.join(output_dir, f"experiment_results_{timestamp}.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "summary": self.summary,
                "results": [
                    {
                        k: v for k, v in result.items()
                        if k != "extracted_data"
                    } | {"extracted_data": result["extracted_data"]}
                    for result in self.results
                ]
            }, f, ensure_ascii=False, indent=2)

        print(f"\n{'=' * 60}")
        print("实验结果已保存:")
        print(f"1. 汇总结果: {summary_path}")
        print(f"2. 详细结果: {detailed_path}")
        print(f"3. JSON格式: {json_path}")

        # 打印汇总统计
        print(f"\n实验汇总统计:")
        print(summary_df.to_string(index=False))

        return summary_df, detailed_df

    def analyze_results(self):
        """分析实验结果"""
        if not self.summary:
            print("没有实验结果可分析")
            return

        print(f"\n{'=' * 60}")
        print("实验结果分析")
        print(f"{'=' * 60}")

        summary_df = pd.DataFrame(self.summary)

        # 按temperature分析
        print("\n1. 不同Temperature的效果对比:")
        temp_groups = summary_df.groupby('temperature')
        for temp, group in temp_groups:
            print(f"\nTemperature = {temp}:")
            for _, row in group.iterrows():
                print(f"  Prompt={row['prompt_type']}: 成功率={row['success_rate']}%, 平均得分={row['avg_score']:.2f}")

        # 按prompt类型分析
        print("\n2. 不同Prompt类型的效果对比:")
        prompt_groups = summary_df.groupby('prompt_type')
        for prompt, group in prompt_groups:
            print(f"\nPrompt类型 = {prompt}:")
            for _, row in group.iterrows():
                print(
                    f"  Temperature={row['temperature']}: 成功率={row['success_rate']}%, 平均得分={row['avg_score']:.2f}")

        # 最佳组合
        best_result = summary_df.loc[summary_df['avg_score'].idxmax()]
        print(f"\n3. 最佳组合:")
        print(f"   Temperature: {best_result['temperature']}")
        print(f"   Prompt类型: {best_result['prompt_type']}")
        print(f"   平均得分: {best_result['avg_score']:.2f}/6")
        print(f"   成功率: {best_result['success_rate']}%")


def get_dashscope_api_key():
    """获取阿里云百炼API Key"""
    # 检查环境变量
    api_key = os.environ.get("DASHSCOPE_API_KEY")

    if not api_key:
        print("=" * 60)
        print("阿里云百炼API配置")
        print("=" * 60)
        print("\n请获取阿里云百炼API Key:")
        print("1. 访问 https://dashscope.aliyun.com/")
        print("2. 注册阿里云账号并完成实名认证")
        print("3. 进入控制台 -> API-KEY管理")
        print("4. 创建API Key")
        print("5. 新用户有免费额度(100万tokens)")

        api_key = input("\n请输入您的DashScope API Key: ").strip()

        if not api_key.startswith("sk-"):
            print("警告: API Key格式可能不正确,通常以'sk-'开头")

    return api_key


def main():
    """主函数"""

    print("=" * 60)
    print("疾病结构化数据提取实验 - 阿里云百炼版")
    print("=" * 60)

    # 1. 获取API Key
    api_key = get_dashscope_api_key()

    if not api_key:
        print("错误: 需要提供DashScope API Key")
        return

    # 2. 选择模型
    print("\n" + "=" * 60)
    print("选择模型:")
    print("1. qwen-max (最强能力,适合复杂任务)")
    print("2. qwen-plus (平衡性能)")
    print("3. qwen-turbo (最快速度)")

    model_choice = input("请选择模型 (默认1): ").strip()
    model_options = {
        "1": "qwen-max",
        "2": "qwen-plus",
        "3": "qwen-turbo"
    }

    model = model_options.get(model_choice, "qwen-max")
    print(f"使用模型: {model}")

    # 3. 初始化提取器
    extractor = DashScopeExtractor(api_key=api_key, model=model)

    # 4. 加载数据
    data_path = r"文件目录\疾病数据_清洗后.csv"

    if not os.path.exists(data_path):
        # 尝试其他可能的路径
        alternatives = [
            r"预处理生成的json文件位置",
            r"文件目录\疾病数据_清洗后.xlsx"
        ]

        for pattern in alternatives:
            if '*' in pattern:
                import glob
                matches = glob.glob(pattern)
                if matches:
                    data_path = matches[0]
                    break
            elif os.path.exists(pattern):
                data_path = pattern
                break
        else:
            # 让用户手动输入
            data_path = input("未找到数据文件,请输入完整路径: ").strip()
            if not os.path.exists(data_path):
                print(f"错误: 文件不存在: {data_path}")
                return

    print(f"使用数据文件: {data_path}")

    processor = DiseaseDataProcessor(data_path)
    if not processor.load_data():
        return

    # 5. 抽取样本
    samples = processor.sample_data(n=50, random_seed=42)
    if not samples:
        print("错误: 未能抽取到样本数据")
        return

    # 6. 显示样本信息
    print(f"\n样本示例 (前3个):")
    for i, sample in enumerate(samples[:3]):
        print(f"\n样本 {sample['id']}:")
        print(f"原始疾病: {sample.get('original_disease', 'N/A')}")
        print(f"文本预览: {sample['disease_text'][:100]}...")

    # 7. 确认开始实验
    confirm = input(f"\n是否开始实验? (将调用API {len(samples) * 9} 次) [y/n]: ").lower()
    if confirm != 'y':
        print("实验已取消")
        return

    # 8. 运行实验
    experiment = ExperimentRunner(extractor)
    experiment.run_experiment(samples)

    # 9. 保存结果
    output_dir = r"输出位置"
    experiment.save_results(output_dir)

    # 10. 分析结果
    experiment.analyze_results()

    print(f"\n{'=' * 60}")
    print("实验完成!")
    print(f"使用模型: {model}")
    print(f"已生成完整的实验报告")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    try:
        import dashscope
    except ImportError:
        import subprocess
        import sys

        subprocess.check_call([sys.executable, "-m", "pip", "install", "dashscope"])
        import dashscope

    main()

出现下图所示界面后输出API

(3)效果对比

100%是数据较少的简单小样本下的结果,非研究重点,比较得分找到最佳组合

结构化后的结果如下

二、接入Neo4j和TuGraph

(1)Neo4j

由于数据源不同,代码仅供参考

"""
Neo4j与大模型接口集成系统 
"""

import os
import json
import logging
from typing import Dict, List, Any, Optional, Tuple
from neo4j import GraphDatabase
import requests
import time
import sys
from datetime import datetime
import re

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('neo4j_llm.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


def test_neo4j_connection(uri: str, username: str, password: str) -> bool:
    """测试Neo4j连接"""
    test_uris = [
        "bolt://localhost:7687",  # 默认URI
        "bolt://127.0.0.1:7687",  # 本地IP
        "bolt://0.0.0.0:7687",  # 全零IP
        "neo4j://localhost:7687",  # neo4j协议
    ]

    for test_uri in test_uris:
        try:
            logger.info(f"测试连接: {test_uri}")
            driver = GraphDatabase.driver(
                test_uri,
                auth=(username, password),
                connection_timeout=5
            )
            with driver.session() as session:
                result = session.run("RETURN 1 as test")
                if result.single()["test"] == 1:
                    logger.info(f"✓ 连接成功: {test_uri}")
                    driver.close()
                    return test_uri
        except Exception as e:
            logger.warning(f"连接失败 {test_uri}: {str(e)[:100]}")
            continue

    return None


class SimpleLLMInterface:
    """简化版LLM接口,避免API问题"""

    def __init__(self):
        self.rules = {
            "查找所有节点": "MATCH (n) RETURN n LIMIT 100",
            "查找所有患者": "MATCH (p:Patient) RETURN p.name, p.age LIMIT 100",
            "查找高血压患者": "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '高血压'}) RETURN p.name, p.age LIMIT 100",
            "查询肺炎症状": "MATCH (d:Disease {name: '肺炎'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name LIMIT 100",
            "查询糖尿病药物": "MATCH (d:Disease {name: '糖尿病'})-[:TREATED_BY]->(m:Medicine) RETURN m.name LIMIT 100",
            "查询胃炎科室": "MATCH (d:Disease {name: '胃炎'})-[:TREATED_IN]->(de:Department) RETURN de.name LIMIT 100",
            "查找同时患有高血压和糖尿病的患者": "MATCH (p:Patient)-[:HAS_DISEASE]->(d1:Disease {name: '高血压'}), (p)-[:HAS_DISEASE]->(d2:Disease {name: '糖尿病'}) RETURN p.name, p.age LIMIT 100"
        }

    def generate_cypher(self, query: str) -> str:
        """基于规则生成Cypher"""
        query_lower = query.lower()

        # 关键词匹配
        if "所有节点" in query_lower or "全部数据" in query_lower:
            return "MATCH (n) RETURN n LIMIT 50"

        elif "患者" in query_lower:
            if "高血压" in query_lower:
                return "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '高血压'}) RETURN p.name, p.age LIMIT 100"
            elif "糖尿病" in query_lower:
                return "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '糖尿病'}) RETURN p.name, p.age LIMIT 100"
            elif "肺炎" in query_lower:
                return "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '肺炎'}) RETURN p.name, p.age LIMIT 100"
            else:
                return "MATCH (p:Patient) RETURN p.name, p.age LIMIT 100"

        elif "症状" in query_lower:
            if "肺炎" in query_lower:
                return "MATCH (d:Disease {name: '肺炎'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name LIMIT 100"
            elif "高血压" in query_lower:
                return "MATCH (d:Disease {name: '高血压'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name LIMIT 100"

        elif "药物" in query_lower or "药品" in query_lower:
            if "糖尿病" in query_lower:
                return "MATCH (d:Disease {name: '糖尿病'})-[:TREATED_BY]->(m:Medicine) RETURN m.name LIMIT 100"
            elif "高血压" in query_lower:
                return "MATCH (d:Disease {name: '高血压'})-[:TREATED_BY]->(m:Medicine) RETURN m.name LIMIT 100"

        elif "科室" in query_lower or "部门" in query_lower:
            if "胃炎" in query_lower:
                return "MATCH (d:Disease {name: '胃炎'})-[:TREATED_IN]->(de:Department) RETURN de.name LIMIT 100"
            elif "肺炎" in query_lower:
                return "MATCH (d:Disease {name: '肺炎'})-[:TREATED_IN]->(de:Department) RETURN de.name LIMIT 100"

        # 默认查询
        return "MATCH (n) RETURN n LIMIT 30"


class Neo4jManager:
    """Neo4j管理类"""

    def __init__(self, uri: str, username: str, password: str):
        self.uri = uri
        self.username = username
        self.password = password
        self.driver = None
        self.connected = False

    def connect(self) -> bool:
        """连接到Neo4j"""
        try:
            logger.info(f"正在连接到Neo4j: {self.uri}")

            # 创建驱动
            self.driver = GraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password),
                max_connection_lifetime=7200,
                connection_timeout=15
            )

            # 测试连接
            with self.driver.session() as session:
                result = session.run("RETURN 'Neo4j Connection Test' as test")
                test_result = result.single()["test"]

                if test_result == "Neo4j Connection Test":
                    self.connected = True
                    logger.info("✓ Neo4j连接成功")
                    return True
                else:
                    logger.error(f"连接测试返回异常: {test_result}")
                    return False

        except Exception as e:
            logger.error(f"连接失败: {e}")
            self.connected = False
            return False

    def execute_query(self, cypher: str, params: Dict = None) -> Tuple[bool, List[Dict], str]:
        """执行Cypher查询"""
        if not self.connected or not self.driver:
            return False, [], "未连接到数据库"

        try:
            with self.driver.session() as session:
                result = session.run(cypher, params or {})
                records = [dict(record) for record in result]
                return True, records, ""

        except Exception as e:
            error_msg = str(e)
            logger.error(f"查询执行失败: {error_msg}")
            return False, [], error_msg

    def get_schema_info(self) -> Dict:
        """获取数据库Schema信息"""
        if not self.connected:
            return {}

        try:
            schema = {
                "node_labels": [],
                "relationship_types": [],
                "properties": {}
            }

            with self.driver.session() as session:
                # 获取节点标签
                result = session.run("CALL db.labels()")
                schema["node_labels"] = [record["label"] for record in result]

                # 获取关系类型
                result = session.run("CALL db.relationshipTypes()")
                schema["relationship_types"] = [record["relationshipType"] for record in result]

                logger.info(f"数据库Schema: 节点标签={schema['node_labels']}, 关系类型={schema['relationship_types']}")
                return schema

        except Exception as e:
            logger.error(f"获取Schema失败: {e}")
            return {}

    def close(self):
        """关闭连接"""
        if self.driver:
            self.driver.close()
            self.connected = False
            logger.info("Neo4j连接已关闭")


def create_medical_demo_data(manager: Neo4jManager) -> bool:
    """创建医疗演示数据"""
    try:
        logger.info("开始创建医疗演示数据...")

        # 检查是否已存在数据
        success, records, error = manager.execute_query("MATCH (n) RETURN count(n) as count LIMIT 1")
        if success and records and records[0].get("count", 0) > 0:
            confirm = input("数据库中已有数据,是否重新创建? (y/n): ").lower()
            if confirm != 'y':
                logger.info("跳过数据创建")
                return True

        # 清空现有数据
        manager.execute_query("MATCH (n) DETACH DELETE n")
        logger.info("已清空现有数据")

        # 创建约束
        constraints = [
            "CREATE CONSTRAINT disease_name_unique IF NOT EXISTS FOR (d:Disease) REQUIRE d.name IS UNIQUE",
            "CREATE CONSTRAINT symptom_name_unique IF NOT EXISTS FOR (s:Symptom) REQUIRE s.name IS UNIQUE",
            "CREATE CONSTRAINT patient_id_unique IF NOT EXISTS FOR (p:Patient) REQUIRE p.id IS UNIQUE",
            "CREATE CONSTRAINT medicine_name_unique IF NOT EXISTS FOR (m:Medicine) REQUIRE m.name IS UNIQUE",
            "CREATE CONSTRAINT dept_name_unique IF NOT EXISTS FOR (d:Department) REQUIRE d.name IS UNIQUE"
        ]

        for constraint in constraints:
            try:
                manager.execute_query(constraint)
            except:
                pass  # 约束可能已存在

        logger.info("约束创建完成")

        # 创建节点
        node_queries = [
            # 疾病节点
            """CREATE (:Disease {name: '高血压', description: '血压持续升高', type: '心血管疾病'})""",
            """CREATE (:Disease {name: '糖尿病', description: '血糖代谢异常', type: '代谢疾病'})""",
            """CREATE (:Disease {name: '肺炎', description: '肺部感染', type: '呼吸系统疾病'})""",
            """CREATE (:Disease {name: '胃炎', description: '胃部炎症', type: '消化系统疾病'})""",
            """CREATE (:Disease {name: '流感', description: '流行性感冒', type: '传染性疾病'})""",

            # 症状节点
            """CREATE (:Symptom {name: '头痛', severity: '中度'})""",
            """CREATE (:Symptom {name: '发热', severity: '高度'})""",
            """CREATE (:Symptom {name: '咳嗽', severity: '中度'})""",
            """CREATE (:Symptom {name: '恶心', severity: '轻度'})""",
            """CREATE (:Symptom {name: '乏力', severity: '轻度'})""",

            # 药物节点
            """CREATE (:Medicine {name: '降压药', type: '处方药', dosage: '10mg'})""",
            """CREATE (:Medicine {name: '胰岛素', type: '处方药', dosage: '5mg'})""",
            """CREATE (:Medicine {name: '抗生素', type: '处方药', dosage: '500mg'})""",
            """CREATE (:Medicine {name: '胃药', type: '非处方药', dosage: '20mg'})""",
            """CREATE (:Medicine {name: '抗病毒药', type: '处方药', dosage: '100mg'})""",

            # 科室节点
            """CREATE (:Department {name: '心血管科', type: '内科', location: '3楼'})""",
            """CREATE (:Department {name: '内分泌科', type: '内科', location: '4楼'})""",
            """CREATE (:Department {name: '呼吸内科', type: '内科', location: '5楼'})""",
            """CREATE (:Department {name: '消化内科', type: '内科', location: '6楼'})""",
            """CREATE (:Department {name: '急诊科', type: '急诊', location: '1楼'})""",

            # 患者节点
            """CREATE (:Patient {id: 'P001', name: '张三', age: 45, gender: '男'})""",
            """CREATE (:Patient {id: 'P002', name: '李四', age: 62, gender: '女'})""",
            """CREATE (:Patient {id: 'P003', name: '王五', age: 38, gender: '男'})""",
            """CREATE (:Patient {id: 'P004', name: '赵六', age: 55, gender: '女'})"""
        ]

        for query in node_queries:
            manager.execute_query(query)
        logger.info("节点创建完成")

        # 创建关系
        relationship_queries = [
            # 疾病-症状关系
            """MATCH (d:Disease {name: '高血压'}), (s:Symptom {name: '头痛'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)""",
            """MATCH (d:Disease {name: '肺炎'}), (s:Symptom {name: '发热'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)""",
            """MATCH (d:Disease {name: '肺炎'}), (s:Symptom {name: '咳嗽'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)""",
            """MATCH (d:Disease {name: '流感'}), (s:Symptom {name: '发热'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)""",
            """MATCH (d:Disease {name: '胃炎'}), (s:Symptom {name: '恶心'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)""",

            # 疾病-药物关系
            """MATCH (d:Disease {name: '高血压'}), (m:Medicine {name: '降压药'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)""",
            """MATCH (d:Disease {name: '糖尿病'}), (m:Medicine {name: '胰岛素'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)""",
            """MATCH (d:Disease {name: '肺炎'}), (m:Medicine {name: '抗生素'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)""",
            """MATCH (d:Disease {name: '胃炎'}), (m:Medicine {name: '胃药'}) CREATE (d)-[:TREATED_BY {effectiveness: '中'}]->(m)""",

            # 疾病-科室关系
            """MATCH (d:Disease {name: '高血压'}), (de:Department {name: '心血管科'}) CREATE (d)-[:TREATED_IN]->(de)""",
            """MATCH (d:Disease {name: '糖尿病'}), (de:Department {name: '内分泌科'}) CREATE (d)-[:TREATED_IN]->(de)""",
            """MATCH (d:Disease {name: '肺炎'}), (de:Department {name: '呼吸内科'}) CREATE (d)-[:TREATED_IN]->(de)""",
            """MATCH (d:Disease {name: '胃炎'}), (de:Department {name: '消化内科'}) CREATE (d)-[:TREATED_IN]->(de)""",

            # 患者-疾病关系
            """MATCH (p:Patient {id: 'P001'}), (d:Disease {name: '高血压'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-01-15'}]->(d)""",
            """MATCH (p:Patient {id: 'P002'}), (d:Disease {name: '糖尿病'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-02-20'}]->(d)""",
            """MATCH (p:Patient {id: 'P003'}), (d:Disease {name: '肺炎'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-03-10'}]->(d)""",
            """MATCH (p:Patient {id: 'P004'}), (d:Disease {name: '胃炎'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-04-05'}]->(d)""",
            """MATCH (p:Patient {id: 'P001'}), (d:Disease {name: '糖尿病'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2022-12-10'}]->(d)"""
        ]

        for query in relationship_queries:
            manager.execute_query(query)
        logger.info("关系创建完成")

        # 统计数据
        success, node_count_result, _ = manager.execute_query("MATCH (n) RETURN count(n) as node_count")
        success, rel_count_result, _ = manager.execute_query("MATCH ()-[r]->() RETURN count(r) as rel_count")

        if node_count_result and rel_count_result:
            node_count = node_count_result[0].get("node_count", 0)
            rel_count = rel_count_result[0].get("rel_count", 0)
            logger.info(f"数据创建完成: 节点数={node_count}, 关系数={rel_count}")
            return True

        return False

    except Exception as e:
        logger.error(f"创建数据失败: {e}")
        return False


def run_interactive_mode(manager: Neo4jManager, llm: SimpleLLMInterface):
    """运行交互模式"""
    print("\n" + "=" * 60)
    print("Neo4j自然语言查询系统")
    print("=" * 60)
    print("支持查询示例:")
    print("1. 查找所有节点")
    print("2. 查找所有患者")
    print("3. 查询高血压患者")
    print("4. 查询肺炎症状")
    print("5. 查询糖尿病药物")
    print("6. 查询胃炎科室")
    print("输入 'help' 查看帮助")
    print("输入 'quit' 退出")
    print("=" * 60)

    while True:
        try:
            print("\n" + "-" * 40)
            query = input("请输入查询语句: ").strip()

            if query.lower() in ['quit', 'exit', 'q']:
                print("退出查询系统")
                break

            if query.lower() in ['help', '?']:
                print("""
查询示例:
1. 查找所有节点
2. 查找所有患者
3. 查询高血压患者
4. 查询肺炎症状
5. 查询糖尿病药物
6. 查询胃炎科室
7. 查找同时患有高血压和糖尿病的患者
                """)
                continue

            if not query:
                continue

            print(f"\n正在处理查询: {query}")

            # 生成Cypher
            cypher = llm.generate_cypher(query)
            print(f"生成的Cypher: {cypher}")

            # 执行查询
            success, results, error = manager.execute_query(cypher)

            if success:
                if results:
                    print(f"✓ 查询成功! 找到 {len(results)} 条结果")
                    print("\n查询结果:")
                    for i, record in enumerate(results[:10], 1):  # 限制显示10条
                        # 格式化显示
                        formatted = []
                        for key, value in record.items():
                            if isinstance(value, dict):
                                # 如果是节点属性
                                if 'name' in value:
                                    formatted.append(f"{key}.name={value['name']}")
                                elif 'id' in value:
                                    formatted.append(f"{key}.id={value['id']}")
                                else:
                                    formatted.append(f"{key}={str(value)[:30]}")
                            else:
                                formatted.append(f"{key}={str(value)[:30]}")
                        print(f"{i}. {' | '.join(formatted)}")

                    if len(results) > 10:
                        print(f"... 还有 {len(results) - 10} 条记录未显示")
                else:
                    print("✓ 查询成功,但未找到匹配结果")
            else:
                print(f"❌ 查询失败: {error}")

        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        except Exception as e:
            print(f"❌ 发生错误: {e}")


def save_query_results(query: str, cypher: str, results: List[Dict], filename: str = None):
    """保存查询结果到文件"""
    if filename is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"query_result_{timestamp}.json"

    data = {
        "timestamp": datetime.now().isoformat(),
        "query": query,
        "cypher": cypher,
        "result_count": len(results),
        "results": results
    }

    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"结果已保存到: {filename}")
        return True
    except Exception as e:
        print(f"保存结果失败: {e}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("Neo4j自然语言查询系统 - 简化版")
    print("=" * 60)

    # 配置信息
    DEFAULT_URI = "bolt://localhost:7687"
    DEFAULT_USER = "neo4j"
    DEFAULT_PASSWORD = "12345678"

    print("当前配置:")
    print(f"  URI: {DEFAULT_URI}")
    print(f"  用户名: {DEFAULT_USER}")

    # 密码输入
    password = input(f"请输入Neo4j密码 (默认: {DEFAULT_PASSWORD}): ").strip()
    if not password:
        password = DEFAULT_PASSWORD

    # 测试连接
    print("\n正在测试Neo4j连接...")
    working_uri = test_neo4j_connection(DEFAULT_URI, DEFAULT_USER, password)

    if not working_uri:
        print("❌ 无法连接到Neo4j数据库")
        print("\n请检查:")
        print("1. Neo4j数据库是否正在运行")
        print("2. 用户名和密码是否正确")
        print("3. 防火墙设置")
        print("\n您可以尝试:")
        print(f"  - 在浏览器中访问: http://localhost:7474")
        print(f"  - 使用用户名: {DEFAULT_USER}")
        print(f"  - 使用您设置的密码")
        return

    print(f"✓ 使用URI: {working_uri}")

    # 初始化管理器
    manager = Neo4jManager(working_uri, DEFAULT_USER, password)
    if not manager.connect():
        print("❌ 连接失败")
        return

    # 检查现有数据
    print("\n检查数据库状态...")
    schema = manager.get_schema_info()
    if schema.get("node_labels"):
        print(f"数据库包含: {len(schema['node_labels'])} 个节点标签")
    else:
        print("数据库为空或无法获取Schema信息")

    # 询问是否创建示例数据
    create_data = input("\n是否创建医疗知识图谱示例数据? (y/n): ").lower()
    if create_data == 'y':
        if create_medical_demo_data(manager):
            print("✓ 示例数据创建成功")
        else:
            print("❌ 示例数据创建失败")

    # 初始化LLM接口
    llm = SimpleLLMInterface()

    # 运行测试查询
    print("\n" + "=" * 60)
    print("运行测试查询")
    print("=" * 60)

    test_queries = [
        "查找所有节点",
        "查找所有患者",
        "查询高血压患者",
        "查询肺炎症状",
        "查询糖尿病药物"
    ]

    for i, query in enumerate(test_queries, 1):
        print(f"\n测试 {i}: {query}")
        cypher = llm.generate_cypher(query)
        print(f"  生成的Cypher: {cypher}")

        success, results, error = manager.execute_query(cypher)
        if success:
            print(f"  ✓ 成功, 找到 {len(results)} 条结果")
        else:
            print(f"  ❌ 失败: {error}")

    # 运行交互模式
    run_interactive_mode(manager, llm)

    # 关闭连接
    manager.close()
    print("\n✓ 程序正常退出")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ 程序异常终止: {e}")
        import traceback

        traceback.print_exc()

(2)TuGraph

docker启动tugraph

"""
TuGraph图数据库与大模型接口集成系统
"""

import os
import json
import requests
import logging
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import sys
import time
import re

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('tugraph_final.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class TuGraphManager:
    """TuGraph图数据库管理器 """

    def __init__(self, host: str = "http://localhost:7070",
                 username: str = "admin",
                 password: str = "73@TuGraph",
                 graph_name: str = "default"):
        self.base_url = host.rstrip('/')
        self.username = username
        self.password = password
        self.graph_name = graph_name
        self.token = None
        self.connected = False

    def connect(self) -> bool:
        """连接到TuGraph"""
        try:
            # 登录获取token
            login_url = f"{self.base_url}/login"
            login_data = {
                "user": self.username,
                "password": self.password
            }

            logger.info(f"正在登录TuGraph: {self.base_url}")
            response = requests.post(login_url, json=login_data, timeout=10)

            if response.status_code == 200:
                data = response.json()
                self.token = data.get("jwt")
                if self.token:
                    self.connected = True
                    logger.info(f"✓ TuGraph连接成功")

                    # 测试查询验证连接
                    test_success, test_result, _ = self.execute_cypher("RETURN 1 as test")
                    if test_success:
                        logger.info("✓ 连接测试通过")
                        return True
                else:
                    logger.error("登录响应中没有token")
                    return False
            else:
                logger.error(f"登录失败: {response.status_code} - {response.text}")
                return False

        except Exception as e:
            logger.error(f"连接异常: {e}")
            return False

    def _get_headers(self) -> Dict:
        """获取请求头"""
        headers = {"Content-Type": "application/json"}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        return headers

    def execute_cypher(self, cypher: str, timeout: int = 30) -> Tuple[bool, List, str]:
        """
        执行Cypher查询
        返回格式: (成功标志, 结果列表, 错误信息)
        """
        if not self.connected:
            return False, [], "未连接到数据库"

        try:
            url = f"{self.base_url}/cypher"
            data = {
                "graph": self.graph_name,
                "script": cypher
            }

            logger.debug(f"执行Cypher: {cypher[:100]}...")
            response = requests.post(
                url,
                headers=self._get_headers(),
                json=data,
                timeout=timeout
            )

            if response.status_code == 200:
                result = response.json()
                # TuGraph返回格式: {"header": [...], "result": [...], "size": n, "elapsed": x}
                if isinstance(result, dict) and "result" in result:
                    return True, result["result"], ""
                else:
                    return True, result if isinstance(result, list) else [result], ""
            else:
                error_msg = f"HTTP {response.status_code}"
                try:
                    error_data = response.json()
                    error_msg += f": {error_data.get('error_message', str(error_data))}"
                except:
                    error_msg += f": {response.text}"
                return False, [], error_msg

        except requests.exceptions.Timeout:
            return False, [], "请求超时"
        except Exception as e:
            return False, [], f"执行异常: {str(e)}"

    def create_vertex_label(self, label_name: str, properties: List[Dict]) -> bool:
        """创建顶点标签"""
        try:
            # TuGraph创建顶点标签的Cypher
            properties_str = ", ".join([f"{prop['name']}:{prop['type']}" for prop in properties])
            cypher = f"CALL db.createVertexLabel('{label_name}', '{properties_str}')"
            success, result, error = self.execute_cypher(cypher)
            if success:
                logger.info(f"✓ 创建顶点标签: {label_name}")
                return True
            else:
                logger.warning(f"创建顶点标签 {label_name} 失败: {error}")
                return False
        except Exception as e:
            logger.error(f"创建顶点标签异常: {e}")
            return False

    def create_edge_label(self, label_name: str, properties: List[Dict]) -> bool:
        """创建边标签"""
        try:
            properties_str = ", ".join([f"{prop['name']}:{prop['type']}" for prop in properties])
            cypher = f"CALL db.createEdgeLabel('{label_name}', '{properties_str}')"
            success, result, error = self.execute_cypher(cypher)
            if success:
                logger.info(f"✓ 创建边标签: {label_name}")
                return True
            else:
                logger.warning(f"创建边标签 {label_name} 失败: {error}")
                return False
        except Exception as e:
            logger.error(f"创建边标签异常: {e}")
            return False


def load_and_filter_experiment_data(data_dir: str) -> Tuple[List[Dict], Dict]:
    """加载并过滤实验数据,只保留有效的疾病信息"""
    data_files = []

    if os.path.exists(data_dir):
        for file in os.listdir(data_dir):
            if file.endswith('.json') and 'experiment_results' in file:
                data_files.append(os.path.join(data_dir, file))

    if not data_files:
        logger.warning("未找到实验数据文件")
        return [], {}

    # 使用最新的文件
    latest_file = max(data_files, key=os.path.getctime)
    logger.info(f"加载数据文件: {latest_file}")

    try:
        with open(latest_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        diseases = []
        stats = {
            'total_samples': 0,
            'valid_diseases': 0,
            'symptoms_found': 0,
            'drugs_found': 0,
            'departments_found': 0
        }

        if 'results' in data:
            stats['total_samples'] = len(data['results'])

            for result in data['results']:
                if 'extracted_data' in result:
                    disease_info = result['extracted_data']
                    disease_name = disease_info.get('disease', '').strip()

                    # 过滤无效的疾病名称
                    if disease_name and len(disease_name) <= 50 and not any(
                            char in disease_name for char in ['<', '>', '&', ';', ':', '"']):
                        stats['valid_diseases'] += 1

                        # 处理症状
                        symptoms_str = disease_info.get('symptom', '')
                        symptoms = []
                        if symptoms_str:
                            # 分割并清理症状
                            symptoms = [s.strip() for s in symptoms_str.split(',') if
                                        s.strip() and len(s.strip()) <= 30]
                            stats['symptoms_found'] += len(symptoms)

                        # 处理药物
                        drugs_str = disease_info.get('drug', '')
                        drugs = []
                        if drugs_str:
                            drugs = [d.strip() for d in drugs_str.split(',') if d.strip() and len(d.strip()) <= 30]
                            stats['drugs_found'] += len(drugs)

                        # 处理科室
                        department = disease_info.get('department', '').strip()
                        if department and len(department) <= 30:
                            stats['departments_found'] += 1
                        else:
                            department = ''

                        diseases.append({
                            'name': disease_name,
                            'part': disease_info.get('part', '').strip()[:30],
                            'department': department,
                            'symptoms': symptoms,
                            'drugs': drugs,
                            'infectious': disease_info.get('infectious', '').strip()[:10]
                        })

        logger.info(f"数据分析结果: 有效疾病={stats['valid_diseases']}, 症状={stats['symptoms_found']}, "
                    f"药物={stats['drugs_found']}, 科室={stats['departments_found']}")

        return diseases, stats

    except Exception as e:
        logger.error(f"加载数据失败: {e}")
        return [], {}


def create_simple_medical_graph(manager: TuGraphManager) -> bool:
    """创建简单的医疗知识图谱(不依赖实验数据)"""
    try:
        logger.info("开始创建简单医疗知识图谱...")

        # 1. 清空现有数据
        logger.info("清空现有数据...")
        manager.execute_cypher("MATCH (n) DETACH DELETE n")

        # 2. 创建顶点标签
        logger.info("创建顶点标签...")

        vertex_labels = [
            {
                'name': 'Disease',
                'properties': [
                    {'name': 'id', 'type': 'STRING'},
                    {'name': 'name', 'type': 'STRING'},
                    {'name': 'description', 'type': 'STRING'},
                    {'name': 'infectious', 'type': 'STRING'}
                ]
            },
            {
                'name': 'Symptom',
                'properties': [
                    {'name': 'id', 'type': 'STRING'},
                    {'name': 'name', 'type': 'STRING'},
                    {'name': 'severity', 'type': 'STRING'}
                ]
            },
            {
                'name': 'Medicine',
                'properties': [
                    {'name': 'id', 'type': 'STRING'},
                    {'name': 'name', 'type': 'STRING'},
                    {'name': 'type', 'type': 'STRING'}
                ]
            },
            {
                'name': 'Patient',
                'properties': [
                    {'name': 'id', 'type': 'STRING'},
                    {'name': 'name', 'type': 'STRING'},
                    {'name': 'age', 'type': 'INT32'},
                    {'name': 'gender', 'type': 'STRING'}
                ]
            },
            {
                'name': 'Department',
                'properties': [
                    {'name': 'id', 'type': 'STRING'},
                    {'name': 'name', 'type': 'STRING'},
                    {'name': 'location', 'type': 'STRING'}
                ]
            }
        ]

        for label in vertex_labels:
            manager.create_vertex_label(label['name'], label['properties'])

        # 3. 创建边标签
        logger.info("创建边标签...")

        edge_labels = [
            {
                'name': 'HAS_SYMPTOM',
                'properties': [
                    {'name': 'severity', 'type': 'STRING'}
                ]
            },
            {
                'name': 'TREATED_BY',
                'properties': [
                    {'name': 'effectiveness', 'type': 'STRING'}
                ]
            },
            {
                'name': 'TREATED_IN',
                'properties': []
            },
            {
                'name': 'HAS_DISEASE',
                'properties': [
                    {'name': 'diagnosed_date', 'type': 'STRING'}
                ]
            }
        ]

        for label in edge_labels:
            manager.create_edge_label(label['name'], label['properties'])

        time.sleep(2)  # 等待标签创建完成

        # 4. 创建节点
        logger.info("创建节点...")

        # 疾病节点
        diseases = [
            "CREATE (:Disease {id: 'D001', name: '高血压', description: '血压持续升高', infectious: '否'})",
            "CREATE (:Disease {id: 'D002', name: '糖尿病', description: '血糖代谢异常', infectious: '否'})",
            "CREATE (:Disease {id: 'D003', name: '肺炎', description: '肺部感染', infectious: '是'})",
            "CREATE (:Disease {id: 'D004', name: '胃炎', description: '胃部炎症', infectious: '否'})",
            "CREATE (:Disease {id: 'D005', name: '流感', description: '流行性感冒', infectious: '是'})"
        ]

        for cypher in diseases:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个疾病节点")

        # 症状节点
        symptoms = [
            "CREATE (:Symptom {id: 'S001', name: '头痛', severity: '中度'})",
            "CREATE (:Symptom {id: 'S002', name: '发热', severity: '高度'})",
            "CREATE (:Symptom {id: 'S003', name: '咳嗽', severity: '中度'})",
            "CREATE (:Symptom {id: 'S004', name: '恶心', severity: '轻度'})",
            "CREATE (:Symptom {id: 'S005', name: '乏力', severity: '轻度'})"
        ]

        for cypher in symptoms:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个症状节点")

        # 药物节点
        medicines = [
            "CREATE (:Medicine {id: 'M001', name: '降压药', type: '处方药'})",
            "CREATE (:Medicine {id: 'M002', name: '胰岛素', type: '处方药'})",
            "CREATE (:Medicine {id: 'M003', name: '抗生素', type: '处方药'})",
            "CREATE (:Medicine {id: 'M004', name: '胃药', type: '非处方药'})",
            "CREATE (:Medicine {id: 'M005', name: '抗病毒药', type: '处方药'})"
        ]

        for cypher in medicines:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个药物节点")

        # 科室节点
        departments = [
            "CREATE (:Department {id: 'DE001', name: '心血管科', location: '3楼'})",
            "CREATE (:Department {id: 'DE002', name: '内分泌科', location: '4楼'})",
            "CREATE (:Department {id: 'DE003', name: '呼吸内科', location: '5楼'})",
            "CREATE (:Department {id: 'DE004', name: '消化内科', location: '6楼'})",
            "CREATE (:Department {id: 'DE005', name: '急诊科', location: '1楼'})"
        ]

        for cypher in departments:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个科室节点")

        # 患者节点
        patients = [
            "CREATE (:Patient {id: 'P001', name: '张三', age: 45, gender: '男'})",
            "CREATE (:Patient {id: 'P002', name: '李四', age: 60, gender: '女'})",
            "CREATE (:Patient {id: 'P003', name: '王五', age: 35, gender: '男'})",
            "CREATE (:Patient {id: 'P004', name: '赵六', age: 50, gender: '女'})"
        ]

        for cypher in patients:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了4个患者节点")

        # 5. 创建关系
        logger.info("创建关系...")

        # 疾病-症状关系
        disease_symptom_relations = [
            "MATCH (d:Disease {id: 'D001'}), (s:Symptom {id: 'S001'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)",
            "MATCH (d:Disease {id: 'D003'}), (s:Symptom {id: 'S002'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)",
            "MATCH (d:Disease {id: 'D003'}), (s:Symptom {id: 'S003'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)",
            "MATCH (d:Disease {id: 'D005'}), (s:Symptom {id: 'S002'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)",
            "MATCH (d:Disease {id: 'D004'}), (s:Symptom {id: 'S004'}) CREATE (d)-[:HAS_SYMPTOM {severity: '常见'}]->(s)"
        ]

        for cypher in disease_symptom_relations:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个疾病-症状关系")

        # 疾病-药物关系
        disease_medicine_relations = [
            "MATCH (d:Disease {id: 'D001'}), (m:Medicine {id: 'M001'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)",
            "MATCH (d:Disease {id: 'D002'}), (m:Medicine {id: 'M002'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)",
            "MATCH (d:Disease {id: 'D003'}), (m:Medicine {id: 'M003'}) CREATE (d)-[:TREATED_BY {effectiveness: '高'}]->(m)",
            "MATCH (d:Disease {id: 'D004'}), (m:Medicine {id: 'M004'}) CREATE (d)-[:TREATED_BY {effectiveness: '中'}]->(m)",
            "MATCH (d:Disease {id: 'D005'}), (m:Medicine {id: 'M005'}) CREATE (d)-[:TREATED_BY {effectiveness: '中'}]->(m)"
        ]

        for cypher in disease_medicine_relations:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个疾病-药物关系")

        # 疾病-科室关系
        disease_dept_relations = [
            "MATCH (d:Disease {id: 'D001'}), (de:Department {id: 'DE001'}) CREATE (d)-[:TREATED_IN]->(de)",
            "MATCH (d:Disease {id: 'D002'}), (de:Department {id: 'DE002'}) CREATE (d)-[:TREATED_IN]->(de)",
            "MATCH (d:Disease {id: 'D003'}), (de:Department {id: 'DE003'}) CREATE (d)-[:TREATED_IN]->(de)",
            "MATCH (d:Disease {id: 'D004'}), (de:Department {id: 'DE004'}) CREATE (d)-[:TREATED_IN]->(de)",
            "MATCH (d:Disease {id: 'D005'}), (de:Department {id: 'DE005'}) CREATE (d)-[:TREATED_IN]->(de)"
        ]

        for cypher in disease_dept_relations:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个疾病-科室关系")

        # 患者-疾病关系
        patient_disease_relations = [
            "MATCH (p:Patient {id: 'P001'}), (d:Disease {id: 'D001'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-01-15'}]->(d)",
            "MATCH (p:Patient {id: 'P002'}), (d:Disease {id: 'D002'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-02-20'}]->(d)",
            "MATCH (p:Patient {id: 'P003'}), (d:Disease {id: 'D003'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-03-10'}]->(d)",
            "MATCH (p:Patient {id: 'P004'}), (d:Disease {id: 'D004'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2023-04-05'}]->(d)",
            "MATCH (p:Patient {id: 'P001'}), (d:Disease {id: 'D002'}) CREATE (p)-[:HAS_DISEASE {diagnosed_date: '2022-12-10'}]->(d)"
        ]

        for cypher in patient_disease_relations:
            manager.execute_cypher(cypher)
        logger.info("✓ 创建了5个患者-疾病关系")

        # 6. 验证数据
        logger.info("验证数据...")

        # 统计节点总数
        success, result, error = manager.execute_cypher("MATCH (n) RETURN count(n) as node_count")
        if success and result and len(result) > 0 and len(result[0]) > 0:
            node_count = result[0][0]
            logger.info(f"总节点数: {node_count}")

        # 统计关系总数
        success, result, error = manager.execute_cypher("MATCH ()-[r]->() RETURN count(r) as rel_count")
        if success and result and len(result) > 0 and len(result[0]) > 0:
            rel_count = result[0][0]
            logger.info(f"总关系数: {rel_count}")

        # 按类型统计节点
        node_types = ["Disease", "Symptom", "Medicine", "Patient", "Department"]
        for node_type in node_types:
            success, result, error = manager.execute_cypher(f"MATCH (n:{node_type}) RETURN count(n) as count")
            if success and result and len(result) > 0 and len(result[0]) > 0:
                logger.info(f"  {node_type}节点数: {result[0][0]}")

        logger.info("✓ 简单医疗知识图谱创建完成")
        return True

    except Exception as e:
        logger.error(f"创建知识图谱失败: {e}")
        import traceback
        traceback.print_exc()
        return False


class SmartCypherGenerator:
    """智能Cypher生成器"""

    def __init__(self):
        self.query_templates = {
            # 基础查询
            "查找所有节点": "MATCH (n) RETURN n LIMIT 20",
            "统计节点数量": "MATCH (n) RETURN count(n) as 总数",
            "查看数据库结构": "CALL db.vertexLabels()",

            # 疾病查询
            "查询疾病信息": "MATCH (d:Disease) RETURN d.id, d.name, d.description LIMIT 10",
            "查询高血压信息": "MATCH (d:Disease {name: '高血压'}) RETURN d.id, d.name, d.description",
            "查询糖尿病信息": "MATCH (d:Disease {name: '糖尿病'}) RETURN d.id, d.name, d.description",
            "查询肺炎信息": "MATCH (d:Disease {name: '肺炎'}) RETURN d.id, d.name, d.description",

            # 症状查询
            "查询症状信息": "MATCH (s:Symptom) RETURN s.id, s.name, s.severity LIMIT 10",
            "查询高血压症状": "MATCH (d:Disease {name: '高血压'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name",
            "查询肺炎症状": "MATCH (d:Disease {name: '肺炎'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name",

            # 药物查询
            "查询药物信息": "MATCH (m:Medicine) RETURN m.id, m.name, m.type LIMIT 10",
            "查询糖尿病药物": "MATCH (d:Disease {name: '糖尿病'})-[:TREATED_BY]->(m:Medicine) RETURN m.name",
            "查询高血压药物": "MATCH (d:Disease {name: '高血压'})-[:TREATED_BY]->(m:Medicine) RETURN m.name",

            # 患者查询
            "查询患者信息": "MATCH (p:Patient) RETURN p.id, p.name, p.age, p.gender LIMIT 10",
            "查询高血压患者": "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '高血压'}) RETURN p.name, p.age, p.gender",
            "查询糖尿病患者": "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '糖尿病'}) RETURN p.name, p.age, p.gender",

            # 科室查询
            "查询科室信息": "MATCH (de:Department) RETURN de.id, de.name, de.location LIMIT 10",
            "查询心血管科室": "MATCH (de:Department {name: '心血管科'}) RETURN de.name, de.location",

            # 关系查询
            "查询所有关系": "MATCH ()-[r]->() RETURN type(r) as 关系类型, count(r) as 数量",
            "查询疾病关联": "MATCH (d:Disease)-[r]->(n) RETURN d.name, type(r), labels(n)[0], n.name LIMIT 20",
        }

    def generate_cypher(self, query: str) -> str:
        """根据自然语言生成Cypher查询"""
        query_lower = query.lower()

        # 直接匹配模板
        for template_query, cypher in self.query_templates.items():
            if template_query.lower() in query_lower or query_lower in template_query.lower():
                return cypher

        # 关键词匹配
        if "所有" in query_lower and "节点" in query_lower:
            return "MATCH (n) RETURN n LIMIT 20"
        elif "统计" in query_lower or "计数" in query_lower or "数量" in query_lower:
            return "MATCH (n) RETURN count(n) as 总数"
        elif "疾病" in query_lower:
            if "症状" in query_lower:
                if "高血压" in query_lower:
                    return "MATCH (d:Disease {name: '高血压'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name"
                elif "肺炎" in query_lower:
                    return "MATCH (d:Disease {name: '肺炎'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name"
                elif "糖尿病" in query_lower:
                    return "MATCH (d:Disease {name: '糖尿病'})-[:HAS_SYMPTOM]->(s:Symptom) RETURN s.name"
                else:
                    return "MATCH (d:Disease)-[:HAS_SYMPTOM]->(s:Symptom) RETURN d.name, s.name LIMIT 10"
            elif "药物" in query_lower or "药" in query_lower:
                if "糖尿病" in query_lower:
                    return "MATCH (d:Disease {name: '糖尿病'})-[:TREATED_BY]->(m:Medicine) RETURN m.name"
                elif "高血压" in query_lower:
                    return "MATCH (d:Disease {name: '高血压'})-[:TREATED_BY]->(m:Medicine) RETURN m.name"
                else:
                    return "MATCH (d:Disease)-[:TREATED_BY]->(m:Medicine) RETURN d.name, m.name LIMIT 10"
            else:
                return "MATCH (d:Disease) RETURN d.id, d.name, d.description LIMIT 10"
        elif "患者" in query_lower:
            if "高血压" in query_lower:
                return "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '高血压'}) RETURN p.name, p.age, p.gender"
            elif "糖尿病" in query_lower:
                return "MATCH (p:Patient)-[:HAS_DISEASE]->(d:Disease {name: '糖尿病'}) RETURN p.name, p.age, p.gender"
            else:
                return "MATCH (p:Patient) RETURN p.id, p.name, p.age, p.gender LIMIT 10"
        elif "症状" in query_lower:
            return "MATCH (s:Symptom) RETURN s.id, s.name, s.severity LIMIT 10"
        elif "药物" in query_lower or "药" in query_lower:
            return "MATCH (m:Medicine) RETURN m.id, m.name, m.type LIMIT 10"
        elif "科室" in query_lower or "部门" in query_lower:
            return "MATCH (de:Department) RETURN de.id, de.name, de.location LIMIT 10"

        # 默认查询
        return "MATCH (n) RETURN n LIMIT 10"

    def format_results(self, results: List, max_display: int = 5) -> str:
        """格式化查询结果"""
        if not results:
            return "无结果"

        output_lines = []
        for i, row in enumerate(results[:max_display]):
            if isinstance(row, list):
                # 简化显示列表
                row_str = []
                for item in row:
                    if isinstance(item, dict):
                        # 显示字典的前几个键值对
                        items = list(item.items())[:3]
                        dict_str = "{" + ", ".join([f"{k}:{v}" for k, v in items]) + "}"
                        if len(item) > 3:
                            dict_str += "..."
                        row_str.append(dict_str)
                    else:
                        row_str.append(str(item))
                output_lines.append(f"{i + 1}. {' | '.join(row_str)}")
            else:
                output_lines.append(f"{i + 1}. {row}")

        if len(results) > max_display:
            output_lines.append(f"... 还有 {len(results) - max_display} 条记录")

        return "\n".join(output_lines)


def demo_queries(manager: TuGraphManager, cypher_generator: SmartCypherGenerator):
    """演示查询"""
    print("\n" + "=" * 60)
    print("演示查询")
    print("=" * 60)

    demo_queries_list = [
        "查找所有节点",
        "统计节点数量",
        "查询疾病信息",
        "查询高血压症状",
        "查询糖尿病药物",
        "查询患者信息",
        "查询科室信息"
    ]

    for query in demo_queries_list:
        print(f"\n查询: {query}")
        cypher = cypher_generator.generate_cypher(query)
        print(f"生成的Cypher: {cypher}")

        start_time = time.time()
        success, results, error = manager.execute_cypher(cypher)
        elapsed_time = time.time() - start_time

        if success:
            if results:
                print(f"✓ 找到 {len(results)} 条结果 (耗时: {elapsed_time:.2f}秒)")
                print(cypher_generator.format_results(results))
            else:
                print(f"✓ 查询成功,但无结果 (耗时: {elapsed_time:.2f}秒)")
        else:
            print(f"❌ 失败: {error}")


def interactive_query_system(manager: TuGraphManager, cypher_generator: SmartCypherGenerator):
    """交互式查询系统"""
    print("\n" + "=" * 60)
    print("TuGraph自然语言查询系统")
    print("=" * 60)
    print("\n支持查询类型:")
    print("1. 数据查询: 查找所有节点, 统计节点数量")
    print("2. 疾病查询: 查询疾病信息, 查询疾病症状")
    print("3. 患者查询: 查询患者信息, 查询特定患者")
    print("4. 药物查询: 查询药物信息, 查询治疗药物")
    print("5. 科室查询: 查询科室信息, 查询就诊科室")
    print("\n命令:")
    print("  help - 查看帮助")
    print("  demo - 运行演示查询")
    print("  clear - 清屏")
    print("  quit - 退出")
    print("=" * 60)

    while True:
        try:
            print("\n" + "-" * 40)
            user_input = input("请输入查询或命令: ").strip()

            if not user_input:
                continue

            if user_input.lower() in ['quit', 'exit', 'q']:
                print("退出系统")
                break

            if user_input.lower() in ['help', '?']:
                print("""
查询示例:
1. 查找所有节点
2. 统计节点数量
3. 查询疾病信息
4. 查询高血压症状
5. 查询糖尿病药物
6. 查询患者信息
7. 查询科室信息
8. 查询高血压患者
9. 查询肺炎症状
10. 查询所有关系
                """)
                continue

            if user_input.lower() == 'demo':
                demo_queries(manager, cypher_generator)
                continue

            if user_input.lower() == 'clear':
                os.system('cls' if os.name == 'nt' else 'clear')
                continue

            print(f"\n正在处理: {user_input}")

            # 生成Cypher
            cypher = cypher_generator.generate_cypher(user_input)
            print(f"生成的Cypher: {cypher}")

            # 执行查询
            start_time = time.time()
            success, results, error = manager.execute_cypher(cypher)
            elapsed_time = time.time() - start_time

            if success:
                if results:
                    print(f"✓ 查询成功! 找到 {len(results)} 条结果 (耗时: {elapsed_time:.2f}秒)")
                    print("\n查询结果:")
                    print(cypher_generator.format_results(results, max_display=8))

                    # 询问是否保存结果
                    save = input("\n是否保存结果到文件? (y/n): ").lower()
                    if save == 'y':
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        filename = f"query_result_{timestamp}.json"
                        data = {
                            "query": user_input,
                            "cypher": cypher,
                            "execution_time": elapsed_time,
                            "result_count": len(results),
                            "results": results
                        }
                        try:
                            with open(filename, 'w', encoding='utf-8') as f:
                                json.dump(data, f, ensure_ascii=False, indent=2)
                            print(f"✓ 结果已保存到: {filename}")
                        except Exception as e:
                            print(f"❌ 保存失败: {e}")
                else:
                    print(f"✓ 查询成功! 但无匹配结果 (耗时: {elapsed_time:.2f}秒)")
            else:
                print(f"❌ 查询失败: {error}")

        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        except Exception as e:
            print(f"❌ 发生错误: {e}")


def main():
    """主函数"""
    print("=" * 60)
    print("TuGraph图数据库与大模型接口集成系统 ")
    print("=" * 60)

    # 配置参数
    TUGraph_HOST = "http://localhost:7070"
    TUGraph_USER = "admin"
    TUGraph_PASSWORD = "73@TuGraph"
    TUGraph_GRAPH = "default"

    print("TuGraph配置:")
    print(f"  主机: {TUGraph_HOST}")
    print(f"  用户名: {TUGraph_USER}")
    print(f"  图名称: {TUGraph_GRAPH}")

    # 初始化TuGraph管理器
    print("\n正在连接到TuGraph...")
    manager = TuGraphManager(
        host=TUGraph_HOST,
        username=TUGraph_USER,
        password=TUGraph_PASSWORD,
        graph_name=TUGraph_GRAPH
    )

    if not manager.connect():
        print("❌ 无法连接到TuGraph数据库")
        print("\n请确保:")
        print("1. TuGraph服务正在运行")
        print("2. 用户名和密码正确")
        print("3. 网络连接正常")
        return

    # 询问是否创建演示数据
    print("\n" + "=" * 60)
    print("数据库初始化")
    print("=" * 60)

    create_data = input("是否创建医疗知识图谱演示数据? (y/n): ").lower()
    if create_data == 'y':
        if create_simple_medical_graph(manager):
            print("✓ 医疗知识图谱创建成功")
        else:
            print("❌ 医疗知识图谱创建失败")
            print("继续使用现有数据...")

    # 初始化Cypher生成器
    cypher_generator = SmartCypherGenerator()

    # 运行演示查询
    demo_queries(manager, cypher_generator)

    # 运行交互式查询系统
    interactive_query_system(manager, cypher_generator)

    print("\n" + "=" * 60)
    print("系统运行完成")
    print("=" * 60)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ 程序异常终止: {e}")
        import traceback

        traceback.print_exc()

更多推荐