44
55import logging
66import os
7+ from typing import Optional
78
89import chainlit as cl
10+ from chainlit .data import get_data_layer
911from chainlit .data .chainlit_data_layer import ChainlitDataLayer
10- from chainlit .types import ThreadDict , CommandDict
12+ from chainlit .types import ThreadDict , CommandDict , Feedback
13+ from chainlit .step import StepDict
1114
1215from willa .chatbot import Chatbot
13- from willa .config import CONFIG
16+ from willa .config import CONFIG , get_langfuse_client
1417from willa .web .cas_provider import CASProvider
1518from willa .web .inject_custom_auth import add_custom_oauth_provider
1619
3437]
3538
3639
40+ async def get_step (self : ChainlitDataLayer , step_id : str ) -> Optional [StepDict ]:
41+ """Get step and related feedback"""
42+ query = """
43+ SELECT s.*,
44+ f.id feedback_id,
45+ f.value feedback_value,
46+ f."comment" feedback_comment
47+ FROM "Step" s LEFT JOIN "Feedback" f ON s.id = f."stepId"
48+ WHERE s.id = $1
49+ """
50+ result = await self .execute_query (query , {"step_id" : step_id })
51+ if not result :
52+ return None
53+ return self ._convert_step_row_to_dict (result [0 ]) # pylint: disable="protected-access"
54+
55+
3756@cl .on_chat_start
3857async def ocs () -> None :
3958 """loaded when new chat is started"""
@@ -48,6 +67,25 @@ async def on_chat_resume(thread: ThreadDict) -> None:
4867# pylint: enable="unused-argument"
4968
5069
70+ @cl .on_feedback
71+ async def on_feedback (feedback : Feedback ) -> None :
72+ """Handle feedback."""
73+ step : Optional [StepDict ] = await get_data_layer ().get_step (feedback .forId )
74+ if step is None :
75+ LOGGER .warning ("Feedback left for unknown step %s" , feedback .forId )
76+ return
77+
78+ trace_id : Optional [str ] = step ['metadata' ].get ('langfuse_trace_id' )
79+ get_langfuse_client ().create_score (
80+ name = 'feedback' ,
81+ value = float (feedback .value ),
82+ session_id = step ['threadId' ] if not trace_id else None ,
83+ trace_id = trace_id ,
84+ data_type = 'BOOLEAN' ,
85+ comment = feedback .comment
86+ )
87+
88+
5189@cl .data_layer
5290def data_layer () -> ChainlitDataLayer :
5391 """Retrieve the data layer to use with Chainlit.
@@ -68,7 +106,11 @@ def _secret() -> str:
68106 database_url = os .environ .get (
69107 'DATABASE_URL' , f"postgresql://{ _pg ('USER' )} :{ _secret ()} @{ _pg ('HOST' )} /{ _pg ('DB' )} "
70108 )
71- return ChainlitDataLayer (database_url = database_url )
109+ dl = ChainlitDataLayer (database_url = database_url )
110+ # pylint: disable="no-value-for-parameter"
111+ dl .get_step = get_step .__get__ (dl ) # type: ignore[attr-defined]
112+ # pylint: enable="no-value-for-parameter"
113+ return dl
72114
73115
74116def _get_history () -> str :
@@ -117,6 +159,7 @@ async def chat(message: cl.Message) -> None:
117159
118160 if 'ai_message' in reply :
119161 await cl .Message (content = reply ['ai_message' ]).send ()
162+ cl .context .current_run .metadata ['langfuse_trace_id' ] = reply ['langfuse_trace_id' ]
120163
121164 if 'tind_message' in reply :
122165 tind_refs = cl .CustomElement (
0 commit comments