Stateful LangGraph Agents with sovseal Memory Nodes

import os
import asyncio
from typing import Annotated, TypedDict, List
from dotenv import load_dotenv
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from sovseal import AgentStateClient, EngineUnavailableError

load_dotenv()

client = AgentStateClient()  # store/recall need no endpoint/api_key


# 1. Define the shared state structure
class AgentGraphState(TypedDict):
    messages: Annotated[List[HumanMessage | AIMessage], add_messages]
    user_id: str


# 2. Build the chatbot node with memory injection
async def query_chatbot(state: AgentGraphState):
    messages = state["messages"]
    last_query = messages[-1].content

    # Recall-on-entry: fetch the top-3 relevant memories
    try:
        hits = client.recall(last_query, top_k=3)
    except EngineUnavailableError:
        hits = []  # fail open — proceed without recalled context
    context_string = "\n".join(h["text"] for h in hits)

    # Inject recalled context into the system prompt
    system_prompt = (
        "You are a personalized assistant. "
        "Use the following context from previous interactions:\n"
        f"{context_string}"
    )

    llm = ChatOpenAI(model="gpt-4o")
    full_messages = [SystemMessage(content=system_prompt)] + messages
    response = await llm.ainvoke(full_messages)

    # Store-on-exit: persist the current exchange
    try:
        client.store(
            f"User asked: '{last_query}' — "
            f"Assistant responded: '{response.content}'"
        )
    except EngineUnavailableError:
        pass  # fail open — engine not installed/running

    return {"messages": [response]}


# 3. Compile the graph
workflow = StateGraph(AgentGraphState)
workflow.add_node("chatbot", query_chatbot)
workflow.add_edge(START, "chatbot")
workflow.add_edge("chatbot", END)

app = workflow.compile()


# 4. Run the graph
async def run():
    inputs = {
        "messages": [HumanMessage(content="I am building a state continuity system.")],
        "user_id": "alice_dev",
    }
    config = {"configurable": {"thread_id": "alice_dev"}}

    async for event in app.astream(inputs, config):
        for value in event.values():
            print("Chatbot:", value["messages"][-1].content)


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