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

import os
import json
import asyncio
import sys
import uuid
from datetime import datetime
from unittest.mock import MagicMock, AsyncMock

# 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

# Token provided by USER
TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJmcmVzaCI6ZmFsc2UsImlhdCI6MTc3NTEwNzEzOS4wODk5NjksImp0aSI6InRjOU9MbndhMkxFckFPaXF1YVk4aXciLCJ0eXBlIjoiYWNjZXNzIiwic3ViIjoiMTUzZjdlMmUtZjBlYS00NDBjLTk5Y2QtN2E3ODJhZThhNjMxIiwibmJmIjoxNzc1MTA3MTM5LCJjc3JmIjoiMmRiNzI3MWEtYTU5Mi00ODU4LWJjMzAtNDY4NGU2ODI1YWNjIiwiZXhwIjoyMDkwNDY3MTM5LCJub25jZSI6ImQzYWNkMDE5YjNiNzkyNWU3ZTY0MDFiMGQyMzhmNDQxIiwidXNlcm5hbWUiOiJkZXZhZG1pbiJ9.YjmpfEFuOAvdOyp9wpPAaykdKGkHFwbNayyVPNDrw_-13-noQwY76lbns-aGF3DG35il2bbjQqNNBIAUFwcD0mDJqeWxaAQkbaL_SGLHpGU6Ul6DGNpyOE4jHxASI7B8DdiwRJkuI-zsA-4uISC8kIkyskOjv-HGFi7mwsxwScbtsWQ5d6dK9xzNCjbwYCVrskJHjEV6aXIYI1VisgbNFyx-hsvbVKKydka1Olu6iurv25ZopDgU0LwRtppALYMURcZYRm9Nrh6ZHLq5306rqUXsKwBbmzKQ8WMExfICvvk1zRPbCNG6-wJCf0E0Ft1XNuFfhbDIM-2N_zS3QCJQGw"

async def test_chat_flow(message: str, location_tag: str = None, conversation_id: str = None):
    print(f"\n=== Testing Chat Flow / Intention Recognition ===")
    print(f"User Message: {message}")
    print(f"Initial Location Tag: {location_tag}")
    print(f"Conversation ID: {conversation_id}")

    # 1. Setup Mock User
    current_user = UserBase(
        user_id=str(uuid.uuid4()),
        name="Test User",
        role="admin"
    )

    # 2. Setup DB Sessions
    async with async_session() as db_session:
        async with collector_async_session() as collector_db_session:

            # 3. Prepare Conversation ID
            conv_uuid = uuid.UUID(conversation_id) if conversation_id else uuid.uuid4()

            # Ensure Conversation exists in DB to prevent IntegrityError (FK constraint)
            if not conversation_id:
                new_conv = ConversationModel(
                    id=conv_uuid,
                    conversation_id=conv_uuid,
                    user_id=uuid.UUID(current_user.user_id),
                    conversation_title=f"Test: {message[:20]}...",
                    created_by=current_user.name,
                    updated_by=current_user.name
                )
                db_session.add(new_conv)
                await db_session.commit()
                print(f"✅ Created dummy conversation in DB: {conv_uuid}")

            # 4. Create Request
            request_as_rag = ChatRequest(
                message=message,
                conversation_id=conv_uuid,
                message_id=uuid.uuid4(),
                attachments=[],
                equipment_id=None,
                flow_type="llm_only",
                model="icon-qwen",
                context_state=None,
                enable_background_analysis=True
            )

            # Mocking CommonParameters (it's a dict according to database/service.py)
            common = {
                "db_session": db_session,
                "page": 1,
                "items_per_page": 5
            }

            # 5. Invoke stream_create
            print("Starting stream...")
            try:
                # We need to pass the request that stream_create expects.
                # stream_create uses request.message and request.conversation_id
                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
                )

                full_response = ""
                async for chunk in generator:
                    if chunk.startswith("data: "):
                        content = chunk[6:].strip()
                        if content:
                            if content.startswith("{"):
                                # This is likely JSON (intention or metadata)
                                try:
                                    data = json.loads(content)
                                    if "intent" in data:
                                        print(f"🎯 [Intention Detected]: {data['intent']} (Confidence: {data.get('confidence')})")
                                        if data.get('suggested_tool'):
                                            print(f"🛠️ [Tool Suggested]: {data['suggested_tool']}")
                                    else:
                                        print(f"[Metadata]: {content}")
                                except:
                                    print(f"[Chunk]: {content}")
                            else:
                                # Normal text
                                print(content, end="", flush=True)
                                full_response += content
                    elif chunk.startswith("event: tool_result"):
                        print(f"\n🛠️ [Tool Result]: {chunk[18:].strip()}")
                    elif chunk.startswith("event: "):
                        print(f"\n[Event]: {chunk.strip()}")

                print("\n\n=== Final Response ===")
                # print(full_response)

                return str(request_as_rag.conversation_id)

            except Exception as e:
                print(f"\n❌ Chat flow failed: {e}")
                import traceback
                traceback.print_exc()
                return None

async def main():
    # Test case 1: Asking for equipment condition (New Conversation)
    print("\n--- TEST 1: New Query ---")
    conv_id = await test_chat_flow("Bagaimana kondisi equipment CHS STACKER/RECLAIMER A saat ini?")

    if conv_id:
        # Test case 2: Follow-up question (Testing history & context persistence)
        print(f"\n--- TEST 2: Follow-up (Context Persistence) ---")
        await test_chat_flow("Apa penyebab utama kegagalannya?", conversation_id=conv_id)

    # Test case 3: Search for location tag by name
    print("\n--- TEST 3: Search by Name ---")
    await test_chat_flow("Cari location tag untuk boiler")

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