self-constructing_graph/agent/cdss/capbility.py

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

class CDSSCapability:
    cdss_helper: CDSSHelper = None
    def __init__(self):
        self.cdss_helper = CDSSHelper()
        #self.cdss_helper.load_local_data()
        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:
            for keyword in chief_complaint:
                results = self.cdss_helper.node_search(
                    keyword, limit=10, node_type="word"
                )
                for item in results:
                    if item['score']>1.9:                        
                        start_nodes.append(item['id'])
            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["diags"][:5]:                       
                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