import chainlit as cl
from qdrant_client import QdrantClient
from rag_graph import RagGraph

@cl.cache
def get_qdrant_client():
  """Create a QdrantClient instance and cache it for restarts during development."""
  from qdrant_client import QdrantClient
  return QdrantClient(path='data/vectors')

@cl.on_chat_start
async def on_chat_start():
  """Create the RAG graph and store it in the user session."""
  qdrant_client = get_qdrant_client()
  rag_graph = RagGraph(qdrant_client)

  cl.user_session.set("rag_graph", rag_graph)

@cl.on_message
async def on_message(question: cl.Message):
  """Stream the response to the user."""
  msg = cl.Message(content="")
  # Send a message to the user to indicate that the response is being generated
  await msg.send()
  
  rag_graph = cl.user_session.get("rag_graph")
   # Update the message when streaming is complete
  await rag_graph.stream(question.content, msg)