Agent RAG complet avec LangGraph et FAISS

docs = ['Python est un langage', 'Django est un framework', 'FastAPI est rapide']
vecs = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
app = build_full_rag_agent(docs, vecs, threshold=1.0)

# Question avec résultat pertinent (vecteur proche)
r = app.invoke({
    'query': 'Qu\'est-ce que Python?',
    'query_vector': [0.9, 0.1, 0],
    'needs_search': False,
    'documents': [],
    'is_relevant': False,
    'answer': '',
})
assert r['needs_search'] == True, 'La query contient un ?, needs_search doit etre True'
assert 'Python est un langage' in r['answer'], f'La réponse doit mentionner Python, got: {r["answer"]}'
assert r['answer'].startswith('Réponse: '), 'La réponse doit commencer par Réponse:'

# Pas une question
r2 = app.invoke({
    'query': 'Bonjour',
    'query_vector': [0.5, 0.5, 0],
    'needs_search': False,
    'documents': [],
    'is_relevant': False,
    'answer': '',
})
assert r2['answer'] == "Ce n'est pas une question.", f'Got: {r2["answer"]}'

# Question mais seuil tres bas (pas pertinent)
app2 = build_full_rag_agent(docs, vecs, threshold=0.001)
r3 = app2.invoke({
    'query': 'C\'est quoi?',
    'query_vector': [0.5, 0.5, 0.5],
    'needs_search': False,
    'documents': [],
    'is_relevant': False,
    'answer': '',
})
assert 'pas trouve' in r3['answer'], f'Doit repondre pas trouve, got: {r3["answer"]}'

Solution officielle

from langgraph.graph import StateGraph, START, END
from typing import TypedDict
import faiss
import numpy as np

class State(TypedDict):
    query: str
    query_vector: list
    needs_search: bool
    documents: list
    is_relevant: bool
    answer: str

def build_full_rag_agent(documents: list[str], vectors: list, threshold: float = 1.0):
    dimension = len(vectors[0])
    index = faiss.IndexFlatL2(dimension)
    index.add(np.array(vectors, dtype='float32'))

    def analyser(state):
        return {"needs_search": "?" in state["query"]}

    def chercher(state):
        distances, indices = index.search(
            np.array([state["query_vector"]], dtype='float32'), k=2
        )
        docs = []
        for idx, dist in zip(indices[0], distances[0]):
            docs.append({"text": documents[idx], "distance": float(dist)})
        return {"documents": docs}

    def evaluer(state):
        min_dist = min(d["distance"] for d in state["documents"])
        return {"is_relevant": min_dist < threshold}

    def repondre_avec_contexte(state):
        textes = " | ".join(d["text"] for d in state["documents"])
        return {"answer": f"Réponse: {textes}"}

    def repondre_sans_contexte(state):
        return {"answer": "Desole, je n'ai pas trouve de réponse pertinente."}

    def repondre_direct(state):
        return {"answer": "Ce n'est pas une question."}

    def route_search(state):
        return "search" if state["needs_search"] else "direct"

    def route_relevance(state):
        return "relevant" if state["is_relevant"] else "not_relevant"

    graph = StateGraph(State)
    graph.add_node("analyser", analyser)
    graph.add_node("chercher", chercher)
    graph.add_node("evaluer", evaluer)
    graph.add_node("repondre_avec_contexte", repondre_avec_contexte)
    graph.add_node("repondre_sans_contexte", repondre_sans_contexte)
    graph.add_node("repondre_direct", repondre_direct)
    graph.add_edge(START, "analyser")
    graph.add_conditional_edges("analyser", route_search, {
        "search": "chercher",
        "direct": "repondre_direct",
    })
    graph.add_edge("chercher", "evaluer")
    graph.add_conditional_edges("evaluer", route_relevance, {
        "relevant": "repondre_avec_contexte",
        "not_relevant": "repondre_sans_contexte",
    })
    graph.add_edge("repondre_avec_contexte", END)
    graph.add_edge("repondre_sans_contexte", END)
    graph.add_edge("repondre_direct", END)
    return graph.compile()