diff --git a/wren-ai-service/.env.dev.example b/wren-ai-service/.env.dev.example index a245ae5ae..1a5dee9da 100644 --- a/wren-ai-service/.env.dev.example +++ b/wren-ai-service/.env.dev.example @@ -65,4 +65,4 @@ LOGGING_LEVEL=INFO LANGFUSE_ENABLE= LANGFUSE_SECRET_KEY= LANGFUSE_PUBLIC_KEY= -LANGFUSE_HOST=https://cloud.langfuse.com \ No newline at end of file +LANGFUSE_HOST=https://cloud.langfuse.com diff --git a/wren-ai-service/poetry.lock b/wren-ai-service/poetry.lock index 9c43b805f..4f944f477 100644 --- a/wren-ai-service/poetry.lock +++ b/wren-ai-service/poetry.lock @@ -1801,6 +1801,30 @@ files = [ [package.dependencies] langchain-core = ">=0.2.10,<0.3.0" +[[package]] +name = "langfuse" +version = "2.35.0" +description = "A client library for accessing langfuse" +optional = false +python-versions = "<4.0,>=3.8.1" +files = [ + {file = "langfuse-2.35.0-py3-none-any.whl", hash = "sha256:e9df2474a01f8e167b7b13674c554915415b27064e48ad207054475f7fa8f82d"}, + {file = "langfuse-2.35.0.tar.gz", hash = "sha256:b1d4b478233eefbc8a6fc63ca00ca82f6afecf2b0fdc1835ca65e751cf901577"}, +] + +[package.dependencies] +backoff = ">=1.10.0" +httpx = ">=0.15.4,<1.0" +idna = ">=3.7,<4.0" +packaging = ">=23.2,<24.0" +pydantic = ">=1.10.7,<3.0" +wrapt = ">=1.14,<2.0" + +[package.extras] +langchain = ["langchain (>=0.0.309)"] +llama-index = ["llama-index (>=0.10.12,<2.0.0)"] +openai = ["openai (>=0.27.8)"] + [[package]] name = "langsmith" version = "0.1.84" @@ -2385,13 +2409,13 @@ files = [ [[package]] name = "packaging" -version = "24.1" +version = "23.2" description = "Core utilities for Python packages" optional = false -python-versions = ">=3.8" +python-versions = ">=3.7" files = [ - {file = "packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124"}, - {file = "packaging-24.1.tar.gz", hash = "sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002"}, + {file = "packaging-23.2-py3-none-any.whl", hash = "sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7"}, + {file = "packaging-23.2.tar.gz", hash = "sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5"}, ] [[package]] @@ -4387,6 +4411,85 @@ MarkupSafe = ">=2.1.1" [package.extras] watchdog = ["watchdog (>=2.3)"] +[[package]] +name = "wrapt" +version = "1.16.0" +description = "Module for decorators, wrappers and monkey patching." +optional = false +python-versions = ">=3.6" +files = [ + {file = "wrapt-1.16.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ffa565331890b90056c01db69c0fe634a776f8019c143a5ae265f9c6bc4bd6d4"}, + {file = "wrapt-1.16.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e4fdb9275308292e880dcbeb12546df7f3e0f96c6b41197e0cf37d2826359020"}, + {file = "wrapt-1.16.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb2dee3874a500de01c93d5c71415fcaef1d858370d405824783e7a8ef5db440"}, + {file = "wrapt-1.16.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2a88e6010048489cda82b1326889ec075a8c856c2e6a256072b28eaee3ccf487"}, + {file = "wrapt-1.16.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ac83a914ebaf589b69f7d0a1277602ff494e21f4c2f743313414378f8f50a4cf"}, + {file = "wrapt-1.16.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:73aa7d98215d39b8455f103de64391cb79dfcad601701a3aa0dddacf74911d72"}, + {file = "wrapt-1.16.0-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:807cc8543a477ab7422f1120a217054f958a66ef7314f76dd9e77d3f02cdccd0"}, + {file = "wrapt-1.16.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:bf5703fdeb350e36885f2875d853ce13172ae281c56e509f4e6eca049bdfb136"}, + {file = "wrapt-1.16.0-cp310-cp310-win32.whl", hash = "sha256:f6b2d0c6703c988d334f297aa5df18c45e97b0af3679bb75059e0e0bd8b1069d"}, + {file = "wrapt-1.16.0-cp310-cp310-win_amd64.whl", hash = "sha256:decbfa2f618fa8ed81c95ee18a387ff973143c656ef800c9f24fb7e9c16054e2"}, + {file = "wrapt-1.16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1a5db485fe2de4403f13fafdc231b0dbae5eca4359232d2efc79025527375b09"}, + {file = "wrapt-1.16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:75ea7d0ee2a15733684badb16de6794894ed9c55aa5e9903260922f0482e687d"}, + {file = "wrapt-1.16.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a452f9ca3e3267cd4d0fcf2edd0d035b1934ac2bd7e0e57ac91ad6b95c0c6389"}, + {file = "wrapt-1.16.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:43aa59eadec7890d9958748db829df269f0368521ba6dc68cc172d5d03ed8060"}, + {file = "wrapt-1.16.0-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:72554a23c78a8e7aa02abbd699d129eead8b147a23c56e08d08dfc29cfdddca1"}, + {file = "wrapt-1.16.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:d2efee35b4b0a347e0d99d28e884dfd82797852d62fcd7ebdeee26f3ceb72cf3"}, + {file = "wrapt-1.16.0-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:6dcfcffe73710be01d90cae08c3e548d90932d37b39ef83969ae135d36ef3956"}, + {file = "wrapt-1.16.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:eb6e651000a19c96f452c85132811d25e9264d836951022d6e81df2fff38337d"}, + {file = "wrapt-1.16.0-cp311-cp311-win32.whl", hash = "sha256:66027d667efe95cc4fa945af59f92c5a02c6f5bb6012bff9e60542c74c75c362"}, + {file = "wrapt-1.16.0-cp311-cp311-win_amd64.whl", hash = "sha256:aefbc4cb0a54f91af643660a0a150ce2c090d3652cf4052a5397fb2de549cd89"}, + {file = "wrapt-1.16.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:5eb404d89131ec9b4f748fa5cfb5346802e5ee8836f57d516576e61f304f3b7b"}, + {file = "wrapt-1.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:9090c9e676d5236a6948330e83cb89969f433b1943a558968f659ead07cb3b36"}, + {file = "wrapt-1.16.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:94265b00870aa407bd0cbcfd536f17ecde43b94fb8d228560a1e9d3041462d73"}, + {file = "wrapt-1.16.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f2058f813d4f2b5e3a9eb2eb3faf8f1d99b81c3e51aeda4b168406443e8ba809"}, + {file = "wrapt-1.16.0-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:98b5e1f498a8ca1858a1cdbffb023bfd954da4e3fa2c0cb5853d40014557248b"}, + {file = "wrapt-1.16.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:14d7dc606219cdd7405133c713f2c218d4252f2a469003f8c46bb92d5d095d81"}, + {file = "wrapt-1.16.0-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:49aac49dc4782cb04f58986e81ea0b4768e4ff197b57324dcbd7699c5dfb40b9"}, + {file = "wrapt-1.16.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:418abb18146475c310d7a6dc71143d6f7adec5b004ac9ce08dc7a34e2babdc5c"}, + {file = "wrapt-1.16.0-cp312-cp312-win32.whl", hash = "sha256:685f568fa5e627e93f3b52fda002c7ed2fa1800b50ce51f6ed1d572d8ab3e7fc"}, + {file = "wrapt-1.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:dcdba5c86e368442528f7060039eda390cc4091bfd1dca41e8046af7c910dda8"}, + {file = "wrapt-1.16.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:d462f28826f4657968ae51d2181a074dfe03c200d6131690b7d65d55b0f360f8"}, + {file = "wrapt-1.16.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a33a747400b94b6d6b8a165e4480264a64a78c8a4c734b62136062e9a248dd39"}, + {file = "wrapt-1.16.0-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b3646eefa23daeba62643a58aac816945cadc0afaf21800a1421eeba5f6cfb9c"}, + {file = "wrapt-1.16.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ebf019be5c09d400cf7b024aa52b1f3aeebeff51550d007e92c3c1c4afc2a40"}, + {file = "wrapt-1.16.0-cp36-cp36m-musllinux_1_1_aarch64.whl", hash = "sha256:0d2691979e93d06a95a26257adb7bfd0c93818e89b1406f5a28f36e0d8c1e1fc"}, + {file = "wrapt-1.16.0-cp36-cp36m-musllinux_1_1_i686.whl", hash = "sha256:1acd723ee2a8826f3d53910255643e33673e1d11db84ce5880675954183ec47e"}, + {file = "wrapt-1.16.0-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:bc57efac2da352a51cc4658878a68d2b1b67dbe9d33c36cb826ca449d80a8465"}, + {file = "wrapt-1.16.0-cp36-cp36m-win32.whl", hash = "sha256:da4813f751142436b075ed7aa012a8778aa43a99f7b36afe9b742d3ed8bdc95e"}, + {file = "wrapt-1.16.0-cp36-cp36m-win_amd64.whl", hash = "sha256:6f6eac2360f2d543cc875a0e5efd413b6cbd483cb3ad7ebf888884a6e0d2e966"}, + {file = "wrapt-1.16.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:a0ea261ce52b5952bf669684a251a66df239ec6d441ccb59ec7afa882265d593"}, + {file = "wrapt-1.16.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7bd2d7ff69a2cac767fbf7a2b206add2e9a210e57947dd7ce03e25d03d2de292"}, + {file = "wrapt-1.16.0-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9159485323798c8dc530a224bd3ffcf76659319ccc7bbd52e01e73bd0241a0c5"}, + {file = "wrapt-1.16.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a86373cf37cd7764f2201b76496aba58a52e76dedfaa698ef9e9688bfd9e41cf"}, + {file = "wrapt-1.16.0-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:73870c364c11f03ed072dda68ff7aea6d2a3a5c3fe250d917a429c7432e15228"}, + {file = "wrapt-1.16.0-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:b935ae30c6e7400022b50f8d359c03ed233d45b725cfdd299462f41ee5ffba6f"}, + {file = "wrapt-1.16.0-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:db98ad84a55eb09b3c32a96c576476777e87c520a34e2519d3e59c44710c002c"}, + {file = "wrapt-1.16.0-cp37-cp37m-win32.whl", hash = "sha256:9153ed35fc5e4fa3b2fe97bddaa7cbec0ed22412b85bcdaf54aeba92ea37428c"}, + {file = "wrapt-1.16.0-cp37-cp37m-win_amd64.whl", hash = "sha256:66dfbaa7cfa3eb707bbfcd46dab2bc6207b005cbc9caa2199bcbc81d95071a00"}, + {file = "wrapt-1.16.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1dd50a2696ff89f57bd8847647a1c363b687d3d796dc30d4dd4a9d1689a706f0"}, + {file = "wrapt-1.16.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:44a2754372e32ab315734c6c73b24351d06e77ffff6ae27d2ecf14cf3d229202"}, + {file = "wrapt-1.16.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e9723528b9f787dc59168369e42ae1c3b0d3fadb2f1a71de14531d321ee05b0"}, + {file = "wrapt-1.16.0-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dbed418ba5c3dce92619656802cc5355cb679e58d0d89b50f116e4a9d5a9603e"}, + {file = "wrapt-1.16.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:941988b89b4fd6b41c3f0bfb20e92bd23746579736b7343283297c4c8cbae68f"}, + {file = "wrapt-1.16.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:6a42cd0cfa8ffc1915aef79cb4284f6383d8a3e9dcca70c445dcfdd639d51267"}, + {file = "wrapt-1.16.0-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:1ca9b6085e4f866bd584fb135a041bfc32cab916e69f714a7d1d397f8c4891ca"}, + {file = "wrapt-1.16.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:d5e49454f19ef621089e204f862388d29e6e8d8b162efce05208913dde5b9ad6"}, + {file = "wrapt-1.16.0-cp38-cp38-win32.whl", hash = "sha256:c31f72b1b6624c9d863fc095da460802f43a7c6868c5dda140f51da24fd47d7b"}, + {file = "wrapt-1.16.0-cp38-cp38-win_amd64.whl", hash = "sha256:490b0ee15c1a55be9c1bd8609b8cecd60e325f0575fc98f50058eae366e01f41"}, + {file = "wrapt-1.16.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9b201ae332c3637a42f02d1045e1d0cccfdc41f1f2f801dafbaa7e9b4797bfc2"}, + {file = "wrapt-1.16.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:2076fad65c6736184e77d7d4729b63a6d1ae0b70da4868adeec40989858eb3fb"}, + {file = "wrapt-1.16.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c5cd603b575ebceca7da5a3a251e69561bec509e0b46e4993e1cac402b7247b8"}, + {file = "wrapt-1.16.0-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b47cfad9e9bbbed2339081f4e346c93ecd7ab504299403320bf85f7f85c7d46c"}, + {file = "wrapt-1.16.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f8212564d49c50eb4565e502814f694e240c55551a5f1bc841d4fcaabb0a9b8a"}, + {file = "wrapt-1.16.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:5f15814a33e42b04e3de432e573aa557f9f0f56458745c2074952f564c50e664"}, + {file = "wrapt-1.16.0-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:db2e408d983b0e61e238cf579c09ef7020560441906ca990fe8412153e3b291f"}, + {file = "wrapt-1.16.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:edfad1d29c73f9b863ebe7082ae9321374ccb10879eeabc84ba3b69f2579d537"}, + {file = "wrapt-1.16.0-cp39-cp39-win32.whl", hash = "sha256:ed867c42c268f876097248e05b6117a65bcd1e63b779e916fe2e33cd6fd0d3c3"}, + {file = "wrapt-1.16.0-cp39-cp39-win_amd64.whl", hash = "sha256:eb1b046be06b0fce7249f1d025cd359b4b80fc1c3e24ad9eca33e0dcdb2e4a35"}, + {file = "wrapt-1.16.0-py3-none-any.whl", hash = "sha256:6906c4100a8fcbf2fa735f6059214bb13b97f75b1a61777fcf6432121ef12ef1"}, + {file = "wrapt-1.16.0.tar.gz", hash = "sha256:5f370f952971e7d17c7d1ead40e49f32345a7f7a5373571ef44d800d06b1899d"}, +] + [[package]] name = "xxhash" version = "3.4.1" @@ -4681,4 +4784,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"] [metadata] lock-version = "2.0" python-versions = ">=3.12.*, <4.0" -content-hash = "223a6ef9ee91ccec9e8d9d2e664cd5a6b8a802b0e75252bdce89d080b79793dd" +content-hash = "ab2603374269e8b84103ee20b426af83db53b69b4322e69283c5ed37336a6bcf" diff --git a/wren-ai-service/pyproject.toml b/wren-ai-service/pyproject.toml index bda99e281..551e3deaf 100644 --- a/wren-ai-service/pyproject.toml +++ b/wren-ai-service/pyproject.toml @@ -23,6 +23,7 @@ orjson = "==3.10.3" sf-hamilton = {version = "==1.63.0", extras = ["visualization"]} aiohttp = "==3.9.5" ollama-haystack = "==0.0.6" +langfuse = "==2.35.0" [tool.poetry.group.dev.dependencies] pytest = "==8.2.0" diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py index 9b0fef5b6..9d278b3be 100644 --- a/wren-ai-service/src/__main__.py +++ b/wren-ai-service/src/__main__.py @@ -9,7 +9,7 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import ORJSONResponse, RedirectResponse import src.globals as container -from src.utils import load_env_vars, setup_custom_logger +from src.utils import init_langfuse, load_env_vars, setup_custom_logger from src.web.v1 import routers env = load_env_vars() @@ -26,6 +26,7 @@ setup_custom_logger( async def lifespan(app: FastAPI): # startup events container.init_globals() + init_langfuse() yield diff --git a/wren-ai-service/src/pipelines/ask/followup_generation.py b/wren-ai-service/src/pipelines/ask/followup_generation.py index 1d6b8aff4..76d7c489e 100644 --- a/wren-ai-service/src/pipelines/ask/followup_generation.py +++ b/wren-ai-service/src/pipelines/ask/followup_generation.py @@ -7,6 +7,7 @@ from hamilton import base from hamilton.experimental.h_async import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline, async_validate @@ -130,6 +131,7 @@ Let's think step by step. ## Start of Pipeline @timer +@observe(capture_input=False) def prompt( query: str, documents: List[Document], @@ -146,12 +148,14 @@ def prompt( @async_timer +@observe(as_type="generation", capture_input=False) async def generate(prompt: dict, generator: Any) -> dict: logger.debug(f"prompt: {prompt}") return await generator.run(prompt=prompt.get("prompt")) @async_timer +@observe(capture_input=False) async def post_process(generate: dict, post_processor: GenerationPostProcessor) -> dict: logger.debug(f"generate: {generate}") return await post_processor.run(generate.get("replies")) @@ -205,6 +209,7 @@ class FollowUpGeneration(BasicPipeline): ) @async_timer + @observe(name="Ask Follow Up Generation") async def run( self, query: str, @@ -227,9 +232,12 @@ class FollowUpGeneration(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() llm_provider, _, _, engine = init_providers() pipeline = FollowUpGeneration(llm_provider=llm_provider, engine=engine) @@ -250,3 +258,5 @@ if __name__ == "__main__": ), ) ) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/ask/generation.py b/wren-ai-service/src/pipelines/ask/generation.py index 018af68a6..29e521b8b 100644 --- a/wren-ai-service/src/pipelines/ask/generation.py +++ b/wren-ai-service/src/pipelines/ask/generation.py @@ -7,6 +7,7 @@ from hamilton import base from hamilton.experimental.h_async import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline, async_validate @@ -90,6 +91,7 @@ Let's think step by step. ## Start of Pipeline @timer +@observe(capture_input=False) def prompt( query: str, documents: List[Document], @@ -106,12 +108,14 @@ def prompt( @async_timer +@observe(as_type="generation", capture_input=False) async def generate(prompt: dict, generator: Any) -> dict: logger.debug(f"prompt: {prompt}") return await generator.run(prompt=prompt.get("prompt")) @async_timer +@observe(capture_input=False) async def post_process(generate: dict, post_processor: GenerationPostProcessor) -> dict: logger.debug(f"generate: {generate}") return await post_processor.run(generate.get("replies")) @@ -163,6 +167,7 @@ class Generation(BasicPipeline): ) @async_timer + @observe(name="Ask Generation") async def run( self, query: str, @@ -185,9 +190,12 @@ class Generation(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() llm_provider, _, _, engine = init_providers() pipeline = Generation( @@ -197,3 +205,5 @@ if __name__ == "__main__": pipeline.visualize("this is a test query", [], []) async_validate(lambda: pipeline.run("this is a test query", [], [])) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/ask/historical_question.py b/wren-ai-service/src/pipelines/ask/historical_question.py index 556a7e482..47059f1b9 100644 --- a/wren-ai-service/src/pipelines/ask/historical_question.py +++ b/wren-ai-service/src/pipelines/ask/historical_question.py @@ -7,6 +7,7 @@ from typing import Any, Dict, List, Optional from hamilton import base from hamilton.experimental.h_async import AsyncDriver from haystack import Document, component +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline, async_validate from src.core.provider import DocumentStoreProvider, EmbedderProvider @@ -56,12 +57,14 @@ class OutputFormatter: ## Start of Pipeline @async_timer +@observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any) -> dict: logger.debug(f"query: {query}") return await embedder.run(query) @async_timer +@observe(capture_input=False) async def retrieval(embedding: dict, retriever: Any) -> dict: res = await retriever.run(query_embedding=embedding.get("embedding")) documents = res.get("documents") @@ -69,12 +72,14 @@ async def retrieval(embedding: dict, retriever: Any) -> dict: @timer +@observe(capture_input=False) def filtered_documents(retrieval: dict, score_filter: ScoreFilter) -> dict: logger.debug(f"retrieval: {retrieval}") return score_filter.run(documents=retrieval.get("documents")) @timer +@observe(capture_input=False) def formatted_output( filtered_documents: dict, output_formatter: OutputFormatter ) -> dict: @@ -126,6 +131,7 @@ class HistoricalQuestion(BasicPipeline): ) @async_timer + @observe(name="Ask Historical Question") async def run(self, query: str): logger.info("Ask HistoricalQuestion pipeline is running...") return await self._pipe.execute( @@ -141,9 +147,12 @@ class HistoricalQuestion(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() _, embedder_provider, document_store_provider, _ = init_providers() @@ -153,3 +162,5 @@ if __name__ == "__main__": pipeline.visualize("this is a query") async_validate(lambda: pipeline.run("this is a query")) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/ask/retrieval.py b/wren-ai-service/src/pipelines/ask/retrieval.py index b01e146a5..152b6ad2e 100644 --- a/wren-ai-service/src/pipelines/ask/retrieval.py +++ b/wren-ai-service/src/pipelines/ask/retrieval.py @@ -5,6 +5,7 @@ from typing import Any from hamilton import base from hamilton.experimental.h_async import AsyncDriver +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline, async_validate from src.core.provider import DocumentStoreProvider, EmbedderProvider @@ -15,12 +16,14 @@ logger = logging.getLogger("wren-ai-service") ## Start of Pipeline @async_timer +@observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any) -> dict: logger.debug(f"query: {query}") return await embedder.run(query) @async_timer +@observe(capture_input=False) async def retrieval(embedding: dict, retriever: Any) -> dict: return await retriever.run(query_embedding=embedding.get("embedding")) @@ -64,6 +67,7 @@ class Retrieval(BasicPipeline): ) @async_timer + @observe(name="Ask Retrieval") async def run(self, query: str): logger.info("Ask Retrieval pipeline is running...") return await self._pipe.execute( @@ -77,9 +81,12 @@ class Retrieval(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() _, embedder_provider, document_store_provider, _ = init_providers() pipeline = Retrieval( @@ -89,3 +96,5 @@ if __name__ == "__main__": pipeline.visualize("this is a query") async_validate(lambda: pipeline.run("this is a query")) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/ask/sql_correction.py b/wren-ai-service/src/pipelines/ask/sql_correction.py index b29f7c93f..b6cd76d36 100644 --- a/wren-ai-service/src/pipelines/ask/sql_correction.py +++ b/wren-ai-service/src/pipelines/ask/sql_correction.py @@ -7,6 +7,7 @@ from hamilton import base from hamilton.experimental.h_async import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline, async_validate @@ -58,6 +59,7 @@ Let's think step by step. ## Start of Pipeline @timer +@observe(capture_input=False) def prompt( documents: List[Document], invalid_generation_results: List[Dict], @@ -74,12 +76,14 @@ def prompt( @async_timer +@observe(as_type="generation", capture_input=False) async def generate(prompt: dict, generator: Any) -> dict: logger.debug(f"prompt: {prompt}") return await generator.run(prompt=prompt.get("prompt")) @async_timer +@observe(capture_input=False) async def post_process(generate: dict, post_processor: GenerationPostProcessor) -> dict: logger.debug(f"generate: {generate}") return await post_processor.run(generate.get("replies")) @@ -131,6 +135,7 @@ class SQLCorrection(BasicPipeline): ) @async_timer + @observe(name="Ask SQL Correction") async def run( self, contexts: List[Document], @@ -151,9 +156,12 @@ class SQLCorrection(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() llm_provider, _, _, engine = init_providers() pipeline = SQLCorrection( @@ -163,3 +171,5 @@ if __name__ == "__main__": pipeline.visualize([], []) async_validate(lambda: pipeline.run([], [])) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/ask_details/generation.py b/wren-ai-service/src/pipelines/ask_details/generation.py index 6b1846fc7..ad0933fdf 100644 --- a/wren-ai-service/src/pipelines/ask_details/generation.py +++ b/wren-ai-service/src/pipelines/ask_details/generation.py @@ -10,6 +10,7 @@ from hamilton import base from hamilton.experimental.h_async import AsyncDriver from haystack import component from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from src.core.engine import ( Engine, @@ -121,18 +122,21 @@ class GenerationPostProcessor: ## Start of Pipeline @timer +@observe(capture_input=False) def prompt(sql: str, prompt_builder: PromptBuilder) -> dict: logger.debug(f"sql: {sql}") return prompt_builder.run(sql=sql) @async_timer +@observe(as_type="generation", capture_input=False) async def generate(prompt: dict, generator: Any) -> dict: logger.debug(f"prompt: {prompt}") return await generator.run(prompt=prompt.get("prompt")) @async_timer +@observe(capture_input=False) async def post_process(generate: dict, post_processor: GenerationPostProcessor) -> dict: logger.debug(f"generate: {generate}") return await post_processor.run(generate.get("replies")) @@ -176,6 +180,7 @@ class Generation(BasicPipeline): ) @async_timer + @observe(name="Ask_Details Generation") async def run(self, sql: str): logger.info("Ask_Details Generation pipeline is running...") return await self._pipe.execute( @@ -190,9 +195,12 @@ class Generation(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context, observe + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() llm_provider, _, _, engine = init_providers() pipeline = Generation( @@ -202,3 +210,5 @@ if __name__ == "__main__": pipeline.visualize("SELECT * FROM table_name") async_validate(lambda: pipeline.run("SELECT * FROM table_name")) + + langfuse_context.flush() diff --git a/wren-ai-service/src/pipelines/indexing/indexing.py b/wren-ai-service/src/pipelines/indexing/indexing.py index 613b8fb3d..bd16049e3 100644 --- a/wren-ai-service/src/pipelines/indexing/indexing.py +++ b/wren-ai-service/src/pipelines/indexing/indexing.py @@ -12,6 +12,7 @@ from hamilton.function_modifiers import extract_fields from haystack import Document, component from haystack.components.writers import DocumentWriter from haystack.document_stores.types import DocumentStore, DuplicatePolicy +from langfuse.decorators import observe from tqdm import tqdm from src.core.pipeline import BasicPipeline, async_validate @@ -313,12 +314,14 @@ class DDLConverter: ## Start of Pipeline @timer +@observe(capture_input=False, capture_output=False) def clean_document_store(mdl_str: str, cleaner: DocumentCleaner) -> Dict[str, Any]: logger.debug(f"input in clean_document_store: {mdl_str}") return cleaner.run(mdl=mdl_str) @timer +@observe(capture_input=False, capture_output=False) @extract_fields(dict(mdl=Dict[str, Any])) def validate_mdl( clean_document_store: Dict[str, Any], validator: MDLValidator @@ -330,12 +333,14 @@ def validate_mdl( @timer +@observe(capture_input=False) def convert_to_ddl(mdl: Dict[str, Any], ddl_converter: DDLConverter) -> Dict[str, Any]: logger.debug(f"input in convert_to_ddl: {mdl}") return ddl_converter.run(mdl=mdl) @async_timer +@observe(capture_input=False, capture_output=False) async def embed_ddl( convert_to_ddl: Dict[str, Any], ddl_embedder: Any ) -> Dict[str, Any]: @@ -344,12 +349,14 @@ async def embed_ddl( @timer +@observe(capture_input=False) def write_ddl(embed_ddl: Dict[str, Any], ddl_writer: DocumentWriter) -> None: logger.debug(f"input in write_ddl: {embed_ddl}") return ddl_writer.run(documents=embed_ddl["documents"]) @timer +@observe(capture_input=False) def convert_to_view( mdl: Dict[str, Any], view_converter: ViewConverter ) -> Dict[str, Any]: @@ -358,6 +365,7 @@ def convert_to_view( @async_timer +@observe(capture_input=False, capture_output=False) async def embed_view( convert_to_view: Dict[str, Any], view_embedder: Any ) -> Dict[str, Any]: @@ -366,6 +374,7 @@ async def embed_view( @timer +@observe(capture_input=False) def write_view(embed_view: Dict[str, Any], view_writer: DocumentWriter) -> None: logger.debug(f"input in write_view: {embed_view}") return view_writer.run(documents=embed_view["documents"]) @@ -427,6 +436,7 @@ class Indexing(BasicPipeline): ) @async_timer + @observe(name="Ask Indexing") async def run(self, mdl_str: str) -> Dict[str, Any]: logger.info("Ask Indexing pipeline is running...") return await self._pipe.execute( @@ -446,9 +456,13 @@ class Indexing(BasicPipeline): if __name__ == "__main__": - from src.utils import load_env_vars + from langfuse.decorators import langfuse_context + + from src.utils import init_langfuse, load_env_vars load_env_vars() + init_langfuse() + _, embedder_provider, document_store_provider, _ = init_providers() pipeline = Indexing( @@ -459,3 +473,5 @@ if __name__ == "__main__": input = '{"models": [], "views": [], "relationships": [], "metrics": []}' pipeline.visualize(input) async_validate(lambda: pipeline.run(input)) + + langfuse_context.flush() diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py index deeb1f8d6..296537038 100644 --- a/wren-ai-service/src/utils.py +++ b/wren-ai-service/src/utils.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Tuple from dotenv import load_dotenv +from langfuse.decorators import langfuse_context from src.core.engine import Engine from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider @@ -121,3 +122,10 @@ def async_timer(func): def remove_trailing_slash(endpoint: str) -> str: return endpoint.rstrip("/") if endpoint.endswith("/") else endpoint + + +def init_langfuse(): + enabled = os.getenv("LANGFUSE_ENABLE", "false") + langfuse_context.configure( + enabled=False if enabled.lower() == "false" else True, + ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 83c537add..201abd927 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2,6 +2,7 @@ import logging from typing import List, Literal, Optional import sqlparse +from langfuse.decorators import observe from pydantic import BaseModel from src.core.pipeline import BasicPipeline @@ -114,6 +115,7 @@ class AskService: self._prepare_semantics_statuses = {} @async_timer + @observe(name="Prepare Semantics") async def prepare_semantics( self, prepare_semantics_request: SemanticsPreparationRequest ): @@ -163,6 +165,7 @@ class AskService: return False @async_timer + @observe(name="Ask Question") async def ask( self, ask_request: AskRequest, diff --git a/wren-ai-service/src/web/v1/services/ask_details.py b/wren-ai-service/src/web/v1/services/ask_details.py index 8eba91b0c..75a0c22c2 100644 --- a/wren-ai-service/src/web/v1/services/ask_details.py +++ b/wren-ai-service/src/web/v1/services/ask_details.py @@ -2,6 +2,7 @@ import logging from typing import List, Literal, Optional from haystack import Pipeline +from langfuse.decorators import observe from pydantic import BaseModel from src.core.engine import add_quotes @@ -64,6 +65,7 @@ class AskDetailsService: self._ask_details_results = {} @async_timer + @observe(name="Ask Details(Breakdown SQL)") async def ask_details( self, ask_details_request: AskDetailsRequest,