Conversational RAG & Memory

import os
from contextlib import asynccontextmanager

from langchain.agents import create_agent
from langchain.agents.middleware import ClearToolUsesEdit, ContextEditingMiddleware, SummarizationMiddleware
from langchain.chat_models import init_chat_model
from langchain.embeddings import init_embeddings
from langchain.tools import ToolRuntime, tool
from langchain_postgres import PGVector
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres.aio import AsyncPostgresStore
from psycopg_pool import AsyncConnectionPool

from docsassistant.security import Caller          # verified identity, see Access Control & Privacy

embeddings = init_embeddings(os.environ["EMBEDDING_MODEL"])
vector_store = PGVector(embeddings=embeddings, collection_name="docsassistant",
                        connection=os.environ["PGVECTOR_URL"], use_jsonb=True)


@tool
def search_docs(query: str, runtime: ToolRuntime[Caller]) -> list[dict]:
    """Search the documentation the current user may read. Returns passages with their chunk_id for citation."""
    docs = vector_store.similarity_search(query, k=6, filter={"acl": {"$in": runtime.context.readable_acls}})
    return [{"chunk_id": d.metadata["chunk_id"], "source": d.metadata["source"], "text": d.page_content}
            for d in docs]


def thread_id(user_id: str, conversation_id: str) -> str:
    return f"{user_id}:{conversation_id}"             # the user is part of the key, not a client-supplied value


def build_agent(checkpointer: AsyncPostgresSaver, store: AsyncPostgresStore):
    return create_agent(
        init_chat_model(os.environ["CHAT_MODEL"]),
        tools=[search_docs],
        system_prompt="Answer questions about the Irurueta documentation from search_docs results only. "
                      "Cite chunk_id values in square brackets.",
        middleware=[SummarizationMiddleware(init_chat_model(os.environ["SUMMARY_MODEL"]),
                                            trigger=("tokens", 6000), keep=("messages", 16)),
                    ContextEditingMiddleware(edits=[ClearToolUsesEdit(trigger=8000, keep=2)])],
        context_schema=Caller,
        checkpointer=checkpointer,
        store=store,
    )


@asynccontextmanager
async def open_agent(migrate: bool = False):
    """Open the async pool and yield the agent; the API enters this from its FastAPI lifespan."""
    async with AsyncConnectionPool(os.environ["CONVERSATIONS_DB_URL"], max_size=20, open=False,
                                   kwargs={"autocommit": True, "prepare_threshold": 0}) as pool:
        checkpointer = AsyncPostgresSaver(pool)
        store = AsyncPostgresStore(pool, index={"dims": 768, "embed": embeddings, "fields": ["text"]})
        if migrate:                                   # True only in the deployment job
            await checkpointer.setup()
            await store.setup()
        yield build_agent(checkpointer, store)


async def ask(agent, caller: Caller, conversation_id: str, question: str) -> str:
    result = await agent.ainvoke({"messages": [{"role": "user", "content": question}]},
                                 config={"configurable": {"thread_id": thread_id(caller.user_id, conversation_id)}},
                                 context=caller)
    return result["messages"][-1].text

package com.example.docsassistant.rag.memory;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.security.access.AccessDeniedException;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.stereotype.Service;

@Service
public class ConversationalAskService {

    private final ChatClient chatClient;
    private final ConversationOwnership ownership;

    public ConversationalAskService(ChatClient.Builder builder, ChatMemoryRepository repository,
                                    RetrievalAugmentationAdvisor ragAdvisor, ConversationOwnership ownership) {
        ChatMemory memory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(repository)      // JdbcChatMemoryRepository from the starter
                .maxMessages(20)
                .build();
        this.chatClient = builder
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(memory).build(), ragAdvisor)
                .build();
        this.ownership = ownership;
    }

    public String ask(Jwt caller, String conversationId, String question) {
        if (!ownership.isOwner(caller.getSubject(), conversationId)) {
            throw new AccessDeniedException("Unknown conversation");
        }
        String memoryKey = caller.getSubject() + ":" + conversationId;
        return chatClient.prompt()
                .user(question)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, memoryKey))
                .call()
                .content();
    }

    public interface ConversationOwnership {
        boolean isOwner(String userId, String conversationId);
    }
}

package com.example.docsassistant.l4j.memory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;

import javax.sql.DataSource;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

/**
 * CREATE TABLE chat_memory (memory_id text PRIMARY KEY, messages jsonb NOT NULL,
 *                           updated_at timestamptz NOT NULL DEFAULT now());
 */
public class JdbcChatMemoryStore implements ChatMemoryStore {

    private final DataSource dataSource;

    public JdbcChatMemoryStore(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String sql = "SELECT messages::text FROM chat_memory WHERE memory_id = ?";
        try (Connection c = dataSource.getConnection(); PreparedStatement ps = c.prepareStatement(sql)) {
            ps.setString(1, memoryId.toString());
            try (ResultSet rs = ps.executeQuery()) {
                return rs.next() ? ChatMessageDeserializer.messagesFromJson(rs.getString(1)) : List.of();
            }
        } catch (SQLException e) {
            throw new IllegalStateException(e);
        }
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String sql = """
                INSERT INTO chat_memory (memory_id, messages) VALUES (?, ?::jsonb)
                ON CONFLICT (memory_id) DO UPDATE SET messages = EXCLUDED.messages, updated_at = now()""";
        try (Connection c = dataSource.getConnection(); PreparedStatement ps = c.prepareStatement(sql)) {
            ps.setString(1, memoryId.toString());
            ps.setString(2, ChatMessageSerializer.messagesToJson(messages));
            ps.executeUpdate();
        } catch (SQLException e) {
            throw new IllegalStateException(e);
        }
    }

    @Override
    public void deleteMessages(Object memoryId) {
        try (Connection c = dataSource.getConnection();
             PreparedStatement ps = c.prepareStatement("DELETE FROM chat_memory WHERE memory_id = ?")) {
            ps.setString(1, memoryId.toString());
            ps.executeUpdate();
        } catch (SQLException e) {
            throw new IllegalStateException(e);
        }
    }
}

package com.example.docsassistant.l4j.memory;

import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

public final class PersistentAssistant {

    public interface DocsAssistant {
        String ask(@MemoryId String memoryKey, @UserMessage String question);
    }

    private PersistentAssistant() {
    }

    public static DocsAssistant create(ChatModel model, RetrievalAugmentor augmentor, ChatMemoryStore store) {
        ChatMemoryProvider memories = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(20)
                .chatMemoryStore(store)
                .build();
        return AiServices.builder(DocsAssistant.class)
                .chatModel(model)
                .chatMemoryProvider(memories)
                .retrievalAugmentor(augmentor)
                .build();
    }

    /** The only way callers build a memory key: from the verified subject plus a conversation they own. */
    public static String memoryKey(String verifiedSubject, String conversationId) {
        return verifiedSubject + ":" + conversationId;
    }
}