# # websocket_server.py
# import json
# from fastapi import WebSocket
# from starlette.websockets import WebSocketState
# from ..webhook.llm_engine import LLMEngine
# from ..webhook.memory import RedisSession
# from typing import List, Optional


# async def stream_response(ws: WebSocket, query: str,session: RedisSession, kb_ids: Optional[List[str]] = None):
#     """Stream response from RAG agent to client."""
#     reply = ""
#     agent = LLMEngine.get_instance().rag_agent
    
#     history = session.get_recent(top_k=5)  # list of {"role", "content"}
#     memory_queries = [m["content"] for m in history if m["role"] == "user"]
    
#     async for token in agent.stream_answer(query,memory_queries, kb_ids):
#         reply += token
#         if ws.client_state == WebSocketState.CONNECTED:
#             await ws.send_text(json.dumps({"type": "text", "token": token, "last": False}))

#     if ws.client_state == WebSocketState.CONNECTED:
#         await ws.send_text(json.dumps({"type": "text", "token": "", "last": True}))
        
#     session.add_message(role="assistant", content=reply)