diviner/service/agent/role.py

# -*- coding: utf-8 -*-
from .__load__ import *
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
class Role(object):

    def init(self, site_id, role_id, uid):
        self.info = Demeter.service('common').one('role', id=role_id)
        if self.info:
            self.model = Demeter.service('common').one('lang_model', id=self.info['lang_model_id'])
        self.db = None
        self.piece = None
        #self.memory()
        # 知识库挂载
        data = Demeter.service('data').init(site_id)
        context = data.load('similarity', {'k':5, 'fetch_k':50, 'filter': {'role_id': role_id, 'uid' : uid}})
        #sample = data.load('similarity', {'k':5, 'fetch_k':50, 'filter': {'role_id': role_id, 'uid': 'sample'}})
        print(context)
        self.piece = {"context": context | self.format_docs, "question": RunnablePassthrough()}
        return self

    # 写入记忆
    def write(self, memory):
        pass

    # 挂载工具
    def tool(self, tool):
        pass

    def set(self, prompts):
        chain = ChatPromptTemplate.from_template(prompts)
        if not self.piece:
            self.piece = chain
        else:
            self.piece = self.piece | chain
        return self

    def out(self, query, type = []):
        if self.info:
            #self.info['persona'] = '你是一个精美时尚杂志社的编辑，根据以下上下文来回答这个问题{context}'
            template = """你是一个精美时尚杂志社的编辑，根据以下上下文来回答这个问题:

            {context}

            Question: {question},请用中文输出答案。
            """

            template = """你是一位专业医生。以下是病人的病例内容，请根据医学规范生成详细分析报告。

            病例内容：
            {context}


            请根据上面提供的病例内容生成报告。根据病人的核心关注需求提供解决方案。
            报告要求：
            1. 核心健康问题汇总
            2. 潜在风险与关联性分析
            3. 综合健康建议
            4. 紧急情况预警
            5. 解决方案

            请以word格式输出，我好直接生成word。
            """
            
            self.set(template)
            self.model = Demeter.service(self.model['channel'], 'llm').load(model='deepseek-r1', streaming=True)
            full_report = ""
            chain = (self.piece | self.model | StrOutputParser())
            for chunk in chain.stream(query):
                print(chunk, end="")
                full_report += chunk
            #self.save_docx(full_report)

    def format_docs(self, docs):
        return "\n\n".join([d.page_content for d in docs])

    def save_docx(self, content):
        patient_id = self.info.get('uid', 'unknown')  # 或者 role_id
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"report_{patient_id}_{timestamp}.docx"
        output_dir = "reports"
        os.makedirs(output_dir, exist_ok=True)
        filepath = os.path.join(output_dir, filename)

        doc = Document()
        doc.add_heading('诊断报告', 0)
        doc.add_paragraph(report_text)
        doc.save(filepath)

        print(f"\n\n📝 报告已保存为：{filepath}")

    # 生成角色
    def create(self, site_id, uid, name, persona, lang_model_id, data, tool):
        db = Demeter.db('role')
        db.site_id = site_id
        db.create_uid = create_uid
        db.owner_uid = owner_uid
        db.persona = persona
        db.lang_model_id = lang_model_id
        id = db.insert()
        if len(data) > 0:
            for key, value in enumerate(data):
                pass