fix(wren-ai-service): langfuse cost display issue

wren-ai-service/src/pipelines/generation/chart_adjustment.py

@@ -16,6 +16,7 @@
     ChartGenerationResults,
     chart_generation_instructions,
 )
+from src.utils import trace_cost
 from src.web.v1.services.chart_adjustment import ChartAdjustmentOption
 
 logger = logging.getLogger("wren-ai-service")
@@ -108,8 +109,13 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate_chart_adjustment(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@trace_cost
+async def generate_chart_adjustment(
+    prompt: dict,
+    generator: Any,
+    generator_name: str,
+) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -152,6 +158,7 @@ def __init__(
                 system_prompt=chart_adjustment_system_prompt,
                 generation_kwargs=CHART_ADJUSTMENT_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "chart_data_preprocessor": ChartDataPreprocessor(),
             "post_processor": ChartGenerationPostProcessor(),
         }

wren-ai-service/src/pipelines/generation/chart_generation.py

@@ -16,6 +16,7 @@
     ChartGenerationResults,
     chart_generation_instructions,
 )
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -79,8 +80,9 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate_chart(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@trace_cost
+async def generate_chart(prompt: dict, generator: Any, generator_name: str) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -125,6 +127,7 @@ def __init__(
                 system_prompt=chart_generation_system_prompt,
                 generation_kwargs=CHART_GENERATION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "chart_data_preprocessor": ChartDataPreprocessor(),
             "post_processor": ChartGenerationPostProcessor(),
         }

wren-ai-service/src/pipelines/generation/data_assistance.py

@@ -10,6 +10,7 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 from src.web.v1.services.ask import AskHistory
 
 logger = logging.getLogger("wren-ai-service")
@@ -68,11 +69,14 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def data_assistance(prompt: dict, generator: Any, query_id: str) -> dict:
+@trace_cost
+async def data_assistance(
+    prompt: dict, generator: Any, query_id: str, generator_name: str
+) -> dict:
     return await generator(
         prompt=prompt.get("prompt"),
         query_id=query_id,
-    )
+    ), generator_name
 
 
 ## End of Pipeline
@@ -90,6 +94,7 @@ def __init__(
                 system_prompt=data_assistance_system_prompt,
                 streaming_callback=self._streaming_callback,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=data_assistance_user_prompt_template
             ),

wren-ai-service/src/pipelines/generation/followup_sql_generation.py

@@ -20,6 +20,7 @@
     sql_generation_system_prompt,
 )
 from src.pipelines.retrieval.sql_functions import SqlFunction
+from src.utils import trace_cost
 from src.web.v1.services import Configuration
 from src.web.v1.services.ask import AskHistory
 
@@ -110,13 +111,17 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
+@trace_cost
 async def generate_sql_in_followup(
-    prompt: dict, generator: Any, histories: list[AskHistory]
+    prompt: dict,
+    generator: Any,
+    histories: list[AskHistory],
+    generator_name: str,
 ) -> dict:
     history_messages = construct_ask_history_messages(histories)
     return await generator(
         prompt=prompt.get("prompt"), history_messages=history_messages
-    )
+    ), generator_name
 
 
 @observe(capture_input=False)
@@ -149,6 +154,7 @@ def __init__(
                 system_prompt=sql_generation_system_prompt,
                 generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=text_to_sql_with_followup_user_prompt_template
             ),

wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py

@@ -14,6 +14,7 @@
     construct_instructions,
     sql_generation_reasoning_system_prompt,
 )
+from src.utils import trace_cost
 from src.web.v1.services import Configuration
 from src.web.v1.services.ask import AskHistory
 
@@ -84,11 +85,17 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate_sql_reasoning(prompt: dict, generator: Any, query_id: str) -> dict:
+@trace_cost
+async def generate_sql_reasoning(
+    prompt: dict,
+    generator: Any,
+    query_id: str,
+    generator_name: str,
+) -> dict:
     return await generator(
         prompt=prompt.get("prompt"),
         query_id=query_id,
-    )
+    ), generator_name
 
 
 @observe()
@@ -113,6 +120,7 @@ def __init__(
                 system_prompt=sql_generation_reasoning_system_prompt,
                 streaming_callback=self._streaming_callback,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=sql_generation_reasoning_user_prompt_template
             ),

wren-ai-service/src/pipelines/generation/intent_classification.py

@@ -15,6 +15,7 @@
 from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
 from src.pipelines.common import build_table_ddl
 from src.pipelines.generation.utils.sql import construct_instructions
+from src.utils import trace_cost
 from src.web.v1.services import Configuration
 from src.web.v1.services.ask import AskHistory
 
@@ -280,8 +281,9 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def classify_intent(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@trace_cost
+async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -348,6 +350,7 @@ def __init__(
                 system_prompt=intent_classification_system_prompt,
                 generation_kwargs=INTENT_CLASSIFICAION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=intent_classification_user_prompt_template
             ),

wren-ai-service/src/pipelines/generation/misleading_assistance.py

@@ -10,6 +10,7 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 from src.web.v1.services.ask import AskHistory
 
 logger = logging.getLogger("wren-ai-service")
@@ -68,11 +69,14 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def misleading_assistance(prompt: dict, generator: Any, query_id: str) -> dict:
+@trace_cost
+async def misleading_assistance(
+    prompt: dict, generator: Any, query_id: str, generator_name: str
+) -> dict:
     return await generator(
         prompt=prompt.get("prompt"),
         query_id=query_id,
-    )
+    ), generator_name
 
 
 ## End of Pipeline
@@ -90,6 +94,7 @@ def __init__(
                 system_prompt=misleading_assistance_system_prompt,
                 streaming_callback=self._streaming_callback,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=misleading_assistance_user_prompt_template
             ),

wren-ai-service/src/pipelines/generation/question_recommendation.py

@@ -11,6 +11,7 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -40,9 +41,10 @@ def prompt(
     )
 
 
-@observe(capture_input=False, as_type="generation")
-async def generate(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@observe(as_type="generation", capture_input=False)
+@trace_cost
+async def generate(prompt: dict, generator: Any, generator_name: str) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -235,6 +237,7 @@ def __init__(
                 system_prompt=system_prompt,
                 generation_kwargs=QUESTION_RECOMMENDATION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
         }
 
         self._final = "normalized"

wren-ai-service/src/pipelines/generation/relationship_recommendation.py

@@ -13,6 +13,7 @@
 from src.core.engine import Engine
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -54,8 +55,9 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@trace_cost
+async def generate(prompt: dict, generator: Any, generator_name: str) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -194,6 +196,7 @@ def __init__(
                 system_prompt=system_prompt,
                 generation_kwargs=RELATIONSHIP_RECOMMENDATION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "engine": engine,
         }
 

wren-ai-service/src/pipelines/generation/semantics_description.py

@@ -11,6 +11,7 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -65,8 +66,9 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate(prompt: dict, generator: Any) -> dict:
-    return await generator(prompt=prompt.get("prompt"))
+@trace_cost
+async def generate(prompt: dict, generator: Any, generator_name: str) -> dict:
+    return await generator(prompt=prompt.get("prompt")), generator_name
 
 
 @observe(capture_input=False)
@@ -214,6 +216,7 @@ def __init__(self, llm_provider: LLMProvider, **_):
                 system_prompt=system_prompt,
                 generation_kwargs=SEMANTICS_DESCRIPTION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
         }
         self._final = "output"
 

wren-ai-service/src/pipelines/generation/sql_answer.py

@@ -10,6 +10,7 @@
 
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -63,8 +64,13 @@ def prompt(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate_answer(prompt: dict, generator: Any, query_id: str) -> dict:
-    return await generator(prompt=prompt.get("prompt"), query_id=query_id)
+@trace_cost
+async def generate_answer(
+    prompt: dict, generator: Any, query_id: str, generator_name: str
+) -> dict:
+    return await generator(
+        prompt=prompt.get("prompt"), query_id=query_id
+    ), generator_name
 
 
 ## End of Pipeline
@@ -85,6 +91,7 @@ def __init__(
                 system_prompt=sql_to_answer_system_prompt,
                 streaming_callback=self._streaming_callback,
             ),
+            "generator_name": llm_provider.get_model(),
         }
 
         super().__init__(

wren-ai-service/src/pipelines/generation/sql_correction.py

@@ -17,6 +17,7 @@
     TEXT_TO_SQL_RULES,
     SQLGenPostProcessor,
 )
+from src.utils import trace_cost
 
 logger = logging.getLogger("wren-ai-service")
 
@@ -70,13 +71,16 @@ def prompts(
 
 
 @observe(as_type="generation", capture_input=False)
-async def generate_sql_corrections(prompts: list[dict], generator: Any) -> list[dict]:
+@trace_cost
+async def generate_sql_corrections(
+    prompts: list[dict], generator: Any, generator_name: str
+) -> list[dict]:
     tasks = []
     for prompt in prompts:
         task = asyncio.ensure_future(generator(prompt=prompt.get("prompt")))
         tasks.append(task)
 
-    return await asyncio.gather(*tasks)
+    return await asyncio.gather(*tasks), generator_name
 
 
 @observe(capture_input=False)
@@ -109,6 +113,7 @@ def __init__(
                 system_prompt=sql_correction_system_prompt,
                 generation_kwargs=SQL_GENERATION_MODEL_KWARGS,
             ),
+            "generator_name": llm_provider.get_model(),
             "prompt_builder": PromptBuilder(
                 template=sql_correction_user_prompt_template
             ),