ExperionCrawler/futurePlan/End-to-End P&ID Graph Pipeline/pid_intelligent_mapper.py

import networkx as nx
import asyncio
import json
from typing import List, Optional, Dict, Any, Tuple
from pydantic import BaseModel, Field
from rapidfuzz import process, fuzz
from openai import AsyncOpenAI

# --- 응답 구조화를 위한 Pydantic 모델 ---
class MappingResult(BaseModel):
    resolved_tag: str = Field(..., description="The final mapped system tag")
    reason: str = Field(..., description="Reason for this mapping based on context")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score from 0 to 1")

class IntelligentMapper:
    def __init__(self, graph: nx.Graph, system_tags: List[str], api_key: str = None):
        self.graph = graph          # Phase 2에서 생성된 NetworkX 그래프
        self.system_tags = system_tags # Experion 시스템의 전체 태그 리스트
        self.client = AsyncOpenAI(api_key=api_key) if api_key else None

    def get_node_context(self, node_id: str) -> str:
        """노드의 주변 위상 정보를 텍스트로 변환"""
        if not self.graph.has_node(node_id):
            return "Node not found in graph"
            
        neighbors = list(self.graph.neighbors(node_id))
        context = []
        for n in neighbors:
            attr = self.graph.nodes[n]
            val = attr.get('value', n)
            typ = attr.get('type', 'Unknown')
            context.append(f"Connected to {val} (Type: {typ})")
        
        return ", ".join(context) if context else "No connected neighbors"

    async def _resolve_generic(self, node_id: str, category_prompt: str) -> MappingResult:
        """공통 매핑 로직 (비동기 + 구조화 응답)"""
        if not self.client:
            return MappingResult(resolved_tag="UNKNOWN", reason="API Key not provided", confidence=0.0)

        # Phase 2에서 'value'에 clean_value가 저장됨
        node_data = self.graph.nodes.get(node_id, {})
        tag_text = node_data.get('value', '')
        
        # 1차 후보 추출 (RapidFuzz)
        candidates = process.extract(tag_text, self.system_tags, scorer=fuzz.WRatio, limit=5)
        context = self.get_node_context(node_id)
        
        prompt = f"""
        {category_prompt}
        P&ID 도면의 태그 '{tag_text}'를 실제 시스템 태그와 매핑해야 합니다.
        위상 맥락: {context}
        후보 리스트: {candidates}
        
        반드시 다음 JSON 형식으로만 응답하세요:
        {{
            "resolved_tag": "태그명 또는 UNKNOWN",
            "reason": "매핑 이유",
            "confidence": 0.0~1.0
        }}
        """
        
        try:
            response = await self.client.chat.completions.create(
                model="gpt-4-turbo",
                messages=[{"role": "user", "content": prompt}],
                response_format={ "type": "json_object" } # JSON 모드 강제
            )
            raw_content = response.choices[0].message.content
            # Pydantic을 통한 유효성 검사
            return MappingResult.model_validate_json(raw_content)
        except Exception as e:
            print(f"Error resolving node {node_id}: {e}")
            return MappingResult(resolved_tag="UNKNOWN", reason=f"Error: {str(e)}", confidence=0.0)

    # --- 전문화된 Worker 함수들 ---
    async def extract_transmitters(self, node_ids: List[str]) -> Dict[str, MappingResult]:
        prompt = "당신은 계측기 전문 엔지니어입니다. 특히 Pressure/Flow/Level Transmitter 매핑에 특화되어 있습니다."
        tasks = [self._resolve_generic(nid, prompt) for nid in node_ids]
        results = await asyncio.gather(*tasks)
        return dict(zip(node_ids, results))

    async def extract_valves(self, node_ids: List[str]) -> Dict[str, MappingResult]:
        prompt = "당신은 밸브 및 액추에이터 전문 엔지니어입니다. 밸브의 개폐 상태 및 제어 태그 매핑에 특화되어 있습니다."
        tasks = [self._resolve_generic(nid, prompt) for nid in node_ids]
        results = await asyncio.gather(*tasks)
        return dict(zip(node_ids, results))

    async def extract_equipment(self, node_ids: List[str]) -> Dict[str, MappingResult]:
        prompt = "당신은 공정 설비 전문 엔지니어입니다. 펌프, 탱크, 열교환기 등의 메인 설비 태그 매핑에 특화되어 있습니다."
        tasks = [self._resolve_generic(nid, prompt) for nid in node_ids]
        results = await asyncio.gather(*tasks)
        return dict(zip(node_ids, results))

def validate_mapping(resolved_tag: str, symbol_type: str, tag_metadata: Dict[str, Any]) -> Tuple[bool, str]:
    """심볼 타입과 실제 태그 메타데이터의 엄격한 일치 여부 검증"""
    if resolved_tag == "UNKNOWN":
        return False, "Tag not resolved"

    # 단순 키워드가 아닌 허용 단위(Unit) 정의
    unit_map = {
        "Pressure Transmitter": ["bar", "psi", "kPa", "Pa"],
        "Flow Meter": ["m3/h", "lpm", "kg/h"],
        "Temperature Sensor": ["°C", "C", "K", "°F"]
    }
    
    actual_unit = tag_metadata.get('unit', '').strip()
    allowed_units = unit_map.get(symbol_type, [])
    
    # 1. 단위 일치 확인 (최우선)
    if actual_unit and actual_unit in allowed_units:
        return True, "Unit Match"
    
    # 2. 단위가 없는 경우 설명(Description) 기반 2차 검증
    actual_desc = tag_metadata.get('description', '').lower()
    expected_keywords = {
        "Pressure Transmitter": ["pressure", "press"],
        "Flow Meter": ["flow", "flowrate"],
        "Temperature Sensor": ["temp", "temperature"]
    }
    
    keywords = expected_keywords.get(symbol_type, [])
    if any(kw in actual_desc for kw in keywords):
        return True, "Description Match (Unit Missing)"
        
    return False, "Mismatch: Symbol type and Tag metadata do not align"