swapped query_pipeline with a chat-enabled inference_pipeline

app.py

@@ -1,4 +1,3 @@
-import os
 from dotenv import load_dotenv
 import pandas as pd
 import streamlit as st
@@ -135,10 +134,12 @@ def display_chat_messages(chat_box, chat_input):
                     st.markdown(message["content"], unsafe_allow_html=True)
 
             st.chat_message("user").markdown(chat_input)
-            st.session_state.messages.append({"role": "user", "content": chat_input})
             with st.chat_message("assistant"):
-                response = st.session_state['document_qa_model'].process_message(chat_input)
+                # process user input and generate response
+                response = st.session_state['document_qa_model'].inference(chat_input, st.session_state.messages)
+
                 st.markdown(response)
+                st.session_state.messages.append({"role": "user", "content": chat_input})
                 st.session_state.messages.append({"role": "assistant", "content": response})
 
 

document_qa_engine.py

@@ -1,4 +1,6 @@
 from typing import List
+
+from haystack.dataclasses import ChatMessage
 from pypdf import PdfReader
 from haystack.utils import Secret
 from haystack import Pipeline, Document, component
@@ -8,9 +10,8 @@ from haystack.components.writers import DocumentWriter
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
-from haystack.components.builders import PromptBuilder
+from haystack.components.builders import DynamicChatPromptBuilder
 from haystack.components.generators.chat import OpenAIChatGenerator, HuggingFaceTGIChatGenerator
-from haystack.components.generators import OpenAIGenerator, HuggingFaceTGIGenerator
 from haystack.document_stores.types import DuplicatePolicy
 
 SENTENCE_RETREIVER_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
@@ -70,34 +71,32 @@ def create_ingestion_pipeline(document_store):
     return pipeline
 
 
-def create_query_pipeline(document_store, model_name, api_key):
-    prompt_builder = PromptBuilder(template=template)
+def create_inference_pipeline(document_store, model_name, api_key):
     if model_name == "local LLM":
-        generator = OpenAIGenerator(model=model_name,
-                                    api_base_url="http://localhost:1234/v1",
-                                    generation_kwargs={"max_tokens": MAX_TOKENS}
-                                    )
+        generator = OpenAIChatGenerator(model=model_name,
+                                        api_base_url="http://localhost:1234/v1",
+                                        generation_kwargs={"max_tokens": MAX_TOKENS}
+                                        )
     elif "gpt" in model_name:
-        generator = OpenAIGenerator(api_key=Secret.from_token(api_key), model=model_name,
-                                    generation_kwargs={"max_tokens": MAX_TOKENS}
-                                    )
+        generator = OpenAIChatGenerator(api_key=Secret.from_token(api_key), model=model_name,
+                                        generation_kwargs={"max_tokens": MAX_TOKENS, "stream": False}
+                                        )
     else:
-        generator = HuggingFaceTGIGenerator(token=Secret.from_token(api_key), model=model_name,
-                                            generation_kwargs={"max_new_tokens": MAX_TOKENS}
-                                            )
-
-    query_pipeline = Pipeline()
-    query_pipeline.add_component("text_embedder",
-                                 SentenceTransformersTextEmbedder(model=SENTENCE_RETREIVER_MODEL))
-    query_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store, top_k=3))
-    query_pipeline.add_component("prompt_builder", prompt_builder)
-    query_pipeline.add_component("generator", generator)
-
-    query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
-    query_pipeline.connect("retriever.documents", "prompt_builder.documents")
-    query_pipeline.connect("prompt_builder", "generator")
+        generator = HuggingFaceTGIChatGenerator(token=Secret.from_token(api_key), model=model_name,
+                                                generation_kwargs={"max_new_tokens": MAX_TOKENS}
+                                                )
+    pipeline = Pipeline()
+    pipeline.add_component("text_embedder",
+                           SentenceTransformersTextEmbedder(model=SENTENCE_RETREIVER_MODEL))
+    pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store, top_k=3))
+    pipeline.add_component("prompt_builder",
+                           DynamicChatPromptBuilder(runtime_variables=["query", "documents"]))
+    pipeline.add_component("llm", generator)
+    pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
+    pipeline.connect("retriever.documents", "prompt_builder.documents")
+    pipeline.connect("prompt_builder.prompt", "llm.messages")
 
-    return query_pipeline
+    return pipeline
 
 
 class DocumentQAEngine:
@@ -109,12 +108,33 @@ class DocumentQAEngine:
         self.model_name = model_name
         document_store = InMemoryDocumentStore()
         self.chunks = []
-        self.query_pipeline = create_query_pipeline(document_store, model_name, api_key)
+        self.inference_pipeline = create_inference_pipeline(document_store, model_name, api_key)
         self.pdf_ingestion_pipeline = create_ingestion_pipeline(document_store)
 
     def ingest_pdf(self, uploaded_file):
         self.pdf_ingestion_pipeline.run({"converter": {"uploaded_file": uploaded_file}})
 
-    def process_message(self, query):
-        response = self.query_pipeline.run({"text_embedder": {"text": query}, "prompt_builder": {"question": query}})
-        return response["generator"]["replies"][0]
+    def inference(self, query, input_messages: List[dict]):
+        system_message = ChatMessage.from_system(
+            "You are a professional HR recruiter that answers questions based on the content of the uploaded CV. in 1 or 2 sentences.")
+        messages = [system_message]
+        for message in input_messages:
+            if message["role"] == "user":
+                messages.append(ChatMessage.from_system(message["content"]))
+            else:
+                messages.append(
+                    ChatMessage.from_user(message["content"]))
+        messages.append(ChatMessage.from_user("""
+        Relevant information from the uploaded CV:
+            {% for doc in documents %}
+                {{ doc.content }}
+            {% endfor %}
+
+            \nQuestion: {{query}}
+            \nAnswer:
+        """))
+        res = self.inference_pipeline.run(data={"text_embedder": {"text": query},
+                                                "prompt_builder": {"prompt_source": messages,
+                                                                   "query": query
+                                                                   }})
+        return res["llm"]["replies"][0].content