[DERCBOT-1037] Use of PromptTemplate + Rewrite the RAG chain using LCEL

Author: assouktim

bot/engine/src/main/kotlin/engine/config/RAGAnswerHandler.kt

@@ -30,10 +30,7 @@ import ai.tock.bot.engine.action.SendSentence
 import ai.tock.bot.engine.action.SendSentenceWithFootnotes
 import ai.tock.bot.engine.dialog.Dialog
 import ai.tock.bot.engine.user.PlayerType
-import ai.tock.genai.orchestratorclient.requests.ChatMessage
-import ai.tock.genai.orchestratorclient.requests.ChatMessageType
-import ai.tock.genai.orchestratorclient.requests.DialogDetails
-import ai.tock.genai.orchestratorclient.requests.RAGQuery
+import ai.tock.genai.orchestratorclient.requests.*
 import ai.tock.genai.orchestratorclient.responses.ObservabilityInfo
 import ai.tock.genai.orchestratorclient.responses.RAGResponse
 import ai.tock.genai.orchestratorclient.responses.TextWithFootnotes
@@ -189,10 +186,14 @@ object RAGAnswerHandler : AbstractProactiveAnswerHandler {
                             )
                         ),
                         questionAnsweringLlmSetting = ragConfiguration.llmSetting,
-                        questionAnsweringPromptInputs = mapOf(
-                            "question" to action.toString(),
-                            "locale" to userPreferences.locale.displayLanguage,
-                            "no_answer" to ragConfiguration.noAnswerSentence
+                        questionAnsweringPrompt = PromptTemplate(
+                            formatter = Formatter.F_STRING.id,
+                            template = ragConfiguration.llmSetting.prompt,
+                            inputs = mapOf(
+                                "question" to action.toString(),
+                                "locale" to userPreferences.locale.displayLanguage,
+                                "no_answer" to ragConfiguration.noAnswerSentence
+                            )
                         ),
                         embeddingQuestionEmSetting = ragConfiguration.emSetting,
                         documentIndexName = indexName,

gen-ai/orchestrator-client/src/main/kotlin/ai/tock/genai/orchestratorclient/requests/RAGQuery.kt

@@ -24,10 +24,10 @@ import ai.tock.genai.orchestratorcore.models.vectorstore.VectorStoreSetting
 
 data class RAGQuery(
     // val condenseQuestionLlmSetting: LLMSetting,
-    // val condenseQuestionPromptInputs: Map<String, String>,
+    // val condenseQuestionPrompt: PromptTemplate,
     val dialog: DialogDetails?,
     val questionAnsweringLlmSetting: LLMSetting,
-    val questionAnsweringPromptInputs: Map<String, String>,
+    val questionAnsweringPrompt: PromptTemplate,
     val embeddingQuestionEmSetting: EMSetting,
     val documentIndexName: String,
     val documentSearchParams: DocumentSearchParamsBase,

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/models/llm/llm_setting.py

@@ -39,8 +39,3 @@ class BaseLLMSetting(BaseModel):
         ge=0,
         le=2,
     )
-    prompt: str = Field(
-        description='The prompt to generate completions for.',
-        examples=['How to learn to ride a bike without wheels!'],
-        min_length=1,
-    )

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/routers/requests/requests.py

@@ -85,6 +85,10 @@ class BaseQuery(BaseModel):
     observability_setting: Optional[ObservabilitySetting] = Field(
         description='The observability settings.', default=None
     )
+    compressor_setting: Optional[DocumentCompressorSetting] = Field(
+        description='Compressor settings, to rerank relevant documents returned by retriever.',
+        default=None,
+    )
 
 
 class QAQuery(BaseQuery):
@@ -159,43 +163,20 @@ class RagQuery(BaseQuery):
     """The RAG query model"""
 
     dialog: Optional[DialogDetails] = Field(description='The user dialog details.')
-    question_answering_prompt_inputs: Any = Field(
-        description='Key-value inputs for the llm prompt when used as a template. Please note that the '
-        'chat_history field must not be specified here, it will be override by the dialog.history field',
-    )
     # condense_question_llm_setting: LLMSetting =
     #   Field(description="LLM setting, used to condense the user's question.")
-    # condense_question_prompt_inputs: Any = (
-    #         Field(
-    #             description='Key-value inputs for the condense question llm prompt, when used as a template.',
-    #         ),
+    # condense_question_prompt: PromptTemplate = Field(
+    #         description='Prompt template, used to create a prompt with inputs for jinja and fstring format'
     #     )
     question_answering_llm_setting: LLMSetting = Field(
         description='LLM setting, used to perform a QA Prompt.'
     )
-    question_answering_prompt_inputs: Any = Field(
-        description='Key-value inputs for the llm prompt when used as a template. Please note that the '
-        'chat_history field must not be specified here, it will be override by the dialog.history field',
-    )
-    embedding_question_em_setting: EMSetting = Field(
-        description="Embedding model setting, used to calculate the user's question vector."
-    )
-    document_index_name: str = Field(
-        description='Index name corresponding to a document collection in the vector database.',
-    )
-    document_search_params: DocumentSearchParams = Field(
-        description='The document search parameters. Ex: number of documents, metadata filter',
-    )
-    observability_setting: Optional[ObservabilitySetting] = Field(
-        description='The observability settings.', default=None
+    question_answering_prompt : PromptTemplate = Field(
+        description='Prompt template, used to create a prompt with inputs for jinja and fstring format'
     )
     guardrail_setting: Optional[GuardrailSetting] = Field(
         description='Guardrail settings, to classify LLM output toxicity.', default=None
     )
-    compressor_setting: Optional[DocumentCompressorSetting] = Field(
-        description='Compressor settings, to rerank relevant documents returned by retriever.',
-        default=None,
-    )
     documents_required: Optional[bool] = Field(
         description='Specifies whether the presence of documents is mandatory for generating answers. '
                     'If set to True, the system will only provide answers when relevant documents are found. '
@@ -223,7 +204,11 @@ class RagQuery(BaseQuery):
                             'secret': 'ab7***************************A1IV4B',
                         },
                         'temperature': 1.2,
-                        'prompt': """Use the following context to answer the question at the end.
+                        'model': 'gpt-3.5-turbo',
+                    },
+                    'question_answering_prompt': {
+                        'formatter': 'f-string',
+                        'template': """Use the following context to answer the question at the end.
 If you don't know the answer, just say {no_answer}.
 
 Context:
@@ -233,12 +218,11 @@ class RagQuery(BaseQuery):
 {question}
 
 Answer in {locale}:""",
-                        'model': 'gpt-3.5-turbo',
-                    },
-                    'question_answering_prompt_inputs': {
-                        'question': 'How to get started playing guitar ?',
-                        'no_answer': "Sorry, I don't know.",
-                        'locale': 'French',
+                        'inputs': {
+                            'question': 'How to get started playing guitar ?',
+                            'no_answer': 'Sorry, I don t know.',
+                            'locale': 'French',
+                        }
                     },
                     'embedding_question_em_setting': {
                         'provider': 'OpenAI',

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/services/completion/completion_service.py

@@ -16,23 +16,14 @@
 
 import logging
 import time
-from typing import Optional
 
-from jinja2 import Template, TemplateError
 from langchain_core.output_parsers import NumberedListOutputParser
 from langchain_core.prompts import PromptTemplate as LangChainPromptTemplate
-from langchain_core.runnables import RunnableConfig
 
-from gen_ai_orchestrator.errors.exceptions.exceptions import (
-    GenAIPromptTemplateException,
-)
 from gen_ai_orchestrator.errors.handlers.openai.openai_exception_handler import (
     openai_exception_handler,
 )
-from gen_ai_orchestrator.models.errors.errors_models import ErrorInfo
 from gen_ai_orchestrator.models.observability.observability_trace import ObservabilityTrace
-from gen_ai_orchestrator.models.prompt.prompt_formatter import PromptFormatter
-from gen_ai_orchestrator.models.prompt.prompt_template import PromptTemplate
 from gen_ai_orchestrator.routers.requests.requests import (
     SentenceGenerationQuery,
 )
@@ -42,6 +33,7 @@
 from gen_ai_orchestrator.services.langchain.factories.langchain_factory import (
     get_llm_factory, create_observability_callback_handler,
 )
+from gen_ai_orchestrator.services.utils.prompt_utility import validate_prompt_template
 
 logger = logging.getLogger(__name__)
 
@@ -90,29 +82,3 @@ async def generate_and_split_sentences(
         )
 
     return SentenceGenerationResponse(sentences=sentences)
-
-
-def validate_prompt_template(prompt: PromptTemplate):
-    """
-    Prompt template validation
-
-    Args:
-        prompt: The prompt template
-
-    Returns:
-        Nothing.
-    Raises:
-        GenAIPromptTemplateException: if template is incorrect
-    """
-    if PromptFormatter.JINJA2 == prompt.formatter:
-        try:
-            Template(prompt.template).render(prompt.inputs)
-        except TemplateError as exc:
-            logger.error('Prompt completion - template validation failed!')
-            logger.error(exc)
-            raise GenAIPromptTemplateException(
-                ErrorInfo(
-                    error=exc.__class__.__name__,
-                    cause=str(exc),
-                )
-            )

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/services/langchain/callbacks/rag_callback_handler.py

@@ -0,0 +1,61 @@
+#   Copyright (C) 2023-2024 Credit Mutuel Arkea
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+"""Retriever callback handler for LangChain."""
+
+import logging
+from typing import Any, Dict, Optional
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain_core.messages import SystemMessage, AIMessage
+from langchain_core.prompt_values import ChatPromptValue, StringPromptValue
+
+logger = logging.getLogger(__name__)
+
+
+class RAGCallbackHandler(BaseCallbackHandler):
+    """Customized RAG callback handler that retrieves data from the chain execution."""
+
+    records: Dict[str, Any] = {
+        'chat_prompt': None,
+        'chat_chain_output': None,
+        'rag_prompt': None,
+        'rag_chain_output': None,
+        'documents': None,
+    }
+
+    def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Print out that we are entering a chain."""
+
+        if kwargs['name'] == 'chat_chain_output' and isinstance(inputs, AIMessage):
+            self.records['chat_chain_output'] = inputs.content
+
+        if kwargs['name'] == 'rag_chain_output' and isinstance(inputs, AIMessage):
+            self.records['rag_chain_output'] = inputs.content
+
+        if kwargs['name'] == 'RunnableAssign<answer>' and 'documents' in inputs:
+            self.records['documents'] = inputs['documents']
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        """Print out that we finished a chain.""" # if outputs is instance of StringPromptValue
+
+        if isinstance(outputs, ChatPromptValue):
+            self.records['chat_prompt'] = next(
+                (msg.content for msg in outputs.messages if isinstance(msg, SystemMessage)), None
+            )
+
+        if isinstance(outputs, StringPromptValue):
+            self.records['rag_prompt'] = outputs.text