class BaseAgent:
    """智能体基类（模拟ModelEngine的BaseAgent）"""
    agent_name = "default_agent"
    agent_desc = "默认智能体，需子类重写业务逻辑"
    def __init__(self, plugins=None):
        self.plugins = plugins or {}  # 关联插件，字典格式管理
    def execute(self, input_data):
        """核心执行方法，必须由子类重写"""
        raise NotImplementedError("请在子类中实现execute方法")

import os
from abc import abstractmethod
from PyPDF2 import PdfReader
from docx import Document
class BasePlugin:
    """插件基类（模拟ModelEngine的BasePlugin）"""
    plugin_name = "default_plugin"
    plugin_version = "1.0.0"
    @abstractmethod
    def run(self, *args, **kwargs):
        """插件核心执行方法，子类重写"""
        pass
# ---------------------- 插件实现（基于本地模型，包含PDF/Word/本地大模型） ----------------------
class PDFPlugin(BasePlugin):
    """PDF解析插件：提取PDF文本内容"""
    plugin_name = "pdf_plugin"
    def run(self, file_path):
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"PDF文件不存在：{file_path}")
        if not file_path.endswith(".pdf"):
            raise ValueError("仅支持后缀为.pdf的文件")
        
        # 读取PDF文本
        reader = PdfReader(file_path)
        text_content = ""
        for page in reader.pages:
            page_text = page.extract_text() or ""
            text_content += page_text + "\n"
        
        return text_content.strip()
class WordPlugin(BasePlugin):
    """Word插件：读取docx文本 + 生成审查报告Word"""
    plugin_name = "word_plugin"
    def run(self, action, file_path, content=None):
        if not os.path.exists(os.path.dirname(file_path)) and os.path.dirname(file_path):
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
        # 动作1：读取docx文件文本
        if action == "read":
            if not file_path.endswith(".docx"):
                raise ValueError("仅支持后缀为.docx的文件")
            doc = Document(file_path)
            text_content = "\n".join([para.text for para in doc.paragraphs])
            return text_content.strip()
        
        # 动作2：生成合规审查报告Word
        elif action == "generate":
            if not content:
                raise ValueError("生成Word报告需要传入content字典参数")
            
            doc = Document()
            # 填充报告内容（企业标准格式）
            doc.add_heading("合同合规审查报告", 0)
            doc.add_paragraph(f"审查合同：{os.path.basename(content['file_path'])}")
            doc.add_paragraph(f"审查时间：{content['review_time']}")
            doc.add_paragraph(f"审查模型：Qwen1.5-1.8B-Chat（本地运行）")
            # 风险点汇总
            doc.add_heading("一、核心风险点汇总", level=1)
            if content["risks"]:
                for i, risk in enumerate(content["risks"], 1):
                    doc.add_paragraph(f"{i}. {risk['item']}", style="List Number")
                    doc.add_paragraph(f"   整改建议：{risk['suggestion']}", style="List Bullet")
            else:
                doc.add_paragraph("✅ 无核心风险点，合同条款符合企业合规要求。")
            # 审查结论
            doc.add_heading("二、审查最终结论", level=1)
            doc.add_paragraph(content["conclusion"])
            # 保存Word文件
            doc.save(file_path)
            return file_path
class LocalLLMPlugin(BasePlugin):
    """本地大模型插件：基于Qwen1.5-1.8B-Chat执行合同合规审查"""
    plugin_name = "local_llm_plugin"
    def __init__(self, model_name="qwen/Qwen1.5-1.8B-Chat", cache_dir="D:\\modelscope\\hub"):
        super().__init__()
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.tokenizer = None
        self.model = None
        self._load_local_model()  # 初始化时自动下载/加载本地模型
    def _load_local_model(self):
        """下载（首次）并加载本地Qwen1.5-1.8B-Chat模型"""
        print("="*50)
        print(f"正在下载/校验模型：{self.model_name}")
        print(f"模型缓存目录：{self.cache_dir}")
        print("提示：首次下载需要约4GB磁盘空间，后续运行直接加载缓存，无需重复下载")
        print("="*50)
        # 步骤1：通过modelscope下载模型到本地缓存
        try:
            local_model_path = snapshot_download(
                self.model_name,
                cache_dir=self.cache_dir
            )
        except Exception as e:
            raise Exception(f"模型下载失败：{e}")
        # 步骤2：通过transformers加载tokenizer和模型
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(
                local_model_path,
                trust_remote_code=True,
                padding_side="right"  # 优化生成效果
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                local_model_path,
                trust_remote_code=True,
                device_map="cpu"  # 自动适配CPU/GPU（有GPU则自动加载，无GPU则用CPU）
            )
            # 模型推理优化（CPU/GPU通用）
            self.model.eval()
        except Exception as e:
            raise Exception(f"模型加载失败：{e}")
        print("="*50)
        print("模型下载/加载完成，可开始合同审查！")
        print("="*50)

from core.agent import BaseAgent
class ParseAgent(BaseAgent):
    """文档解析智能体：解析PDF/Word合同文本"""
    agent_name = "parse_agent"
    agent_desc = "解析PDF/Word格式合同，提取纯文本内容"
    def execute(self, input_data):
        file_path = input_data["file_path"]
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"合同文件不存在：{file_path}")
        # 根据文件后缀选择对应插件
        if file_path.endswith(".pdf"):
            pdf_plugin = self.plugins.get("pdf_plugin")
            return pdf_plugin.run(file_path)
        elif file_path.endswith(".docx"):
            word_plugin = self.plugins.get("word_plugin")
            return word_plugin.run(action="read", file_path=file_path, content=None)
        else:
            raise ValueError("不支持的文件格式，仅支持.pdf和.docx")

from core.agent import BaseAgent
class ReviewAgent(BaseAgent):
    """合规审查智能体：调用本地模型审查合同风险"""
    agent_name = "review_agent"
    agent_desc = "基于本地大模型，审查合同合规风险并生成修改建议"
    def execute(self, input_data):
        contract_text = input_data
        if not contract_text:
            raise ValueError("合同文本为空，无法进行审查")
        # 调用本地大模型插件执行审查
        llm_plugin = self.plugins.get("local_llm_plugin")
        return llm_plugin.run(contract_text)

from core.agent import BaseAgent
import datetime
class ReportAgent(BaseAgent):
    """报告生成智能体：生成标准化Word审查报告"""
    agent_name = "report_agent"
    agent_desc = "整合审查结果，生成企业标准格式Word审查报告"
    def execute(self, input_data):
        review_result = input_data["review_result"]
        file_path = input_data["file_path"]
        # 构造报告内容
        report_content = {
            "file_path": file_path,
            "review_time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "risks": review_result["risks"],
            "conclusion": review_result["conclusion"]
        }
        # 调用Word插件生成报告
        word_plugin = self.plugins.get("word_plugin")
        output_file_name = f"合同审查报告_{os.path.basename(file_path).replace('.', '_')}.docx"
        output_file_path = os.path.join("contract_reports", output_file_name)
        return word_plugin.run(
            action="generate",
            file_path=output_file_path,
            content=report_content
        )

import os
from core.plugin import PDFPlugin, WordPlugin, LLMPlugin
from agents.parse_agent import ParseAgent
from agents.review_agent import ReviewAgent
from agents.report_agent import ReportAgent
def main(contract_file_path):
    """
    主流程：合同文件 → 解析 → 审查 → 生成报告
    :param contract_file_path: 合同文件路径（.pdf或.docx）
    """
    # 步骤1：初始化插件（包含本地大模型插件）
    plugins = {
        "pdf_plugin": PDFPlugin(),
        "word_plugin": WordPlugin(),
        "local_llm_plugin": LocalLLMPlugin(
            model_name="qwen/Qwen1.5-1.8B-Chat",
            cache_dir="D:\\modelscope\\hub"
        )
    }
    # 步骤2：初始化智能体并关联插件
    parse_agent = ParseAgent(plugins=plugins)
    review_agent = ReviewAgent(plugins=plugins)
    report_agent = ReportAgent(plugins=plugins)
    # 步骤3：执行全流程
    try:
        print("\n【步骤1/3】正在解析合同文件...")
        contract_text = parse_agent.execute({"file_path": contract_file_path})
        print("\n【步骤2/3】正在进行合规审查（本地模型推理中，稍等）...")
        review_result = review_agent.execute(contract_text)
        print(f"===review_result===:{review_result}")
        print("\n【步骤3/3】正在生成Word审查报告...")
        report_file_path = report_agent.execute({
            "review_result": review_result,
            "file_path": contract_file_path
        })
        # 步骤4：输出结果
        print("\n" + "="*60)
        print("合同审查全流程完成！")
        print(f"审查报告保存路径：{os.path.abspath(report_file_path)}")
        print(f"风险点数量：{len(review_result['risks'])}")
        for i, risk in enumerate(review_result['risks'], 1):
            print(f"  {i}. 风险点：{risk['item']}")
            print(f"     建议：{risk['suggestion']}")
        print("="*60)
    except Exception as e:
        print(f"\n流程执行失败：{e}")
# ---------------------- 运行入口（直接执行此文件即可） ----------------------
if __name__ == "__main__":
    # 替换为你的合同文件路径（支持.pdf或.docx）
    TEST_CONTRACT_PATH = "测试合同.pdf"  # 可修改为："测试合同.docx"
    # 运行主流程
    main(TEST_CONTRACT_PATH)