Merge pull request #277 from marcklingen/fix-langfuse-filter

fix: langfuse filter pipeline cost tracking
This commit is contained in:
Timothy Jaeryang Baek 2024-09-27 20:05:12 +02:00 committed by GitHub
commit c1dd8987d4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,23 +1,28 @@
""" """
title: Langfuse Filter Pipeline title: Langfuse Filter Pipeline
author: open-webui author: open-webui
date: 2024-05-30 date: 2024-09-27
version: 1.3 version: 1.4
license: MIT license: MIT
description: A filter pipeline that uses Langfuse. description: A filter pipeline that uses Langfuse.
requirements: langfuse requirements: langfuse
""" """
from typing import List, Optional from typing import List, Optional
from schemas import OpenAIChatMessage
import os import os
import uuid import uuid
from utils.pipelines.main import get_last_user_message, get_last_assistant_message from utils.pipelines.main import get_last_assistant_message
from pydantic import BaseModel from pydantic import BaseModel
from langfuse import Langfuse from langfuse import Langfuse
from langfuse.api.resources.commons.errors.unauthorized_error import UnauthorizedError from langfuse.api.resources.commons.errors.unauthorized_error import UnauthorizedError
def get_last_assistant_message_obj(messages: List[dict]) -> dict:
for message in reversed(messages):
if message["role"] == "assistant":
return message
return {}
class Pipeline: class Pipeline:
class Valves(BaseModel): class Valves(BaseModel):
@ -109,21 +114,28 @@ class Pipeline:
async def outlet(self, body: dict, user: Optional[dict] = None) -> dict: async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
print(f"outlet:{__name__}") print(f"outlet:{__name__}")
print(f"Received body: {body}")
if body["chat_id"] not in self.chat_generations: if body["chat_id"] not in self.chat_generations:
return body return body
generation = self.chat_generations[body["chat_id"]] generation = self.chat_generations[body["chat_id"]]
assistant_message = get_last_assistant_message(body["messages"]) assistant_message = get_last_assistant_message(body["messages"])
# Extract usage information
info = assistant_message.get("info", {}) # Extract usage information for models that support it
usage = None usage = None
if "prompt_tokens" in info and "completion_tokens" in info: assistant_message_obj = get_last_assistant_message_obj(body["messages"])
usage = { if assistant_message_obj:
"input": info["prompt_tokens"], info = assistant_message_obj.get("info", {})
"output": info["completion_tokens"], if isinstance(info, dict):
"unit": "TOKENS", input_tokens = info.get("prompt_eval_count") or info.get("prompt_tokens")
} output_tokens = info.get("eval_count") or info.get("completion_tokens")
if input_tokens is not None and output_tokens is not None:
usage = {
"input": input_tokens,
"output": output_tokens,
"unit": "TOKENS",
}
# Update generation # Update generation
generation.end( generation.end(