knowledge/agent/cdss/capbility.py

import time

from agent.cdss.models.schemas import CDSSDict, CDSSInput,CDSSInt,CDSSOutput,CDSSText
from agent.cdss.libs.cdss_helper2 import CDSSHelper
import logging
logger = logging.getLogger(__name__)

class CDSSCapability:
    cdss_helper: CDSSHelper = None
    def __init__(self):
        self.cdss_helper = CDSSHelper()
        logger.debug("CDSSCapability initialized")
    
    def process(self, input: CDSSInput, embeding_search:bool = True) -> CDSSOutput:        
        start_nodes = []
        chief_complaint = input.get_value("chief_complaint")
        logger.info(f"process input: {input}")
        output = CDSSOutput()
        if chief_complaint:
            start_time = time.time()
            for keyword in chief_complaint:
                results = self.cdss_helper.node_search(
                    keyword, limit=10, node_type="症状"
                )
                for item in results:
                    if item['score']>1.9:
                        start_nodes.append(item['id'])
            end_time = time.time()
            print(f"node_search执行完成，耗时：{end_time - start_time:.2f}秒")
            logger.info(f"cdss start from {start_nodes}")    
            result = self.cdss_helper.cdss_travel(input, start_nodes,max_hops=2)

            # for item in result["details"]:
            #     name, data = item
            #     output.departments.value[name] = data

            for item in result["score_diags"]:
                output.diagnosis.value[item[0]] = item[1]


            # for item in result["checks"][:5]:
            #     item[1]['score'] = item[1]['count'] / result["total_checks"]
            #     output.checks.value[item[0]] = item[1]
            # for item in result["drugs"][:5]:
            #     item[1]['score'] = item[1]['count'] / result["total_checks"]
            #     output.drugs.value[item[0]] = item[1]
                #print(f"\t药品：{item[0]} {item[1]['count'] / result["total_drugs"] * 100:.2f} %")
        return output