From 0f3453228fd937a87ae073ea9d4195d02920f6c2 Mon Sep 17 00:00:00 2001 From: panda_zhu <183920368@qq.com> Date: Sat, 15 Jun 2024 16:07:29 +0800 Subject: [PATCH] Enable langfuse for file_chat, knowledge_base_chat, search_engine_chat, and agent_chat (#3229) * Enable langchain-chatchat to support langfuse * Enable langchain-chatchat to support langfuse --- server/chat/agent_chat.py | 13 ++++++++++++- server/chat/chat.py | 12 ++++++++++++ server/chat/file_chat.py | 14 +++++++++++++- server/chat/knowledge_base_chat.py | 14 +++++++++++++- server/chat/search_engine_chat.py | 13 ++++++++++++- 5 files changed, 62 insertions(+), 4 deletions(-) diff --git a/server/chat/agent_chat.py b/server/chat/agent_chat.py index 41bf5baba7..c25450cd61 100644 --- a/server/chat/agent_chat.py +++ b/server/chat/agent_chat.py @@ -48,11 +48,22 @@ async def agent_chat_iterator( if isinstance(max_tokens, int) and max_tokens <= 0: max_tokens = None + callbacks = [callback] + # Enable langchain-chatchat to support langfuse + import os + langfuse_secret_key = os.environ.get('LANGFUSE_SECRET_KEY') + langfuse_public_key = os.environ.get('LANGFUSE_PUBLIC_KEY') + langfuse_host = os.environ.get('LANGFUSE_HOST') + if langfuse_secret_key and langfuse_public_key and langfuse_host: + from langfuse import Langfuse + from langfuse.callback import CallbackHandler + langfuse_handler = CallbackHandler() + callbacks.append(langfuse_handler) model = get_ChatOpenAI( model_name=model_name, temperature=temperature, max_tokens=max_tokens, - callbacks=[callback], + callbacks=callbacks, ) kb_list = {x["kb_name"]: x for x in get_kb_details()} diff --git a/server/chat/chat.py b/server/chat/chat.py index 5783829c58..3f978c5106 100644 --- a/server/chat/chat.py +++ b/server/chat/chat.py @@ -47,6 +47,18 @@ async def chat_iterator() -> AsyncIterable[str]: query=query) callbacks.append(conversation_callback) + # Enable langchain-chatchat to support langfuse + import os + langfuse_secret_key = os.environ.get('LANGFUSE_SECRET_KEY') + langfuse_public_key = os.environ.get('LANGFUSE_PUBLIC_KEY') + langfuse_host = os.environ.get('LANGFUSE_HOST') + if langfuse_secret_key and langfuse_public_key and langfuse_host : + from langfuse import Langfuse + from langfuse.callback import CallbackHandler + langfuse_handler = CallbackHandler() + callbacks.append(langfuse_handler) + + if isinstance(max_tokens, int) and max_tokens <= 0: max_tokens = None diff --git a/server/chat/file_chat.py b/server/chat/file_chat.py index a3bbdfc6ef..23ed6c97b1 100644 --- a/server/chat/file_chat.py +++ b/server/chat/file_chat.py @@ -118,11 +118,23 @@ async def knowledge_base_chat_iterator() -> AsyncIterable[str]: if isinstance(max_tokens, int) and max_tokens <= 0: max_tokens = None + callbacks = [callback] + # Enable langchain-chatchat to support langfuse + import os + langfuse_secret_key = os.environ.get('LANGFUSE_SECRET_KEY') + langfuse_public_key = os.environ.get('LANGFUSE_PUBLIC_KEY') + langfuse_host = os.environ.get('LANGFUSE_HOST') + if langfuse_secret_key and langfuse_public_key and langfuse_host: + from langfuse import Langfuse + from langfuse.callback import CallbackHandler + langfuse_handler = CallbackHandler() + callbacks.append(langfuse_handler) + model = get_ChatOpenAI( model_name=model_name, temperature=temperature, max_tokens=max_tokens, - callbacks=[callback], + callbacks=callbacks, ) embed_func = EmbeddingsFunAdapter() embeddings = await embed_func.aembed_query(query) diff --git a/server/chat/knowledge_base_chat.py b/server/chat/knowledge_base_chat.py index 8a9197f5d4..edf2a455e0 100644 --- a/server/chat/knowledge_base_chat.py +++ b/server/chat/knowledge_base_chat.py @@ -70,11 +70,23 @@ async def knowledge_base_chat_iterator( if isinstance(max_tokens, int) and max_tokens <= 0: max_tokens = None + callbacks = [callback] + # Enable langchain-chatchat to support langfuse + import os + langfuse_secret_key = os.environ.get('LANGFUSE_SECRET_KEY') + langfuse_public_key = os.environ.get('LANGFUSE_PUBLIC_KEY') + langfuse_host = os.environ.get('LANGFUSE_HOST') + if langfuse_secret_key and langfuse_public_key and langfuse_host : + from langfuse import Langfuse + from langfuse.callback import CallbackHandler + langfuse_handler = CallbackHandler() + callbacks.append(langfuse_handler) + model = get_ChatOpenAI( model_name=model_name, temperature=temperature, max_tokens=max_tokens, - callbacks=[callback], + callbacks=callbacks, ) docs = await run_in_threadpool(search_docs, query=query, diff --git a/server/chat/search_engine_chat.py b/server/chat/search_engine_chat.py index 42bef3c23b..c35c3c102a 100644 --- a/server/chat/search_engine_chat.py +++ b/server/chat/search_engine_chat.py @@ -154,11 +154,22 @@ async def search_engine_chat_iterator(query: str, if isinstance(max_tokens, int) and max_tokens <= 0: max_tokens = None + callbacks = [callback] + # Enable langchain-chatchat to support langfuse + import os + langfuse_secret_key = os.environ.get('LANGFUSE_SECRET_KEY') + langfuse_public_key = os.environ.get('LANGFUSE_PUBLIC_KEY') + langfuse_host = os.environ.get('LANGFUSE_HOST') + if langfuse_secret_key and langfuse_public_key and langfuse_host: + from langfuse import Langfuse + from langfuse.callback import CallbackHandler + langfuse_handler = CallbackHandler() + callbacks.append(langfuse_handler) model = get_ChatOpenAI( model_name=model_name, temperature=temperature, max_tokens=max_tokens, - callbacks=[callback], + callbacks=callbacks, ) docs = await lookup_search_engine(query, search_engine_name, top_k, split_result=split_result)