src/core.py

import sys
import os

_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _root not in sys.path:
    sys.path.insert(0, _root)

from langchain_core.messages import HumanMessage, AIMessage
from langgraph.checkpoint.memory import MemorySaver

from src.agent.graph import build_graph

_graph = build_graph(checkpointer=MemorySaver())

# Sessions whose history has already been seeded from Qdrant this process run
_restored_sessions: set[str] = set()


async def restore_from_qdrant(session_id: str) -> None:
    """Pre-seed LangGraph MemorySaver with history stored in Qdrant.
    Safe to call on a brand-new session (returns early if no history).
    """
    if session_id in _restored_sessions:
        return
    _restored_sessions.add(session_id)

    from src.memory.session import get_store
    history = get_store().get_messages(session_id)
    if not history:
        return

    lc_messages: list = []
    for m in history:
        role, content = m.get("role"), m.get("content", "")
        if role == "user":
            lc_messages.append(HumanMessage(content=content))
        elif role == "assistant":
            lc_messages.append(AIMessage(content=content))

    if lc_messages:
        config = {"configurable": {"thread_id": session_id}}
        await _graph.aupdate_state(config, {"messages": lc_messages, "session_id": session_id})


async def final_answer(query: str, session_id: str) -> str:
    # LangSmith tracing is enabled automatically via env vars:
    # LANGCHAIN_TRACING_V2=true, LANGCHAIN_API_KEY, LANGCHAIN_PROJECT
    config = {
        "configurable": {"thread_id": session_id},
        "metadata": {"session_id": session_id},
    }
    result = await _graph.ainvoke(
        {"messages": [HumanMessage(content=query)], "session_id": session_id},
        config=config,
    )
    last = result["messages"][-1]
    return getattr(last, "content", str(last))


if __name__ == "__main__":
    import asyncio
    import uuid

    async def _main():
        sid   = str(uuid.uuid4())
        query = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else input("Query: ")
        print(await final_answer(query, sid))

    def _suppress_ssl_noise(loop, ctx):
        msg = ctx.get("message", "")
        if "SSL" in msg or "transport" in msg.lower():
            return
        loop.default_exception_handler(ctx)

    loop = asyncio.new_event_loop()
    loop.set_exception_handler(_suppress_ssl_noise)
    try:
        loop.run_until_complete(_main())
    finally:
        loop.close()