Serving RAG as an API

import json
import re
from collections.abc import AsyncIterable
from contextlib import asynccontextmanager
from typing import Annotated

from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from fastapi.sse import EventSourceResponse, ServerSentEvent
from langchain_core.messages import AIMessage, ToolMessage
from opentelemetry import trace
from pydantic import BaseModel, ConfigDict, Field
from pydantic.alias_generators import to_camel

from docsassistant.agent import open_agent, thread_id
from docsassistant.conversations import conversations          # ownership table: create(), owns()
from docsassistant.security import Caller, caller_from_token


@asynccontextmanager
async def lifespan(app: FastAPI):
    async with open_agent() as agent:                          # async checkpointer and store on one pool
        app.state.agent = agent
        yield


app = FastAPI(title="DocsAssistant API", version="1.4.0", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["https://albertoirurueta.github.io"],
                   allow_methods=["POST", "GET"], allow_headers=["Authorization", "Content-Type", "Idempotency-Key"])
bearer = HTTPBearer(auto_error=False)


class Camel(BaseModel):
    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class Filters(Camel):
    component: str | None = None
    version: str | None = None


class AskRequest(Camel):
    question: str = Field(min_length=1, max_length=4000)
    conversation_id: str | None = None
    filters: Filters | None = None        # narrowing only; merged with the ACL filter as on Access Control & Privacy


def current_caller(creds: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer)]) -> Caller:
    if creds is None:                                         # no anonymous access: bearerAuth is required
        raise HTTPException(status_code=401, detail="missing token", headers={"WWW-Authenticate": "Bearer"})
    try:
        return caller_from_token(creds.credentials)
    except Exception:
        raise HTTPException(status_code=401, detail="invalid token", headers={"WWW-Authenticate": "Bearer"})


def resolve_conversation(caller: Caller, conversation_id: str | None) -> str:
    if conversation_id is None:
        return conversations.create(caller.user_id)
    if not conversations.owns(caller.user_id, conversation_id):
        raise HTTPException(status_code=404, detail="conversation-not-found")    # same answer for foreign IDs
    return conversation_id


def owned_conversation(req: AskRequest, caller: Annotated[Caller, Depends(current_caller)]) -> str:
    """A dependency, so that a 401 or 404 is returned before the stream starts."""
    return resolve_conversation(caller, req.conversation_id)


@app.post("/v1/ask", response_class=EventSourceResponse)
async def ask_stream(req: AskRequest, request: Request, caller: Annotated[Caller, Depends(current_caller)],
                     conversation_id: Annotated[str, Depends(owned_conversation)]) -> AsyncIterable[ServerSentEvent]:
    agent = request.app.state.agent
    trace_id = format(trace.get_current_span().get_span_context().trace_id, "032x")
    config = {"configurable": {"thread_id": thread_id(caller.user_id, conversation_id)}}
    passages: dict[str, dict] = {}
    answer: list[str] = []
    yield ServerSentEvent(event="status", data={"stage": "retrieving"})
    try:
        async for part in agent.astream({"messages": [{"role": "user", "content": req.question}]}, config,
                                        context=caller, stream_mode=["messages", "updates"], version="v2"):
            if part["type"] == "messages":
                chunk, _ = part["data"]
                if isinstance(chunk, AIMessage) and chunk.text:
                    answer.append(chunk.text)
                    yield ServerSentEvent(event="token", data={"text": chunk.text})
            elif part["type"] == "updates" and "tools" in part["data"]:
                for msg in part["data"]["tools"]["messages"]:
                    if isinstance(msg, ToolMessage):
                        for p in json.loads(msg.text):
                            passages[p["chunk_id"]] = p
                yield ServerSentEvent(event="status", data={"stage": "writing", "passages": len(passages)})
        cited = [cid for cid in dict.fromkeys(re.findall(r"\[(c\d+)\]", "".join(answer))) if cid in passages]
        yield ServerSentEvent(event="citations", data=[{"marker": i + 1, "chunkId": cid,
                                                        "source": passages[cid]["source"]}
                                                       for i, cid in enumerate(cited)])
        yield ServerSentEvent(event="done", data={"conversationId": conversation_id, "traceId": trace_id})
    except Exception as exc:                                  # the status is already 200: report as an event
        yield ServerSentEvent(event="error", data={"type": type(exc).__name__})

event: status
data: {"stage": "retrieving"}

event: status
data: {"stage": "writing", "passages": 1}

event: token
data: {"text": "ef_search sets the candidate list size [c1]. Not [c9]."}

event: citations
data: [{"marker": 1, "chunkId": "c1", "source": "vector-rag/ann-indexes"}]

event: done
data: {"conversationId": "conv1", "traceId": "00000000000000000000000000000000"}

package com.example.docsassistant.rag.api;

import java.util.Map;
import java.util.UUID;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.security.core.annotation.AuthenticationPrincipal;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("/v1")
public class AskController {

    public record AskRequest(String question, String conversationId) {
    }

    private final ChatClient chatClient;          // built with memory and RAG advisors

    public AskController(ChatClient docsChatClient) {
        this.chatClient = docsChatClient;
    }

    @PostMapping(path = "/ask", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<Object>> ask(@RequestBody AskRequest req, @AuthenticationPrincipal Jwt jwt) {
        String conversationId = req.conversationId() != null ? req.conversationId() : UUID.randomUUID().toString();
        var acl = new FilterExpressionBuilder().in("acl", readableAcls(jwt)).build();

        Flux<ServerSentEvent<Object>> tokens = chatClient.prompt()
                .user(req.question())
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, jwt.getSubject() + ":" + conversationId)
                        .param(VectorStoreDocumentRetriever.FILTER_EXPRESSION, acl))
                .stream()
                .content()
                .map(text -> ServerSentEvent.<Object>builder(Map.of("text", text)).event("token").build());

        return Flux.concat(
                        Flux.just(ServerSentEvent.<Object>builder(Map.of("stage", "retrieving")).event("status").build()),
                        tokens,
                        Flux.just(ServerSentEvent.<Object>builder(Map.of("conversationId", conversationId))
                                .event("done").build()))
                .onErrorResume(e -> Flux.just(ServerSentEvent.<Object>builder(Map.of("type", e.getClass().getSimpleName()))
                        .event("error").build()));
    }

    private static Object[] readableAcls(Jwt jwt) {
        // the policy from Access Control & Privacy
        return com.example.docsassistant.rag.acl.AclAwareAskService.readableAcls(jwt).toArray();
    }
}

package com.example.docsassistant.l4j.api;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import com.example.docsassistant.rag.acl.AclAwareAskService;
import dev.langchain4j.invocation.InvocationParameters;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import org.springframework.http.HttpStatus;
import org.springframework.security.core.annotation.AuthenticationPrincipal;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.server.ResponseStatusException;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

@RestController
public class StreamingAskController {

    public interface StreamingDocsAssistant {
        TokenStream ask(@MemoryId String memoryKey, @UserMessage String question, InvocationParameters caller);
    }

    public record AskRequest(String question, String conversationId) {
    }

    public interface Conversations {                  // the ownership table, like conversations in the Python API
        void register(String userId, String conversationId);

        boolean isOwner(String userId, String conversationId);
    }

    private final StreamingDocsAssistant assistant;   // AiServices with a StreamingChatModel
    private final Conversations conversations;

    public StreamingAskController(StreamingDocsAssistant assistant, Conversations conversations) {
        this.assistant = assistant;
        this.conversations = conversations;
    }

    @PostMapping("/v1/ask/l4j")
    public SseEmitter ask(@RequestBody AskRequest req, @AuthenticationPrincipal Jwt jwt) {
        String conversationId = ownedConversation(jwt.getSubject(), req.conversationId());
        SseEmitter emitter = new SseEmitter(120_000L);
        List<String> readable = AclAwareAskService.readableAcls(jwt);   // the policy from Access Control & Privacy
        assistant.ask(jwt.getSubject() + ":" + conversationId, req.question(),
                        InvocationParameters.from("readableAcls", readable))
                .onRetrieved(contents -> send(emitter, "status", Map.of("stage", "writing", "passages", contents.size())))
                .onPartialResponse(text -> send(emitter, "token", Map.of("text", text)))
                .onCompleteResponse(response -> {
                    send(emitter, "done", Map.of("conversationId", conversationId));
                    emitter.complete();
                })
                .onError(error -> {
                    send(emitter, "error", Map.of("type", error.getClass().getSimpleName()));
                    emitter.complete();
                })
                .start();
        return emitter;
    }

    private String ownedConversation(String userId, String requested) {
        if (requested == null) {
            String created = UUID.randomUUID().toString();
            conversations.register(userId, created);
            return created;
        }
        if (!conversations.isOwner(userId, requested)) {           // same answer for foreign and unknown IDs
            throw new ResponseStatusException(HttpStatus.NOT_FOUND, "conversation-not-found");
        }
        return requested;
    }

    private static void send(SseEmitter emitter, String event, Object data) {
        try {
            emitter.send(SseEmitter.event().name(event).data(data));
        } catch (IOException clientGone) {
            emitter.completeWithError(clientGone);
        }
    }
}