hybrid_chat_stream.py

"""

Hybrid Chat Streaming Endpoint

Real-time SSE streaming for scenarios + RAG

"""
from typing import AsyncGenerator
import asyncio
from datetime import datetime

from stream_utils import (
    format_sse, stream_text_slowly,
    EVENT_STATUS, EVENT_TOKEN, EVENT_DONE, EVENT_ERROR, EVENT_METADATA
)


async def hybrid_chat_stream(

    request,

    conversation_service,

    intent_classifier,

    scenario_engine,

    advanced_rag,

    embedding_service,

    qdrant_service,

    hf_token,

    lead_storage

) -> AsyncGenerator[str, None]:
    """

    Stream chat responses in real-time (SSE format)

    

    Yields SSE events:

    - status: "Đang suy nghĩ...", "Đang tìm kiếm..."

    - token: Individual text chunks

    - metadata: Context, session info

    - done: Completion signal

    - error: Error messages

    """
    try:
        # === SESSION MANAGEMENT ===
        session_id = request.session_id
        if not session_id:
            session_id = conversation_service.create_session(
                metadata={"user_agent": "api", "created_via": "stream"},
                user_id=request.user_id
            )
            yield format_sse(EVENT_METADATA, {"session_id": session_id})
        
        # === INTENT CLASSIFICATION ===
        yield format_sse(EVENT_STATUS, "Đang phân tích câu hỏi...")
        
        scenario_state = conversation_service.get_scenario_state(session_id) or {}
        intent = intent_classifier.classify(request.message, scenario_state)
        
        # === ROUTING ===
        if intent.startswith("scenario:"):
            # Scenario flow with simulated streaming
            async for sse_event in handle_scenario_stream(
                intent, request.message, session_id,
                scenario_state, scenario_engine, conversation_service, lead_storage
            ):
                yield sse_event
        
        elif intent == "rag:with_resume":
            # Quick RAG answer + resume scenario
            yield format_sse(EVENT_STATUS, "Đang tra cứu...")
            async for sse_event in handle_rag_stream(
                request, advanced_rag, embedding_service, qdrant_service
            ):
                yield sse_event
            
            # Resume hint
            async for chunk in stream_text_slowly(
                "\n\n---\nVậy nha! Quay lại câu hỏi trước nhé ^^",
                chars_per_chunk=5,
                delay_ms=15
            ):
                yield chunk
        
        else:  # Pure RAG
            yield format_sse(EVENT_STATUS, "Đang tìm kiếm trong tài liệu...")
            async for sse_event in handle_rag_stream(
                request, advanced_rag, embedding_service, qdrant_service
            ):
                yield sse_event
        
        # === SAVE HISTORY ===
        # Note: We'll save the full response after streaming completes
        # This requires buffering on the server side
        
        # === DONE ===
        yield format_sse(EVENT_DONE, {
            "session_id": session_id,
            "timestamp": datetime.utcnow().isoformat()
        })
    
    except Exception as e:
        yield format_sse(EVENT_ERROR, str(e))


async def handle_scenario_stream(

    intent, user_message, session_id,

    scenario_state, scenario_engine, conversation_service, lead_storage

) -> AsyncGenerator[str, None]:
    """

    Handle scenario with simulated typing effect

    """
    # Get scenario response (sync)
    if intent == "scenario:continue":
        result = scenario_engine.next_step(
            scenario_id=scenario_state["active_scenario"],
            current_step=scenario_state["scenario_step"],
            user_input=user_message,
            scenario_data=scenario_state.get("scenario_data", {})
        )
    else:
        scenario_type = intent.split(":", 1)[1]
        result = scenario_engine.start_scenario(scenario_type)
    
    # Update state
    if result.get("end_scenario"):
        conversation_service.clear_scenario(session_id)
    elif result.get("new_state"):
        conversation_service.set_scenario_state(session_id, result["new_state"])
    
    # Execute actions
    if result.get("action") and lead_storage:
        action = result['action']
        scenario_data = result.get('new_state', {}).get('scenario_data', {})
        
        if action == "send_pdf_email":
            lead_storage.save_lead(
                event_name=scenario_data.get('step_1_input', 'Unknown'),
                email=scenario_data.get('step_5_input'),
                interests={"group": scenario_data.get('group_size'), "wants_pdf": True},
                session_id=session_id
            )
        elif action == "save_lead_phone":
            lead_storage.save_lead(
                event_name=scenario_data.get('step_1_input', 'Unknown'),
                email=scenario_data.get('step_5_input'),
                phone=scenario_data.get('step_8_input'),
                interests={"group": scenario_data.get('group_size'), "wants_reminder": True},
                session_id=session_id
            )
    
    # Stream response with typing effect
    response_text = result["message"]
    async for chunk in stream_text_slowly(
        response_text,
        chars_per_chunk=4,  # Faster for scenarios
        delay_ms=15
    ):
        yield chunk
    
    yield format_sse(EVENT_METADATA, {
        "mode": "scenario",
        "scenario_active": not result.get("end_scenario")
    })


async def handle_rag_stream(

    request, advanced_rag, embedding_service, qdrant_service

) -> AsyncGenerator[str, None]:
    """

    Handle RAG with real LLM streaming

    """
    # RAG search (sync part)
    context_used = []
    if request.use_rag:
        query_embedding = embedding_service.encode_text(request.message)
        results = qdrant_service.search(
            query_embedding=query_embedding,
            limit=request.top_k,
            score_threshold=request.score_threshold,
            ef=256
        )
        context_used = results
    
    # Build context
    if context_used:
        context_str = "\n\n".join([
            f"[{i+1}] {r['metadata'].get('text', '')[:500]}"
            for i, r in enumerate(context_used[:3])
        ])
    else:
        context_str = "Không tìm thấy thông tin liên quan."
    
    # Simple response (for now - can integrate with real LLM streaming later)
    if context_used:
        response_text = f"Dựa trên tài liệu, {context_used[0]['metadata'].get('text', '')[:300]}..."
    else:
        response_text = "Xin lỗi, tôi không tìm thấy thông tin về câu hỏi này."
    
    # Simulate streaming (will be replaced with real HF streaming)
    async for chunk in stream_text_slowly(
        response_text,
        chars_per_chunk=3,
        delay_ms=20
    ):
        yield chunk
    
    yield format_sse(EVENT_METADATA, {
        "mode": "rag",
        "context_count": len(context_used)
    })


# TODO: Implement real HF InferenceClient streaming
# This requires updating advanced_rag.py to support stream=True