Turning chat into a runnable function

core_api/src/runnables.py

@@ -15,6 +15,33 @@
 from redbox.models import Chunk
 
 
+def make_chat_runnable(
+    system_prompt: str,
+    llm: ChatLiteLLM,
+) -> Runnable:
+    """Takes a system prompt and LLM returns a chat runnable.
+
+    Runnable takes input of a dict keyed to question and messages.
+
+    Runnable returns a string.
+    """
+    chat_history = [
+        ("system", system_prompt),
+        ("placeholder", "{messages}"),
+        ("user", "{question}"),
+    ]
+
+    return (
+        {
+            "question": itemgetter("question"),
+            "messages": itemgetter("messages"),
+        }
+        | ChatPromptTemplate.from_messages(chat_history)
+        | llm
+        | StrOutputParser()
+    )
+
+
 def make_stuff_document_runnable(
     system_prompt: str,
     llm: ChatLiteLLM,

core_api/tests/test_runnables.py

@@ -2,12 +2,35 @@
 
 from core_api.src.format import format_chunks, get_file_chunked_to_tokens
 from core_api.src.runnables import (
+    make_chat_runnable,
     make_es_retriever,
     make_rag_runnable,
     make_stuff_document_runnable,
 )
 
 
+def test_make_chat_runnable(mock_llm):
+    chain = make_chat_runnable(
+        system_prompt="Your job is chat.",
+        llm=mock_llm,
+    )
+
+    previous_history = [
+        {"text": "Lorem ipsum dolor sit amet.", "role": "user"},
+        {"text": "Consectetur adipiscing elit.", "role": "ai"},
+        {"text": "Donec cursus nunc tortor.", "role": "user"},
+    ]
+
+    response = chain.invoke(
+        input={
+            "question": "How are you today?",
+            "messages": [(msg["role"], msg["text"]) for msg in previous_history],
+        }
+    )
+
+    assert response == "<<TESTING>>"
+
+
 def test_format_chunks(stored_file_chunks):
     formatted_documents = format_chunks(chunks=stored_file_chunks)