Pipe: возврат в формате OpenAI chat.completion с usage для Token Tracker (v2.8.0)

openrouter_auto_router_v2.7.1.py

@@ -1,14 +1,15 @@
 """
 title: OpenRouter Auto Router
 author: OpenAI
-version: 2.7.1
+version: 2.8.0
-description: Автоматическая маршрутизация OpenRouter с поддержкой tools.
+description: Автоматическая маршрутизация OpenRouter с поддержкой tools и прокидыванием usage для Token Tracker.
 """
 
 import asyncio
 import inspect
 import json
 import re
+import time
 
 import httpx
 from pydantic import BaseModel, Field
@@ -315,6 +316,74 @@ class Pipe:
 
         return payload
 
+    def _accumulate_usage(self, total_usage, data):
+        """Суммирует usage по всем итерациям tool-calling."""
+        if not isinstance(data, dict):
+            return
+
+        usage = data.get("usage")
+        if not isinstance(usage, dict):
+            return
+
+        total_usage["prompt_tokens"] += int(usage.get("prompt_tokens", 0) or 0)
+        total_usage["completion_tokens"] += int(usage.get("completion_tokens", 0) or 0)
+        total_usage["total_tokens"] += int(usage.get("total_tokens", 0) or 0)
+
+        cost = usage.get("cost")
+        if isinstance(cost, (int, float)):
+            total_usage["cost"] += float(cost)
+
+        prompt_details = usage.get("prompt_tokens_details") or {}
+        if isinstance(prompt_details, dict):
+            total_usage["cached_tokens"] += int(
+                prompt_details.get("cached_tokens", 0) or 0
+            )
+            total_usage["cache_write_tokens"] += int(
+                prompt_details.get("cache_write_tokens", 0) or 0
+            )
+
+        completion_details = usage.get("completion_tokens_details") or {}
+        if isinstance(completion_details, dict):
+            total_usage["reasoning_tokens"] += int(
+                completion_details.get("reasoning_tokens", 0) or 0
+            )
+
+    def _build_usage_payload(self, total_usage):
+        """Собирает usage в формате OpenAI chat.completion для Token Tracker."""
+        return {
+            "prompt_tokens": total_usage["prompt_tokens"],
+            "completion_tokens": total_usage["completion_tokens"],
+            "total_tokens": total_usage["total_tokens"],
+            "cost": total_usage["cost"],
+            "prompt_tokens_details": {
+                "cached_tokens": total_usage["cached_tokens"],
+                "cache_write_tokens": total_usage["cache_write_tokens"],
+            },
+            "completion_tokens_details": {
+                "reasoning_tokens": total_usage["reasoning_tokens"],
+            },
+        }
+
+    def _build_router_content(self, answer, route, model):
+        """Собирает финальный текст ответа с префиксом маршрутизатора."""
+        if not self.valves.SHOW_ROUTER_INFO:
+            return answer
+
+        route_names = {
+            "default": "💰 Обычный запрос",
+            "coding": "👨‍💻 Код / Linux / Docker",
+            "vision": "🖼️ Изображение / скриншот",
+            "complex": "🧠 Сложная задача",
+        }
+        route_name = route_names.get(route, route)
+
+        return (
+            f"**🤖 Auto Router**  \n"
+            f"Маршрут: **{route_name}**  \n"
+            f"Модель: `{model}`\n\n"
+            f"{answer}"
+        )
+
     async def _run_tool_callable(self, callable_fn, arguments, name):
         """Вызывает sync/async callable и ограничивает результат."""
         try:
@@ -453,7 +522,10 @@ class Pipe:
         __user__: dict = None,
         __tools__: dict = None,
     ):
-        """Маршрутизирует запрос и выполняет цикл tool calling."""
+        """Маршрутизирует запрос и выполняет цикл tool calling.
+
+        Возвращает OpenAI-совместимый dict формата chat.completion,
+        чтобы Token Tracker Filter мог извлечь usage (prompt/completion/cost)."""
         if not isinstance(body, dict):
             return "❌ **Некорректный формат запроса.**"
 
@@ -478,6 +550,20 @@ class Pipe:
         final_content = ""
         stopped_reason = ""
 
+        # Аккумуляторы usage по всем итерациям tool-calling
+        total_usage = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+            "cost": 0.0,
+            "reasoning_tokens": 0,
+            "cached_tokens": 0,
+            "cache_write_tokens": 0,
+        }
+        response_id = f"auto-router-{int(time.time())}"
+        response_created = int(time.time())
+        response_model_id = model
+
         timeout = httpx.Timeout(self.valves.REQUEST_TIMEOUT)
 
         try:
@@ -488,6 +574,16 @@ class Pipe:
                     )
                     data = await self._request(payload, client)
 
+                    # Суммируем usage и обновляем метаданные ответа
+                    if isinstance(data, dict):
+                        self._accumulate_usage(total_usage, data)
+                        if isinstance(data.get("id"), str):
+                            response_id = data["id"]
+                        if isinstance(data.get("created"), (int, float)):
+                            response_created = int(data["created"])
+                        if isinstance(data.get("model"), str):
+                            response_model_id = data["model"]
+
                     try:
                         message = data["choices"][0]["message"]
                     except (KeyError, IndexError, TypeError):
@@ -577,20 +673,21 @@ class Pipe:
             trace_text = self._format_tool_trace(tool_trace)
             answer = f"{answer}\n\n---\n\n{trace_text}" if answer else trace_text
 
-        if not self.valves.SHOW_ROUTER_INFO:
+        content = self._build_router_content(answer, route, model)
-            return answer
 
-        route_names = {
+        # Возвращаем OpenAI chat.completion dict — Token Tracker Filter
-            "default": "💰 Обычный запрос",
+        # прочитает usage и покажет токены/стоимость/баланс.
-            "coding": "👨‍💻 Код / Linux / Docker",
+        return {
-            "vision": "🖼️ Изображение / скриншот",
+            "id": response_id,
-            "complex": "🧠 Сложная задача",
+            "object": "chat.completion",
-        }
+            "created": response_created,
-        route_name = route_names.get(route, route)
+            "model": response_model_id,
-
+            "choices": [
-        return (
+                {
-            f"**🤖 Auto Router**  \n"
+                    "index": 0,
-            f"Маршрут: **{route_name}**  \n"
+                    "message": {"role": "assistant", "content": content},
-            f"Модель: `{model}`\n\n"
+                    "finish_reason": "stop",
-            f"{answer}"
+                }
-        )
+            ],
+            "usage": self._build_usage_payload(total_usage),
+        }