Query Understanding & Routing

package com.example.docsassistant.rag.routing;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class QueryUnderstandingConfig {

    @Bean
    RetrievalAugmentationAdvisor understandingAdvisor(ChatClient.Builder builder,
            @Qualifier("routingDocumentRetriever") DocumentRetriever routingRetriever) {
        // a small, fast model is enough for these steps; clone() keeps the main client's defaults untouched
        ChatClient.Builder small = builder.clone();
        return RetrievalAugmentationAdvisor.builder()
                .queryTransformers(
                        CompressionQueryTransformer.builder().chatClientBuilder(small).build(),
                        RewriteQueryTransformer.builder().chatClientBuilder(small)
                                .targetSearchSystem("technical documentation search").build())
                .documentRetriever(routingRetriever)      // the router defined further down
                .build();
    }
}

import os
from typing import Literal

from langchain.chat_models import init_chat_model
from pydantic import BaseModel, Field

Component = Literal["algebra", "statistics", "sorting", "backend", "database", "ai", "vector-rag"]


class SearchPlan(BaseModel):
    """A search query plus optional filters taken from the user's question."""
    query: str = Field(description="The question rewritten as a short search query, without the filter words")
    component: Component | None = Field(None, description="Only if the user names a section or library")
    version: str | None = Field(None, pattern=r"^\d+\.\d+$", description="Only if the user names a version")


planner = init_chat_model(os.environ["ROUTER_MODEL"], temperature=0).with_structured_output(SearchPlan)


def plan(question: str) -> SearchPlan:
    return planner.invoke(
        "Extract a search plan. Use a filter only when the user explicitly names it.\n\nQuestion: " + question)


def to_filter(p: SearchPlan, readable_acls: list[str]) -> dict:
    """PGVector filter: the ACL clause always comes from identity and is always present."""
    clauses: list[dict] = [{"acl": {"$in": readable_acls}}]
    if p.component:
        clauses.append({"component": {"$eq": p.component}})
    if p.version:
        clauses.append({"version": {"$eq": p.version}})
    return {"$and": clauses}

import os

from langchain.chat_models import init_chat_model
from pydantic import BaseModel, Field


class SubQuestions(BaseModel):
    questions: list[str] = Field(min_length=1, max_length=4,
                                 description="Self-contained sub-questions; a single item if no split is needed")


decomposer = init_chat_model(os.environ["ROUTER_MODEL"], temperature=0).with_structured_output(SubQuestions)


def decompose(question: str) -> list[str]:
    result = decomposer.invoke(
        "Split the question into the smallest set of self-contained sub-questions needed to answer it. "
        "Do not split simple questions.\n\nQuestion: " + question)
    return result.questions

import math
import os
from typing import Literal

from langchain.chat_models import init_chat_model
from langchain.embeddings import init_embeddings
from pydantic import BaseModel

RouteName = Literal["docs", "sql", "graph", "none"]

EXEMPLARS: dict[str, list[str]] = {
    "docs": ["How do I configure the solver?", "What does ef_search control?", "Explain OAuth refresh tokens"],
    "sql": ["How many pages were updated this month?", "List the ten newest pages in the backend section"],
    "graph": ["Which components depend on sorting?", "What concepts connect Kafka and CDC?"],
    "none": ["Write a poem", "What is the weather tomorrow?"],
}


def _cos(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    return dot / (math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b)))


class Route(BaseModel):
    route: RouteName


class Router:
    def __init__(self, embeddings, llm=None, threshold: float = 0.78):
        self.embeddings = embeddings
        self.llm = llm.with_structured_output(Route) if llm else None
        self.threshold = threshold
        self.vectors = {r: embeddings.embed_documents(qs) for r, qs in EXEMPLARS.items()}

    def route(self, question: str) -> RouteName:
        q = self.embeddings.embed_query(question)
        best_route, best = max(((r, max(_cos(q, v) for v in vs)) for r, vs in self.vectors.items()),
                               key=lambda pair: pair[1])
        if best >= self.threshold or self.llm is None:
            return best_route
        return self.llm.invoke(
            "Choose the back end for this question: docs (explanations), sql (counts and lists over the page "
            "catalogue), graph (dependencies and relationships) or none (out of scope).\n\n" + question).route


router = Router(init_embeddings(os.environ["EMBEDDING_MODEL"]),
                init_chat_model(os.environ["ROUTER_MODEL"], temperature=0))

package com.example.docsassistant.rag.routing;

import java.util.List;
import java.util.Map;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

@Component
public class RoutingDocumentRetriever implements DocumentRetriever {

    enum RouteName { DOCS, SQL, GRAPH, NONE }

    record Route(RouteName route) {
    }

    private final ChatClient router;
    private final Map<RouteName, DocumentRetriever> backends;

    public RoutingDocumentRetriever(ChatClient.Builder builder,
                                    @Qualifier("docsRetriever") DocumentRetriever docsRetriever,
                                    @Qualifier("sqlRetriever") DocumentRetriever sqlRetriever,
                                    @Qualifier("graphRetriever") DocumentRetriever graphRetriever) {
        this.router = builder.clone().build();
        this.backends = Map.of(RouteName.DOCS, docsRetriever, RouteName.SQL, sqlRetriever,
                RouteName.GRAPH, graphRetriever);
    }

    @Override
    public List<Document> retrieve(Query query) {
        Route choice = router.prompt()
                .system("Choose the back end: DOCS (explanations), SQL (counts and lists over the page catalogue), "
                        + "GRAPH (dependencies and relationships) or NONE (out of scope).")
                .user(query.text())
                .call()
                .entity(Route.class);
        DocumentRetriever backend = choice == null ? null : backends.get(choice.route());
        return backend == null ? List.of() : backend.retrieve(query);
    }
}

package com.example.docsassistant.l4j.routing;

import java.util.LinkedHashMap;
import java.util.Map;

import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.router.LanguageModelQueryRouter;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;

public final class RoutedAssistant {

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

    private RoutedAssistant() {
    }

    public static DocsAssistant create(ChatModel chatModel, ChatModel smallModel, ContentRetriever docs,
                                       ContentRetriever catalogueSql, ContentRetriever graph) {
        Map<ContentRetriever, String> routes = new LinkedHashMap<>();
        routes.put(docs, "Explanations and how-to answers from the documentation pages");
        routes.put(catalogueSql, "Counts, lists and dates about pages in the documentation catalogue");
        routes.put(graph, "Dependencies and relationships between components and concepts");

        return AiServices.builder(DocsAssistant.class)
                .chatModel(chatModel)
                .chatMemoryProvider(id -> MessageWindowChatMemory.withMaxMessages(20))
                .retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                        .queryTransformer(new CompressingQueryTransformer(smallModel))   // uses the chat memory
                        .queryRouter(LanguageModelQueryRouter.builder()
                                .chatModel(smallModel)
                                .retrieverToDescription(routes)
                                .fallbackStrategy(LanguageModelQueryRouter.FallbackStrategy.ROUTE_TO_ALL)
                                .build())
                        .build())
                .build();
    }
}

import sqlglot
from sqlglot import exp

ALLOWED_SCHEMA = "catalogue"                                       # the schema that holds the views
ALLOWED_TABLES = {"page_catalogue", "component_catalogue"}          # views, not base tables
# connectives and a few harmless functions; anything else -- including every function sqlglot does not know
# (exp.Anonymous), such as pg_sleep, dblink_exec or set_config -- is rejected
ALLOWED_FUNCTIONS = {exp.And, exp.Or, exp.Not, exp.Count, exp.Sum, exp.Avg, exp.Min, exp.Max, exp.Lower, exp.Upper,
                     exp.Coalesce, exp.Cast, exp.Case, exp.If, exp.CurrentDate, exp.DateTrunc, exp.Extract}


class UnsafeSql(ValueError):
    pass


def validate(sql: str, max_rows: int = 200) -> str:
    statements = sqlglot.parse(sql, read="postgres")
    if len(statements) != 1 or not isinstance(statements[0], exp.Select):
        raise UnsafeSql("Only a single SELECT statement is allowed")
    tree = statements[0]
    if any(isinstance(node, (exp.Insert, exp.Update, exp.Delete, exp.Drop, exp.Create, exp.Command))
           for node in tree.walk()):
        raise UnsafeSql("Data-modifying statements are not allowed")
    functions = {f.name if isinstance(f, exp.Anonymous) else f.sql_name()
                 for f in tree.find_all(exp.Func) if type(f) not in ALLOWED_FUNCTIONS}
    if functions:
        raise UnsafeSql(f"Functions not allowed: {sorted(functions)}")
    ctes = {cte.alias_or_name for cte in tree.find_all(exp.CTE)}
    if any(t.catalog or t.db not in ("", ALLOWED_SCHEMA) for t in tree.find_all(exp.Table)):
        raise UnsafeSql(f"Only tables in the {ALLOWED_SCHEMA} schema are allowed")
    tables = {t.name for t in tree.find_all(exp.Table)} - ctes
    if not tables <= ALLOWED_TABLES:
        raise UnsafeSql(f"Tables not allowed: {sorted(tables - ALLOWED_TABLES)}")
    if tree.args.get("limit") is None:
        tree = tree.limit(max_rows)
    return tree.sql(dialect="postgres")


def run_read_only(conn, sql: str) -> list[tuple]:
    """conn is a psycopg connection opened as the read-only text2sql role."""
    with conn.transaction():
        conn.execute("SET TRANSACTION READ ONLY")
        conn.execute("SET LOCAL statement_timeout = '2s'")
        return conn.execute(validate(sql)).fetchall()