ZYM/ACM-OpenWebUI-plugins
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

able to submit record

Browse Source
This commit is contained in:
ZhuangYumin
2025-06-26 18:45:01 +08:00
parent 834aacc49f
commit d66b060cba
1 changed files with 132 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
openrouter_provider.py
View File

@ -1,8 +1,8 @@
"""
title: OpenRouter
version: 0.1.2
version: 0.2.0
license: MIT
description: Adds support for OpenRouter, including citations and reasoning tokens
description: Adds support for OpenRouter, including citations, reasoning tokens, and API call reporting
author: rburmorrison
author_url: https://github.com/rburmorrison
"""
@ -10,6 +10,7 @@ author_url: https://github.com/rburmorrison
import re
import requests
import json
import time
from typing import List, Union, Generator, Iterator
from pydantic import BaseModel, Field
@ -60,6 +61,14 @@ class Pipe:
MODEL_PREFIX: str = Field(
default="", description="Optional prefix for model names in Open WebUI"
)
REPORT_API_URL: str = Field(
default="",
description="URL to report API",
)
REPORT_API_KEY: str = Field(
default="",
description="API key to report API",
)
def __init__(self):
self.type = "manifold" # Multiple models
@ -107,7 +116,82 @@ class Pipe:
print(f"Error fetching models: {e}")
return [{"id": "error", "name": f"Error: {str(e)}"}]
def pipe(self, body: dict) -> Union[str, Generator, Iterator]:
def _fetch_generation_details(self, generation_id: str) -> dict:
"""Fetch generation details from OpenRouter API"""
try:
cnt = 0
max_allowed = 10
while cnt < max_allowed:
cnt += 1
headers = {"Authorization": f"Bearer {self.valves.OPENROUTER_API_KEY}"}
response = requests.get(
"https://openrouter.ai/api/v1/generation",
headers=headers,
params={"id": generation_id},
timeout=30
)
if response.status_code == 200:
return response.json()
elif response.status_code == 404:
print(f"Generation not found, retrying... ({cnt}/{max_allowed})")
time.sleep(1)
continue
else:
print(f"Error fetching generation details: {response.status_code}")
return {}
except Exception as e:
print(f"Error fetching generation details: {e}")
return {}
def _report_api_call(self, generation_data: dict, user_email: str, model_id: str):
"""Report API call to upstream reporting service"""
if not self.valves.REPORT_API_URL or not self.valves.REPORT_API_KEY:
return
try:
data = generation_data.get("data", {})
# Extract required fields for reporting
timestamp = int(time.time())
input_tokens = data.get("tokens_prompt", 0)
output_tokens = data.get("tokens_completion", 0)
cost_usd = data.get("total_cost", 0.0)
# Prepare API call record
api_call_record = {
"timestamp": timestamp,
"model_id": model_id,
"user_email": user_email,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cost_usd": cost_usd
}
# Send to reporting API
headers = {
"Authorization": f"Bearer {self.valves.REPORT_API_KEY}",
"Content-Type": "application/json"
}
report_url = f"{self.valves.REPORT_API_URL.rstrip('/')}/api/record_api_call"
response = requests.post(
report_url,
headers=headers,
json=api_call_record,
timeout=30
)
if response.status_code == 200:
print(f"Successfully reported API call for user {user_email}")
else:
print(f"Failed to report API call: {response.status_code}")
except Exception as e:
print(f"Error reporting API call: {e}")
def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]:
"""Process the request and handle reasoning tokens if supported"""
# Clone the body for OpenRouter
payload = body.copy()
@ -115,6 +199,10 @@ class Pipe:
# Print incoming body for debugging
print(f"Original request body: {json.dumps(body)[:500]}...")
# Extract user email and model ID for reporting
user_email = __user__.get("email", "") if __user__ else ""
model_id = "test_model"
# Make sure the model ID is properly extracted from the pipe format
if "model" in payload and payload["model"] and "." in payload["model"]:
# Extract the model ID from the format like "openrouter.model-id"
@ -148,9 +236,9 @@ class Pipe:
try:
if body.get("stream", False):
return self.stream_response(url, headers, payload)
return self.stream_response(url, headers, payload, user_email, model_id)
else:
return self.non_stream_response(url, headers, payload)
return self.non_stream_response(url, headers, payload, user_email, model_id)
except requests.exceptions.RequestException as e:
print(f"Request failed: {e}")
return f"Error: Request failed: {e}"
@ -158,7 +246,7 @@ class Pipe:
print(f"Error in pipe method: {e}")
return f"Error: {e}"
def non_stream_response(self, url, headers, payload):
def non_stream_response(self, url, headers, payload, user_email, model_id):
"""Handle non-streaming responses and wrap reasoning in <think> tags if present"""
try:
print(
@ -190,6 +278,20 @@ class Pipe:
res = response.json()
print(f"OpenRouter response keys: {list(res.keys())}")
# Extract generation ID for reporting
generation_id = res.get("id", "")
if generation_id and user_email and model_id:
try:
generation_data = self._fetch_generation_details(generation_id)
if generation_data:
self._report_api_call(generation_data, user_email, model_id)
except Exception as e:
print(f"Error reporting API call: {e}")
return f"Error: {e}"
else:
print(f"No generation ID found for reporting")
return f"Error: No generation ID found for reporting"
# Check if we have choices in the response
if not res.get("choices") or len(res["choices"]) == 0:
return ""
@ -204,8 +306,8 @@ class Pipe:
content = message.get("content", "")
reasoning = message.get("reasoning", "")
print(f"Found reasoning: {bool(reasoning)} ({len(reasoning)} chars)")
print(f"Found content: {bool(content)} ({len(content)} chars)")
print(f"Found reasoning: {bool(reasoning)} ({len(reasoning) if reasoning is not None else 0} chars)")
print(f"Found content: {bool(content)} ({len(content) if content is not None else 0} chars)")
# If we have both reasoning and content
if reasoning and content:
@ -219,7 +321,7 @@ class Pipe:
print(f"Error in non_stream_response: {e}")
return f"Error: {e}"
def stream_response(self, url, headers, payload):
def stream_response(self, url, headers, payload, user_email, model_id):
"""Stream reasoning tokens in real-time with proper tag management"""
try:
response = requests.post(
@ -240,6 +342,7 @@ class Pipe:
# State tracking
in_reasoning_state = False # True if we've output the opening <think> tag
latest_citations = [] # The latest citations list
generation_id = "" # Track generation ID for reporting
# Process the response stream
for line in response.iter_lines():
@ -250,15 +353,35 @@ class Pipe:
if not line_text.startswith("data: "):
continue
elif line_text == "data: [DONE]":
# Handle citations at the end
if latest_citations:
citation_list = [f"1. {l}" for l in latest_citations]
citation_list_str = "\n".join(citation_list)
yield f"\n\n---\nCitations:\n{citation_list_str}"
# Handle reporting at the end
if generation_id and user_email and model_id:
try:
generation_data = self._fetch_generation_details(generation_id)
if generation_data:
self._report_api_call(generation_data, user_email, model_id)
except Exception as e:
print(f"Error reporting API call: {e}")
return f"Error: {e}"
yield "" ## trick
else:
print(f"No generation ID found for reporting")
return f"Error: No generation ID found for reporting"
continue
try:
chunk = json.loads(line_text[6:])
# Extract generation ID from any chunk that has it
if not generation_id and "id" in chunk:
generation_id = chunk["id"]
print(f"Extracted generation ID for reporting: {generation_id}")
if "choices" in chunk and chunk["choices"]:
choice = chunk["choices"][0]
citations = chunk.get("citations") or []