""" title: Weave (Wandb) Context and Completion Logging description: Weave (by Weights and Biases) context and completion logging for OpenWebUI author: m-rgba (@martinmark) author_url: https://github.com/m-rgba/openwebui-weave version: 0.1 """ import datetime import tiktoken import requests from pydantic import BaseModel, Field from typing import List, Optional class Filter: class Valves(BaseModel): priority: int = Field( default=0, description="Priority level for the filter operations." ) wandb_project_name: Optional[str] = Field( default=None, description="Weights & Biases project name for initialization in the format of `username/project_name`.", ) wandb_api_key: Optional[str] = Field( default=None, description="Weights & Biases API key for login." ) pass def __init__(self): self.valves = self.Valves() pass def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict: # print(f"inlet:{__name__}") # print(f"inlet:body:{body}") # print(f"inlet:user:{__user__}") url = "https://trace.wandb.ai/call/start" headers = {"content-type": "application/json"} payload = { "start": { "project_id": self.valves.wandb_project_name, "op_name": __name__, "started_at": datetime.datetime.now().isoformat(), "inputs": { "messages": body.get("messages", []), "model": body.get("model", ""), "metadata": body.get("metadata", {}), "user": __user__, }, "attributes": {}, } } response = requests.post( url, headers=headers, json=payload, auth=("api", self.valves.wandb_api_key) ) if response.status_code == 200: # print("Start request successful") data = response.json() self.trace_id = data.get("id") else: print("Start request failed with status:", response.status_code) print(response.text) return body def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict: # print(f"outlet:{__name__}") # print(f"outlet:body:{body}") # print(f"outlet:user:{__user__}") # Extract the last assistant message and its metadata messages = body.get("messages", []) last_assistant_message_obj = next( (msg for msg in reversed(messages) if msg.get("role") == "assistant"), None, ) last_assistant_message = ( last_assistant_message_obj.get("content") if last_assistant_message_obj else None ) # Check if usage data is available in the last assistant message usage_data = ( last_assistant_message_obj.get("info", {}).get("usage") if last_assistant_message_obj else None ) model = body.get("model", "gpt-4o") if usage_data: usage_source = "API" # Use the provided usage data input_tokens = usage_data.get("prompt_tokens", 0) output_tokens = usage_data.get("completion_tokens", 0) else: usage_source = "Calculated" # Initialize tiktoken encoder and token buffers based on model # Approximate buffer from tiktoken as described here: # https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken#6-counting-tokens-for-chat-completions-api-calls try: encoding = tiktoken.encoding_for_model(model) except KeyError: encoding = tiktoken.get_encoding("cl100k_base") # Calculate input tokens from all messages input_tokens = 0 # print(messages) for message in messages: # Exclude last assistant response for input tokens if message is not last_assistant_message_obj: message_content = message.get("content", "") input_tokens += len(encoding.encode(str(message_content))) input_tokens = input_tokens + 4 # Calculate output tokens output_tokens = 0 if last_assistant_message: output_tokens = len(encoding.encode(last_assistant_message)) output_tokens = output_tokens + 3 # Build the payload payload = { "end": { "project_id": self.valves.wandb_project_name, "id": self.trace_id, "ended_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), "output": { "choices": [ {"message": {"content": last_assistant_message}}, ], "usage_source": usage_source, }, "summary": { "usage": { model: { "prompt_tokens": input_tokens, "completion_tokens": output_tokens, "total_tokens": input_tokens + output_tokens, "requests": 1, } } }, } } # Send end request url = "https://trace.wandb.ai/call/end" headers = {"content-type": "application/json"} response = requests.post( url, headers=headers, json=payload, auth=("api", self.valves.wandb_api_key) ) if response.status_code == 200: # print("End request successful") data = response.json() # print(f"outlet:data:{data}") else: print("End request failed with status:", response.status_code) print(response.text) return body