Spike/langgraph

core-api/core_api/build_chains.py

@@ -58,7 +58,7 @@ def make_document_context():
     @chain
     def map_operation(input_dict):
         system_map_prompt = env.ai.map_system_prompt
-        prompt_template = PromptTemplate.from_template(env.ai.map_question_prompt)
+        prompt_template = PromptTemplate.from_template(env.ai.chat_map_question_prompt)
 
         formatted_map_question_prompt = prompt_template.format(question=input_dict["question"])
 
@@ -222,14 +222,14 @@ def make_document_context():
     @chain
     def map_operation(input_dict):
         system_map_prompt = env.ai.map_system_prompt
-        prompt_template = PromptTemplate.from_template(env.ai.map_question_prompt)
+        prompt_template = PromptTemplate.from_template(env.ai.chat_map_question_prompt)
 
         formatted_map_question_prompt = prompt_template.format(question=input_dict["question"])
 
         map_prompt = ChatPromptTemplate.from_messages(
             [
                 ("system", system_map_prompt),
-                ("human", formatted_map_question_prompt + env.ai.map_document_prompt),
+                ("human", formatted_map_question_prompt + env.ai.chat_map_question_prompt),
             ]
         )
 

notebooks/summarise.ipynb

@@ -355,7 +355,7 @@
                 "    @chain\n",
                 "    def map_operation(input_dict):\n",
                 "        system_map_prompt = env.ai.map_system_prompt\n",
-                "        prompt_template = PromptTemplate.from_template(env.ai.map_question_prompt)\n",
+                "        prompt_template = PromptTemplate.from_template(env.ai.chat_map_question_prompt)\n",
                 "\n",
                 "        formatted_map_question_prompt = prompt_template.format(question=input_dict[\"question\"])\n",
                 "\n",

redbox-core/pyproject.toml

@@ -14,6 +14,7 @@ readme = "../README.md"
 python = ">=3.12,<3.13"
 pydantic = "^2.7.1"
 elasticsearch = "^8.14.0"
+langchain-community = "^0.2.6"
 langchain = "^0.2.11"
 langchain_openai = "^0.1.9"
 tiktoken = "^0.7.0"
@@ -22,12 +23,14 @@ pydantic-settings = "^2.3.4"
 langchain-elasticsearch = "^0.2.2"
 pytest-dotenv = "^0.5.2"
 kneed = "^0.8.5"
+langgraph = "^0.1.9"
 
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.2"
 moto = "^5.0.10"
 pytest-cov = "^5.0.0"
+pytest-asyncio = "^0.23.6"
 
 [build-system]
 requires = ["poetry-core"]

redbox-core/redbox/chains/components.py

@@ -0,0 +1,87 @@
+from langchain_elasticsearch import ElasticsearchRetriever
+from langchain_core.embeddings import Embeddings, FakeEmbeddings
+from langchain_openai import AzureChatOpenAI
+from langchain_openai.embeddings import AzureOpenAIEmbeddings, OpenAIEmbeddings
+from langchain_core.utils import convert_to_secret_str
+from langchain_core.runnables import ConfigurableField
+import tiktoken
+
+from redbox.models.settings import Settings
+from redbox.graph.retriever import AllElasticsearchRetriever, ParameterisedElasticsearchRetriever
+
+
+def get_chat_llm(env: Settings):
+    return AzureChatOpenAI(
+        api_key=convert_to_secret_str(env.azure_openai_api_key),
+        azure_endpoint=env.azure_openai_endpoint,
+        model=env.azure_openai_model,
+    )
+
+
+def get_tokeniser() -> tiktoken.Encoding:
+    return tiktoken.get_encoding("cl100k_base")
+
+
+def get_azure_embeddings(env: Settings):
+    return AzureOpenAIEmbeddings(
+        azure_endpoint=env.azure_openai_endpoint,
+        api_version=env.azure_api_version_embeddings,
+        model=env.azure_embedding_model,
+        max_retries=env.embedding_max_retries,
+        retry_min_seconds=env.embedding_retry_min_seconds,
+        retry_max_seconds=env.embedding_retry_max_seconds,
+    )
+
+
+def get_openai_embeddings(env: Settings):
+    return OpenAIEmbeddings(
+        api_key=convert_to_secret_str(env.openai_api_key),
+        base_url=env.embedding_openai_base_url,
+        model=env.embedding_openai_model,
+        chunk_size=env.embedding_max_batch_size,
+    )
+
+
+def get_embeddings(env: Settings) -> Embeddings:
+    if env.embedding_backend == "azure":
+        return get_azure_embeddings(env)
+    elif env.embedding_backend == "openai":
+        return get_openai_embeddings(env)
+    elif env.embedding_backend == "fake":
+        return FakeEmbeddings(size=3072)  # TODO
+    else:
+        raise Exception("No configured embedding model")
+
+
+def get_all_chunks_retriever(env: Settings) -> ElasticsearchRetriever:
+    return AllElasticsearchRetriever(
+        es_client=env.elasticsearch_client(),
+        index_name=f"{env.elastic_root_index}-chunk",
+    )
+
+
+def get_parameterised_retriever(env: Settings, embeddings: Embeddings = None) -> ElasticsearchRetriever:
+    """Creates an Elasticsearch retriever runnable.
+
+    Runnable takes input of a dict keyed to question, file_uuids and user_uuid.
+
+    Runnable returns a list of Chunks.
+    """
+    default_params = {
+        "size": env.ai.rag_k,
+        "num_candidates": env.ai.rag_num_candidates,
+        "match_boost": 1,
+        "knn_boost": 1,
+        "similarity_threshold": 0,
+    }
+    return ParameterisedElasticsearchRetriever(
+        es_client=env.elasticsearch_client(),
+        index_name=f"{env.elastic_root_index}-chunk",
+        params=default_params,
+        embedding_model=embeddings or get_embeddings(env),
+        embedding_field_name=env.embedding_document_field_name,
+    ).configurable_fields(
+        params=ConfigurableField(
+            id="params", name="Retriever parameters", description="A dictionary of parameters to use for the retriever."
+        )
+    )

redbox-core/redbox/chains/graph.py

@@ -0,0 +1,117 @@
+import logging
+import re
+
+from langchain.schema import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import Runnable, RunnableLambda, RunnableParallel, chain
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.vectorstores import VectorStoreRetriever
+from tiktoken import Encoding
+
+from redbox.api.format import format_documents
+from redbox.models import ChatRoute, Settings
+from redbox.models.chain import ChainState
+from redbox.models.errors import QuestionLengthError
+
+log = logging.getLogger()
+re_keyword_pattern = re.compile(r"@(\w+)")
+
+
+def build_get_docs(env: Settings, retriever: VectorStoreRetriever):
+    return RunnableParallel({"documents": retriever})
+
+
+@chain
+def set_route(state: ChainState):
+    """
+    Choose an approach to chatting based on the current state
+    """
+    # Match keyword
+    route_match = re_keyword_pattern.search(state["query"].question)
+    if route_match:
+        selected = route_match.group()[1:]
+    elif len(state["query"].file_uuids) > 0:
+        selected = ChatRoute.chat_with_docs.value
+    else:
+        selected = ChatRoute.chat.value
+    log.info(f"Based on user query [{selected}] selected")
+    return {"route_name": selected}
+
+
+def make_chat_prompt_from_messages_runnable(
+    system_prompt: str,
+    question_prompt: str,
+    input_token_budget: int,
+    tokeniser: Encoding,
+):
+    system_prompt_message = [("system", system_prompt)]
+    prompts_budget = len(tokeniser.encode(system_prompt)) - len(tokeniser.encode(question_prompt))
+    token_budget = input_token_budget - prompts_budget
+
+    @chain
+    def chat_prompt_from_messages(state: ChainState):
+        """
+        Create a ChatPrompTemplate as part of a chain using 'chat_history'.
+        Returns the PromptValue using values in the input_dict
+        """
+        log.debug("Setting chat prompt")
+        chat_history_budget = token_budget - len(tokeniser.encode(state["query"].question))
+
+        if chat_history_budget <= 0:
+            raise QuestionLengthError
+
+        truncated_history: list[dict[str, str]] = []
+        for msg in state["query"].chat_history[::-1]:
+            chat_history_budget -= len(tokeniser.encode(msg["text"]))
+            if chat_history_budget <= 0:
+                break
+            else:
+                truncated_history.insert(0, msg)
+
+        return ChatPromptTemplate.from_messages(
+            system_prompt_message
+            + [(msg["role"], msg["text"]) for msg in truncated_history]
+            + [("user", question_prompt)]
+        ).invoke(state["query"].dict() | state.get("prompt_args", {}))
+
+    return chat_prompt_from_messages
+
+
+@chain
+def set_prompt_args(state: ChainState):
+    log.debug("Setting prompt args")
+    return {
+        "prompt_args": {
+            "formatted_documents": format_documents(state.get("documents") or []),
+        }
+    }
+
+
+def build_llm_chain(
+    llm: BaseChatModel, tokeniser: Encoding, env: Settings, system_prompt: str, question_prompt: str
+) -> Runnable:
+    return RunnableParallel(
+        {
+            "response": make_chat_prompt_from_messages_runnable(
+                system_prompt=system_prompt,
+                question_prompt=question_prompt,
+                input_token_budget=env.ai.context_window_size - env.llm_max_tokens,
+                tokeniser=tokeniser,
+            )
+            | llm.with_config(tags=["response"])
+            | StrOutputParser(),
+        }
+    )
+
+
+def get_no_docs_available(env: Settings):
+    return RunnableLambda(
+        lambda _: {
+            "response": env.response_no_doc_available,
+        }
+    )
+
+
+def empty_node(state: ChainState):
+    log.info(f"Empty Node: {state}")
+    return None