dev-rcfa-memory/be-ircfa/tests/test_rcfa_performance.py

import os
import json
import asyncio
import sys
import uuid
import time
import re
from datetime import datetime

# Set PYTHONPATH to include the current directory
sys.path.append(os.getcwd())

from src.ircfa_interaction.schema import ChatRequest
from src.auth.model import UserBase
from src.ircfa_interaction.stream2 import stream_create
from src.ircfa_interaction.model import ChatHistory as ChatHistoryModel, Conversation as ConversationModel
from src.database.core import async_session, collector_async_session
from sqlalchemy import delete

TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJmcmVzaCI6ZmFsc2UsImlhdCI6MTc3NTEwNzEzOS4wODk5NjksImp0aSI6InRjOU9MbndhMkxFckFPaXF1YVk4aXciLCJ0eXBlIjoiYWNjZXNzIiwic3ViIjoiMTUzZjdlMmUtZjBlYS00NDBjLTk5Y2QtN2E3ODJhZThhNjMxIiwibmJmIjoxNzc1MTA3MTM5LCJjc3JmIjoiMmRiNzI3MWEtYTU5Mi00ODU4LWJjMzAtNDY4NGU2ODI1YWNjIiwiZXhwIjoyMDkwNDY3MTM5LCJub25jZSI6ImQzYWNkMDE5YjNiNzkyNWU3ZTY0MDFiMGQyMzhmNDQxIiwidXNlcm5hbWUiOiJkZXZhZG1pbiJ9.YjmpfEFuOAvdOyp9wpPAaykdKGkHFwbNayyVPNDrw_-13-noQwY76lbns-aGF3DG35il2bbjQqNNBIAUFwcD0mDJqeWxaAQkbaL_SGLHpGU6Ul6DGNpyOE4jHxASI7B8DdiwRJkuI-zsA-4uISC8kIkyskOjv-HGFi7mwsxwScbtsWQ5d6dK9xzNCjbwYCVrskJHjEV6aXIYI1VisgbNFyx-hsvbVKKydka1Olu6iurv25ZopDgU0LwRtppALYMURcZYRm9Nrh6ZHLq5306rqUXsKwBbmzKQ8WMExfICvvk1zRPbCNG6-wJCf0E0Ft1XNuFfhbDIM-2N_zS3QCJQGw"
LOCATION_TAG = "3FW-H050"  # HP Feedwater Heater No. 5

QUESTIONS = [
    "Bagaimana kondisi equipment HP Feedwater No.5 saat ini?",
    "Bagaimana jika kondisi sensor selama 7 hari yang lalu?",
    "Terjadi kenaikan pada sensor Delta Feed Water Flow, kira-kira apa saja kemungkinan failure mode-nya?",
    "Apa penyebab dari failure mode tersebut?",
    "Tindakan apa yang perlu dilakukan untuk memastikan penyebab tersebut?",
    "Apa Langkah mitigasi untuk menyelesaikan permasalahan ini?",
    "Langkah apa saja yang ekonomis dan feasible dari mitigasi diatas?",
    "Jika saya tidak melakukan semuanya, apa yang akan terjadi pada pembangkit saya?"
]

ENGLISH_TECHNICAL_TERMS = [
    "Pitting Corrosion",
    "Stress Corrosion Cracking",
    "Intergranular Corrosion",
    "Corrosion Thinning",
    "Tube Vibration",
    "Inlet Erosion",
    "Steam Impingement",
    "Tube Plugging",
    "Feedwater Heater",
    "Delta Feed Water Flow",
    "Failure Mode",
    "Root Cause",
    "Mitigation",
    "Containment",
    "Corrective",
    "Preventive",
    "Fault Tree Analysis",
    "Overheating",
    "Water Chemistry",
    "Flow-induced vibration"
]

def validate_gfm_tables(text: str) -> dict:
    has_tables = "|" in text and "-" in text
    has_html_tables = any(tag in text.lower() for tag in ["<table", "<tr", "<td", "<th", "<tbody", "<thead"])
    has_br_tag = "<br" in text.lower()
    
    return {
        "has_markdown_tables": has_tables,
        "has_html_tables": has_html_tables,
        "has_br_tags": has_br_tag,
        "is_compliant": (not has_html_tables) and (not has_br_tag)
    }

def extract_english_terms(text: str) -> list:
    found = []
    for term in ENGLISH_TECHNICAL_TERMS:
        pattern = re.compile(rf'\b{re.escape(term)}\b', re.IGNORECASE)
        if pattern.search(text):
            found.append(term)
    return found

async def run_evaluation():
    print("🚀 Initializing RCFA Performance Evaluation...")
    print(f"Target Location Tag: {LOCATION_TAG} (HP Feedwater Heater No.5)")
    
    # 1. Setup Mock User
    current_user = UserBase(
        user_id=str(uuid.uuid4()),
        name="Evaluator Bot",
        role="admin"
    )

    conv_uuid = uuid.uuid4()
    results = []

    async with async_session() as db_session:
        async with collector_async_session() as collector_db_session:
            # Create the initial conversation in database
            new_conv = ConversationModel(
                id=conv_uuid,
                conversation_id=conv_uuid,
                user_id=uuid.UUID(current_user.user_id),
                conversation_title=f"Eval: {QUESTIONS[0][:20]}...",
                created_by=current_user.name,
                updated_by=current_user.name
            )
            db_session.add(new_conv)
            await db_session.commit()
            print(f"✅ Created conversation in database: {conv_uuid}")

            try:
                for idx, message in enumerate(QUESTIONS):
                    turn_num = idx + 1
                    print(f"\n=========================================")
                    print(f"👉 TURN {turn_num}/8: {message}")
                    print(f"=========================================")

                    # Create actual ChatRequest schema
                    request_as_rag = ChatRequest(
                        conversation_id=conv_uuid,
                        message=message,
                        equipment_id=None,
                        attachments=None,
                        context_state=None,
                        enable_background_analysis=True
                    )

                    common = {
                        "db_session": db_session,
                        "page": 1,
                        "items_per_page": 5
                    }

                    # Start timing
                    start_time = time.time()
                    ttft = None
                    full_response = ""
                    phase_transitions = []
                    usage_metadata = {}

                    generator = stream_create(
                        db_session=db_session,
                        collector_db_session=collector_db_session,
                        request=request_as_rag,
                        upload_type=None,
                        location_tag=LOCATION_TAG,
                        current_user=current_user,
                        token=TOKEN,
                        common=common
                    )

                    async for chunk in generator:
                        if chunk.startswith("data: "):
                            content = chunk[6:].strip()
                            if content:
                                if content.startswith("{"):
                                    # JSON metadata
                                    try:
                                        data = json.loads(content)
                                        if "intent" in data:
                                            print(f"🎯 [Intent]: {data['intent']}")
                                    except:
                                        pass
                                else:
                                    # Text content
                                    if ttft is None:
                                        ttft = round(time.time() - start_time, 2)
                                        print(f"⚡ TTFT: {ttft}s")
                                    
                                    print(content, end="", flush=True)
                                    full_response += content

                        elif chunk.startswith("event: phase"):
                            phase = chunk.replace("event: phase\ndata: ", "").strip()
                            print(f"\n🏷️ [Phase Event]: {phase}")
                            phase_transitions.append({
                                "timestamp": datetime.now().isoformat(),
                                "phase": phase
                            })

                        elif chunk.startswith("event: usage_metadata"):
                            meta_str = chunk.replace("event: usage_metadata\ndata: ", "").strip()
                            try:
                                usage_metadata = json.loads(meta_str)
                                print(f"\n📊 [Usage Metadata]: {usage_metadata}")
                            except:
                                pass

                    total_latency = round(time.time() - start_time, 2)
                    print(f"\n🏁 Finished. Latency: {total_latency}s")

                    # Performance validation
                    gfm_validation = validate_gfm_tables(full_response)
                    english_terms = extract_english_terms(full_response)
                    has_export_button = "<export-rcfa-button>" in full_response

                    # Record results for this turn
                    turn_data = {
                        "turn": turn_num,
                        "question": message,
                        "ttft_seconds": ttft,
                        "total_latency_seconds": total_latency,
                        "prompt_tokens": usage_metadata.get("prompt_tokens", 0),
                        "response_tokens": usage_metadata.get("response_tokens", 0),
                        "total_tokens": usage_metadata.get("tokens_used", 0),
                        "last_detected_phase": phase_transitions[-1]["phase"] if phase_transitions else "Unknown",
                        "phase_transitions": phase_transitions,
                        "gfm_compliance": gfm_validation,
                        "english_terms_found": english_terms,
                        "english_terms_count": len(english_terms),
                        "has_export_button": has_export_button,
                        "response_length": len(full_response),
                        "response_text": full_response
                    }
                    results.append(turn_data)

            finally:
                # Cleanup the test records in database to avoid cluttering
                print("\n🧹 Cleaning up test database records...")
                await db_session.execute(delete(ChatHistoryModel).where(ChatHistoryModel.conversation_id == conv_uuid))
                await db_session.execute(delete(ConversationModel).where(ConversationModel.id == conv_uuid))
                await db_session.commit()
                print("✅ Database cleanup complete.")

    # Save results to a file
    report_dir = "docs"
    os.makedirs(report_dir, exist_ok=True)
    report_file = os.path.join(report_dir, "rcfa_performance_report.json")
    with open(report_file, "w") as f:
        json.dump(results, f, indent=2)
    print(f"💾 Saved JSON report to: {report_file}")

if __name__ == "__main__":
    asyncio.run(run_evaluation())