rag/pipeline.py

from dataclasses import dataclass, field

from llama_index.core.agent.workflow import (
    AgentInput,
    AgentOutput,
    AgentStream,
    FunctionAgent,
    ToolCall,
    ToolCallResult,
)
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.schema import NodeWithScore
from llama_index.core.tools import FunctionTool

from rag.config import MAX_AGENT_ITERATIONS, MAX_SEARCHES_PER_TURN
from rag.prompts import SEARCH_TOOL_DESCRIPTION, SYSTEM_PROMPT, render_context
from rag.providers import make_llm

# How many prior exchanges we replay to the model. Enough for it to follow up on
# an earlier question, without the transcript taking up the space the excerpts need.
HISTORY_TURNS = 6


@dataclass
class Search:
    """One pass of the agent loop: what it thought, what it asked, and what came back.

    `nodes` is filled in slightly after the rest, once retrieval returns, so that
    we can show the query on screen while the search is still running. `book` is
    the book the model scoped its search to, or None when it searched everything.
    """

    query: str
    book: str | None = None
    thought: str = ""
    nodes: list[NodeWithScore] = field(default_factory=list)


@dataclass
class TurnState:
    searches: list[Search] = field(default_factory=list)
    pending_thought: str = ""
    answer: str = ""

    @property
    def nodes(self) -> list[NodeWithScore]:
        return [node for search in self.searches for node in search.nodes]


def recent(history: list[ChatMessage] | None) -> list[ChatMessage]:
    """The last `HISTORY_TURNS` exchanges, in the order they happened."""
    return (history or [])[-HISTORY_TURNS * 2 :]


def make_search_tool(get_retriever, state: "TurnState") -> FunctionTool:

    def search_rust_docs(query: str, book: str | None = None) -> str:
        # We record the search before running it so that the query can be shown
        # on screen while the retrieval is still in progress.
        search = Search(query=query, book=book, thought=state.pending_thought)
        state.pending_thought = ""
        state.searches.append(search)

        if len(state.searches) > MAX_SEARCHES_PER_TURN:
            return (
                f"Search limit of {MAX_SEARCHES_PER_TURN} reached for this question. "
                "Answer from the excerpts you already have, and say what is missing "
                "if they do not cover the question."
            )
        search.nodes = get_retriever(book).retrieve(query)
        return render_context(search.nodes)

    return FunctionTool.from_defaults(
        fn=search_rust_docs,
        name="search_rust_docs",
        description=SEARCH_TOOL_DESCRIPTION,
    )


async def get_response_stream(
    question: str,
    get_retriever,
    provider: str,
    api_key: str,
    model: str,
    history: list[ChatMessage] | None = None,
):
    """Yield the `TurnState` as it fills, while the agent searches and then answers.

    The LLM is constructed inside this generator so that a bad key raises where
    the caller is already handling streaming errors. This is needed because the
    providers fail at different points: GoogleGenAI validates the key in its
    __init__, while OpenAI and Anthropic only fail once streaming has started.
    """
    state = TurnState()
    llm = make_llm(provider, api_key, model)
    agent = FunctionAgent(
        llm=llm,
        tools=[make_search_tool(get_retriever, state)],
        system_prompt=SYSTEM_PROMPT,
    )

    handler = agent.run(
        user_msg=question,
        chat_history=recent(history),
        max_iterations=MAX_AGENT_ITERATIONS,
        early_stopping_method="generate",
    )
    buffer = ""

    async for event in handler.stream_events():
        if isinstance(event, AgentInput):
            # A new LLM turn is starting, so anything buffered belongs to the
            # previous one.
            buffer = ""
        elif isinstance(event, AgentStream):
            buffer += event.delta
            state.answer = buffer
            yield state
        elif isinstance(event, AgentOutput):
            if event.tool_calls:
                # This turn decided to search, so whatever it wrote was reasoning
                # rather than an answer. We hand it to the tool, which attaches it
                # to the search that follows it.
                state.pending_thought = buffer.strip()
                state.answer = ""
                buffer = ""
            yield state
        elif isinstance(event, (ToolCall, ToolCallResult)):
            yield state

    await handler
    yield state