Doing self-route and handling RAG response in stream to avoid double loop

Author: jamesrichards4

Diff for: django_app/tests/test_consumers.py

@@ -16,8 +16,8 @@
 from websockets import WebSocketClientProtocol
 from websockets.legacy.client import Connect
 
-from redbox.graph.root import FINAL_RESPONSE_TAG, ROUTE_NAME_TAG, SOURCE_DOCUMENTS_TAG
 from redbox.models.chain import LLMCallMetadata, RequestMetadata
+from redbox.models.graph import FINAL_RESPONSE_TAG, ROUTE_NAME_TAG, SOURCE_DOCUMENTS_TAG
 from redbox_app.redbox_core import error_messages
 from redbox_app.redbox_core.consumers import ChatConsumer
 from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File, User
@@ -484,7 +484,7 @@ def mocked_connect(uploaded_file: File) -> Connect:
             "data": {"chunk": Token(content="Good afternoon, ")},
         },
         {"event": "on_chat_model_stream", "tags": [FINAL_RESPONSE_TAG], "data": {"chunk": Token(content="Mr. Amor.")}},
-        {"event": "on_chain_end", "tags": [ROUTE_NAME_TAG], "data": {"output": "gratitude"}},
+        {"event": "on_chain_end", "tags": [ROUTE_NAME_TAG], "data": {"output": {"route_name": "gratitude"}}},
         {
             "event": "on_retriever_end",
             "tags": [SOURCE_DOCUMENTS_TAG],
@@ -529,7 +529,7 @@ def mocked_connect_with_naughty_citation(uploaded_file: File) -> CannedGraphLLM:
             "tags": [FINAL_RESPONSE_TAG],
             "data": {"chunk": Token(content="Good afternoon, Mr. Amor.")},
         },
-        {"event": "on_chain_end", "tags": [ROUTE_NAME_TAG], "data": {"output": "gratitude"}},
+        {"event": "on_chain_end", "tags": [ROUTE_NAME_TAG], "data": {"output": {"route_name": "gratitude"}}},
         {
             "event": "on_retriever_end",
             "tags": [SOURCE_DOCUMENTS_TAG],

Diff for: redbox-core/redbox/app.py

@@ -1,15 +1,14 @@
 from langchain_core.vectorstores import VectorStoreRetriever
 
+from redbox.models.graph import FINAL_RESPONSE_TAG, ROUTE_NAME_TAG, SOURCE_DOCUMENTS_TAG
 from redbox.graph.root import get_root_graph
 from redbox.models.chain import RedboxState
 from redbox.models.chat import ChatRoute
+from redbox.models.graph import RedboxEventType
 from redbox.models.settings import Settings
 from redbox.chains.components import get_all_chunks_retriever, get_metadata_retriever, get_parameterised_retriever
-from redbox.graph.root import (
+from redbox.models.graph import (
     ROUTABLE_KEYWORDS,
-    ROUTE_NAME_TAG,
-    FINAL_RESPONSE_TAG,
-    SOURCE_DOCUMENTS_TAG,
 )
 
 
@@ -25,7 +24,6 @@ def __init__(
         metadata_retriever: VectorStoreRetriever | None = None,
         env: Settings | None = None,
         debug: bool = False,
-        interrupt_after: list[str] = [],
     ):
         _env = env or Settings()
         _all_chunks_retriever = all_chunks_retriever or get_all_chunks_retriever(_env)
@@ -54,11 +52,13 @@ async def run(
                 content = event["data"]["output"]
                 if isinstance(content, str):
                     await response_tokens_callback(content)
+            elif kind == "on_custom_event" and event["name"] == RedboxEventType.response_tokens.value:
+                await response_tokens_callback(event["data"])
             elif kind == "on_chain_end" and ROUTE_NAME_TAG in tags:
-                await route_name_callback(event["data"]["output"])
+                await route_name_callback(event["data"]["output"]["route_name"])
             elif kind == "on_retriever_end" and SOURCE_DOCUMENTS_TAG in tags:
                 await documents_callback(event["data"]["output"])
-            elif kind == "on_custom_event" and event["name"] == "on_metadata_generation":
+            elif kind == "on_custom_event" and event["name"] == RedboxEventType.on_metadata_generation.value:
                 await metadata_tokens_callback(event["data"])
             elif kind == "on_chain_end" and event["name"] == "LangGraph":
                 final_state = RedboxState(**event["data"]["output"])

Diff for: redbox-core/redbox/chains/runnables.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterator
+from typing import Any, Callable, Iterator, Iterable
 import re
 from operator import itemgetter
 
@@ -9,9 +9,13 @@
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import Runnable, chain, RunnableLambda
+from langchain_core.runnables import Runnable, chain, RunnableLambda, RunnableGenerator
+from langchain_core.callbacks.manager import dispatch_custom_event
+
 from tiktoken import Encoding
 
+from redbox.models.graph import RedboxEventType
+
 from redbox.api.format import format_documents
 from redbox.chains.components import get_tokeniser
 from redbox.models.chain import ChainChatMessage, RedboxState
@@ -67,25 +71,54 @@ def _chat_prompt_from_messages(state: RedboxState) -> Runnable:
     return _chat_prompt_from_messages
 
 
-def build_llm_chain(prompt_set: PromptSet, llm: BaseChatModel, final_response_chain: bool = False) -> Runnable:
+def build_llm_chain(
+    prompt_set: PromptSet,
+    llm: BaseChatModel,
+    output_parser: Runnable | Callable = None,
+    final_response_chain: bool = False,
+) -> Runnable:
     """Builds a chain that correctly forms a text and metadata state update.
 
     Permits both invoke and astream_events.
     """
     model_name = getattr(llm, "model_name", "unknown-model")
     _llm = llm.with_config(tags=["response_flag"]) if final_response_chain else llm
-
+    _output_parser = output_parser if output_parser else StrOutputParser()
     return (
         build_chat_prompt_from_messages_runnable(prompt_set)
         | {
             "prompt": RunnableLambda(lambda prompt: prompt.to_string()),
-            "response": _llm | StrOutputParser(),
+            "response": _llm | _output_parser,
             "model": lambda x: model_name,
         }
         | {"text": itemgetter("response"), "metadata": to_request_metadata}
     )
 
 
+def build_self_route_output_parser(final_response_chain: bool = False):
+    def _self_route_output_parser(chunks: Iterable[AIMessageChunk]) -> Iterable[str]:
+        current_content = ""
+        tokens_to_pass = 4
+        token_count = 0
+        for chunk in chunks:
+            current_content += chunk.content
+            token_count += 1
+            if "unanswerable" in current_content:
+                yield current_content
+                return
+            elif token_count > tokens_to_pass:
+                break
+        if final_response_chain:
+            dispatch_custom_event(RedboxEventType.response_tokens, current_content)
+        yield current_content
+        for chunk in chunks:
+            if final_response_chain:
+                dispatch_custom_event(RedboxEventType.response_tokens, chunk.content)
+            yield chunk.content
+
+    return RunnableGenerator(_self_route_output_parser)
+
+
 class CannedChatLLM(BaseChatModel):
     """A custom chat model that returns its text as if an LLM returned it.
 

Diff for: redbox-core/redbox/graph/nodes/processes.py

@@ -1,16 +1,18 @@
 import logging
+from operator import add
 import re
 import json
-from typing import Any, Callable
+from typing import Any
 from uuid import uuid4
 from functools import reduce
 
 from langchain.schema import StrOutputParser
-from langchain_core.runnables import RunnableParallel, RunnablePassthrough
+from langchain_core.runnables import Runnable, RunnableLambda, RunnableParallel
 from langchain_core.vectorstores import VectorStoreRetriever
 
 from redbox.chains.components import get_tokeniser, get_chat_llm
 from redbox.chains.runnables import build_llm_chain, CannedChatLLM
+from redbox.models.graph import ROUTE_NAME_TAG
 from redbox.models import ChatRoute, Settings
 from redbox.models.chain import RedboxState, RequestMetadata
 from redbox.transform import combine_documents, structure_documents
@@ -29,7 +31,7 @@
 
 def build_retrieve_pattern(
     retriever: VectorStoreRetriever, final_source_chain: bool = False
-) -> Callable[[RedboxState], dict[str, Any]]:
+) -> Runnable[RedboxState, dict[str, Any]]:
     """Returns a function that uses state["request"] and state["text"] to set state["documents"]."""
     retriever_chain = RunnableParallel({"documents": retriever | structure_documents})
 
@@ -38,16 +40,13 @@ def build_retrieve_pattern(
     else:
         _retriever = retriever_chain
 
-    def _retrieve(state: RedboxState) -> dict[str, Any]:
-        return _retriever.invoke(state)
-
-    return _retrieve
+    return _retriever
 
 
 def build_chat_pattern(
     prompt_set: PromptSet,
     final_response_chain: bool = False,
-) -> Callable[[RedboxState], dict[str, Any]]:
+) -> Runnable[RedboxState, dict[str, Any]]:
     """Returns a function that uses state["request"] to set state["text"]."""
 
     def _chat(state: RedboxState) -> dict[str, Any]:
@@ -64,7 +63,7 @@ def _chat(state: RedboxState) -> dict[str, Any]:
 def build_merge_pattern(
     prompt_set: PromptSet,
     final_response_chain: bool = False,
-) -> Callable[[RedboxState], dict[str, Any]]:
+) -> Runnable[RedboxState, dict[str, Any]]:
     """Returns a function that uses state["request"] and state["documents"] to return one item in state["documents"].
 
     When combined with chunk send, will replace each Document with what's returned from the LLM.
@@ -75,6 +74,7 @@ def build_merge_pattern(
     """
     tokeniser = get_tokeniser()
 
+    @RunnableLambda
     def _merge(state: RedboxState) -> dict[str, Any]:
         llm = get_chat_llm(Settings(), state["request"].ai_settings)
 
@@ -117,46 +117,55 @@ def _merge(state: RedboxState) -> dict[str, Any]:
 
 def build_stuff_pattern(
     prompt_set: PromptSet,
+    output_parser: Runnable = None,
     final_response_chain: bool = False,
-) -> Callable[[RedboxState], dict[str, Any]]:
+) -> Runnable[RedboxState, dict[str, Any]]:
     """Returns a function that uses state["request"] and state["documents"] to set state["text"]."""
 
+    @RunnableLambda
     def _stuff(state: RedboxState) -> dict[str, Any]:
         llm = get_chat_llm(Settings(), state["request"].ai_settings)
 
-        return build_llm_chain(prompt_set=prompt_set, llm=llm, final_response_chain=final_response_chain).invoke(state)
+        events = []
+
+        for event in build_llm_chain(
+            prompt_set=prompt_set, llm=llm, output_parser=output_parser, final_response_chain=final_response_chain
+        ).stream(state):
+            events.append(event)
+
+        if len(events) == 0:
+            return None
+        else:
+            return reduce(add, events)
 
     return _stuff
 
 
 ## Utility patterns
 
 
-def build_set_route_pattern(route: ChatRoute) -> Callable[[RedboxState], dict[str, Any]]:
-    """Returns a function that sets state["route_name"]."""
+def build_set_route_pattern(route: ChatRoute) -> Runnable[RedboxState, dict[str, Any]]:
+    """Returns a Runnable that sets state["route_name"]."""
 
     def _set_route(state: RedboxState) -> dict[str, Any]:
-        set_route_chain = (RunnablePassthrough() | StrOutputParser()).with_config(tags=["route_flag"])
-
-        return {"route_name": set_route_chain.invoke(route.value)}
+        return {"route_name": route}
 
-    return _set_route
+    return RunnableLambda(_set_route).with_config(tags=[ROUTE_NAME_TAG])
 
 
+@RunnableLambda
 def set_self_route_from_llm_answer(state: RedboxState):
-    llm_response = state["text"].lower()
-    if llm_response == "true":
-        route = ChatRoute.search
-    elif llm_response == "false":
-        route = ChatRoute.chat_with_docs_map_reduce
+    llm_response = state["text"]
+    if "unanswerable" in llm_response[: min(20, len(llm_response))]:
+        return {"route_name": ChatRoute.chat_with_docs_map_reduce}
     else:
-        route = ChatRoute.search
-    return {"route_name": route.value}
+        return {"route_name": ChatRoute.search}
 
 
-def build_passthrough_pattern() -> Callable[[RedboxState], dict[str, Any]]:
-    """Returns a function that uses state["request"] to set state["text"]."""
+def build_passthrough_pattern() -> Runnable[RedboxState, dict[str, Any]]:
+    """Returns a Runnable that uses state["request"] to set state["text"]."""
 
+    @RunnableLambda
     def _passthrough(state: RedboxState) -> dict[str, Any]:
         return {
             "text": state["request"].question,
@@ -166,10 +175,11 @@ def _passthrough(state: RedboxState) -> dict[str, Any]:
 
 
 def build_set_text_pattern(text: str, final_response_chain: bool = False):
-    """Returns a function that can arbitrarily set state["text"] to a value."""
+    """Returns a Runnable that can arbitrarily set state["text"] to a value."""
     llm = CannedChatLLM(text=text)
     _llm = llm.with_config(tags=["response_flag"]) if final_response_chain else llm
 
+    @RunnableLambda
     def _set_text(state: RedboxState) -> dict[str, Any]:
         set_text_chain = _llm | StrOutputParser()
 
@@ -179,6 +189,7 @@ def _set_text(state: RedboxState) -> dict[str, Any]:
 
 
 def build_set_metadata_pattern():
+    @RunnableLambda
     def _set_metadata_pattern(state: RedboxState):
         flat_docs = flatten_document_state(state.get("documents", {}))
         return {
@@ -192,10 +203,11 @@ def _set_metadata_pattern(state: RedboxState):
 
 
 def build_error_pattern(text: str, route_name: str | None):
+    @RunnableLambda
     def _error_pattern(state: RedboxState):
-        return build_set_text_pattern(text, final_response_chain=True)(state) | build_set_route_pattern(route_name)(
-            state
-        )
+        return build_set_text_pattern(text, final_response_chain=True).invoke(state) | build_set_route_pattern(
+            route_name
+        ).invoke(state)
 
     return _error_pattern