[GEN AI] Improving Langfuse's tracks

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

@@ -179,7 +179,14 @@ object RAGAnswerHandler : AbstractProactiveAnswerHandler {
             try {
                 val response = ragService.rag(
                     query = RAGQuery(
-                        history = getDialogHistory(dialog),
+                        dialog = DialogDetails(
+                            dialogId = dialog.id.toString(),
+                            userId = dialog.playerIds.firstOrNull { PlayerType.user == it.type }?.id,
+                            history = getDialogHistory(dialog),
+                            tags = listOf(
+                                "connector:${underlyingConnector.connectorType.id}"
+                            )
+                        ),
                         questionAnsweringLlmSetting = ragConfiguration.llmSetting,
                         questionAnsweringPromptInputs = mapOf(
                             "question" to action.toString(),

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

@@ -24,7 +24,7 @@ import ai.tock.genai.orchestratorcore.models.vectorstore.VectorStoreSetting
 data class RAGQuery(
     // val condenseQuestionLlmSetting: LLMSetting,
     // val condenseQuestionPromptInputs: Map<String, String>,
-    val history: List<ChatMessage> = emptyList(),
+    val dialog: DialogDetails,
     val questionAnsweringLlmSetting: LLMSetting,
     val questionAnsweringPromptInputs: Map<String, String>,
     val embeddingQuestionEmSetting: EMSetting,
@@ -34,6 +34,13 @@ data class RAGQuery(
     val observabilitySetting: ObservabilitySetting?
 )
 
+data class DialogDetails(
+    val dialogId: String? = null,
+    val userId: String? = null,
+    val history: List<ChatMessage> = emptyList(),
+    val tags: List<String> = emptyList(),
+)
+
 data class ChatMessage(
     val text: String,
     val type: ChatMessageType,

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/models/rag/rag_models.py

@@ -48,7 +48,7 @@ def __eq__(self, other):
         )
 
     def __hash__(self):
-        return hash((self.title, self.url, self.content))
+        return hash((self.title, str(self.url or ''), self.content))
 
 
 class Footnote(Source):

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

@@ -133,16 +133,22 @@ class VectorStoreProviderSettingStatusQuery(BaseModel):
         default=None,
     )
 
+class DialogDetails(BaseModel):
+    """The dialog details model"""
+
+    dialog_id: Optional[str] = Field(description="The dialog ID.", default=None, examples=["uuid-0123"])
+    user_id: Optional[str] = Field(description="The user ID.", default=None, examples=["[email protected]"])
+    history: list[ChatMessage] = Field(description="Conversation history, used to reformulate the user's question.")
+    tags: list[str] = Field(description='List of tags', examples=[["my-Tag"]])
+
 
 class RagQuery(BaseQuery):
     """The RAG query model"""
 
-    history: list[ChatMessage] = Field(
-        description="Conversation history, used to reformulate the user's question."
-    )
+    dialog: 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 history field',
+        '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.")
@@ -156,7 +162,7 @@ class RagQuery(BaseQuery):
     )
     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 history field',
+        '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."
@@ -182,13 +188,15 @@ class RagQuery(BaseQuery):
         'json_schema_extra': {
             'examples': [
                 {
-                    'history': [
-                        {'text': 'Hello, how can I do this?', 'type': 'HUMAN'},
-                        {
-                            'text': 'you can do this with the following method ....',
-                            'type': 'AI',
-                        },
-                    ],
+                    'dialog' : {
+                        'history': [
+                            {'text': 'Hello, how can I do this?', 'type': 'HUMAN'},
+                            {
+                                'text': 'you can do this with the following method ....',
+                                'type': 'AI',
+                            },
+                        ]
+                    },
                     'question_answering_llm_setting': {
                         'provider': 'OpenAI',
                         'api_key': {

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

@@ -79,7 +79,7 @@ async def generate_and_split_sentences(
         config = {"callbacks": [
             create_observability_callback_handler(
                 observability_setting=query.observability_setting,
-                trace_name=ObservabilityTrace.SENTENCE_GENERATION
+                trace_name=ObservabilityTrace.SENTENCE_GENERATION.value
             )]}
 
     sentences = await chain.ainvoke(query.prompt.inputs, config=config)

gen-ai/orchestrator-server/src/main/python/server/src/gen_ai_orchestrator/services/langchain/factories/langchain_factory.py

@@ -23,7 +23,7 @@
 """
 
 import logging
-from typing import Optional
+from typing import Optional, Any
 
 from langchain_core.embeddings import Embeddings
 from langfuse.callback import CallbackHandler as LangfuseCallbackHandler
@@ -338,22 +338,21 @@ def get_callback_handler_factory(
 
 def create_observability_callback_handler(
     observability_setting: Optional[ObservabilitySetting],
-    trace_name: ObservabilityTrace,
+    **kwargs: Any
 ) -> Optional[LangfuseCallbackHandler]:
     """
     Create the Observability Callback Handler
 
     Args:
         observability_setting: The Observability Settings
-        trace_name: The trace name
 
     Returns:
         The Observability Callback Handler
     """
     if observability_setting is not None:
         return get_callback_handler_factory(
-            setting=observability_setting
-        ).get_callback_handler(trace_name=trace_name.value)
+            setting=observability_setting,
+        ).get_callback_handler(**kwargs)
 
     return None
 

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

@@ -93,15 +93,17 @@ async def execute_qa_chain(query: RagQuery, debug: bool) -> RagResponse:
 
     conversational_retrieval_chain = create_rag_chain(query=query)
 
+    message_history = ChatMessageHistory()
+    if query.dialog:
+        for msg in query.dialog.history:
+            if ChatMessageType.HUMAN == msg.type:
+                message_history.add_user_message(msg.text)
+            else:
+                message_history.add_ai_message(msg.text)
+
     logger.debug(
-        'RAG chain - Use chat history: %s', 'Yes' if len(query.history) > 0 else 'No'
+        'RAG chain - Use chat history: %s', 'Yes' if len(message_history.messages) > 0 else 'No'
     )
-    message_history = ChatMessageHistory()
-    for msg in query.history:
-        if ChatMessageType.HUMAN == msg.type:
-            message_history.add_user_message(msg.text)
-        else:
-            message_history.add_ai_message(msg.text)
 
     inputs = {
         **query.question_answering_prompt_inputs,
@@ -123,7 +125,10 @@ async def execute_qa_chain(query: RagQuery, debug: bool) -> RagResponse:
         callback_handlers.append(
             create_observability_callback_handler(
                 observability_setting=query.observability_setting,
-                trace_name=ObservabilityTrace.RAG,
+                trace_name=ObservabilityTrace.RAG.value,
+                session_id=query.dialog.dialog_id,
+                user_id=query.dialog.user_id,
+                tags=query.dialog.tags,
             )
         )
 

gen-ai/orchestrator-server/src/main/python/server/tests/services/test_rag_chain.py

@@ -99,13 +99,13 @@ async def test_rag_chain(
     """Test the full execute_qa_chain method by mocking all external calls."""
     # Build a test RagQuery
     query_dict = {
-        'history': [
-            {'text': 'Hello, how can I do this?', 'type': 'HUMAN'},
-            {
-                'text': 'you can do this with the following method ....',
-                'type': 'AI',
-            },
-        ],
+        'dialog': {
+            'history': [
+                {'text': 'Hello, how can I do this?', 'type': 'HUMAN'},
+                {'text': 'you can do this with the following method ....', 'type': 'AI'}
+            ],
+            'tags': []
+        },
         'question_answering_llm_setting': {
             'provider': 'OpenAI',
             'api_key': {'type': 'Raw', 'value': 'ab7***************************A1IV4B'},