Skip to content

Commit 9394a6f

Browse files
chatbot: Add Langfuse feedback support
Co-authored-by: Dave Zuckerman <dzuckerm@berkeley.edu>
1 parent c850d74 commit 9394a6f

File tree

2 files changed

+48
-3
lines changed

2 files changed

+48
-3
lines changed

willa/chatbot/chatbot.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
LANGFUSE_HANDLER = CallbackHandler()
2222
"""The Langfuse callback handler."""
2323

24+
2425
class Chatbot: # pylint: disable=R0903
2526
"""An instance of a Willa chatbot.
2627
@@ -89,6 +90,7 @@ def ask(self, question: str) -> dict[str, str]:
8990

9091
if ai_message:
9192
answers["ai_message"] = str(ai_message[-1].content)
93+
answers["langfuse_trace_id"] = str(LANGFUSE_HANDLER.last_trace_id)
9294

9395
if len(answers) == 0:
9496
return {"no_result": "I'm sorry, I couldn't generate a response."}

willa/web/app.py

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,16 @@
44

55
import logging
66
import os
7+
from typing import Optional
78

89
import chainlit as cl
10+
from chainlit.data import get_data_layer
911
from 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

1215
from willa.chatbot import Chatbot
13-
from willa.config import CONFIG
16+
from willa.config import CONFIG, get_langfuse_client
1417
from willa.web.cas_provider import CASProvider
1518
from willa.web.inject_custom_auth import add_custom_oauth_provider
1619

@@ -34,6 +37,22 @@
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
3857
async 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
5290
def 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

74116
def _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

Comments
 (0)