inference working

packages/node/meshnet_node/torch_server.py

@@ -11,6 +11,7 @@ import urllib.error
 import urllib.parse
 import urllib.request
 import uuid
+from typing import Any
 
 from .model_backend import (
     InsufficientVRAMError,
@@ -232,7 +233,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                     )
                 else:
                     text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
-                    self._send_openai_response(text, model_name, False)
+                    self._send_openai_response(text, model_name, False, messages)
             except Exception as exc:
                 self._send_json(500, {"error": f"generation failed: {exc}"})
             return
@@ -250,7 +251,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             return
         remaining_route = self._get_remaining_route(model_name)
         result_text = self._run_downstream_pipeline(payload, remaining_route)
-        self._send_openai_response(result_text, model_name, stream)
+        self._send_openai_response(result_text, model_name, stream, messages)
 
     def _get_remaining_route(self, model: str) -> list[str]:
         server: _TorchHTTPServer = self.server  # type: ignore[assignment]
@@ -367,10 +368,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
         self.wfile.write(b"data: [DONE]\n\n")
         self.wfile.flush()
 
-    def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
+    def _send_openai_response(
+        self,
+        text: str,
+        model: str,
+        stream: bool,
+        messages: list[dict] | None = None,
+    ) -> None:
         chunk_id = "chatcmpl-node"
         created = int(time.time())
         if not stream:
+            usage = _usage_for_response(self.server.backend, messages or [], text)  # type: ignore[attr-defined]
             self._send_json(200, {
                 "id": chunk_id,
                 "object": "chat.completion",
@@ -381,7 +389,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                     "message": {"role": "assistant", "content": text},
                     "finish_reason": "stop",
                 }],
-                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+                "usage": usage,
             })
             return
         self.send_response(200)
@@ -412,6 +420,52 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
         self.wfile.flush()
 
 
+def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
+    prompt_tokens = _backend_token_count(
+        backend,
+        "count_prompt_tokens",
+        messages,
+        fallback=_fallback_message_token_count(messages),
+    )
+    completion_tokens = _backend_token_count(
+        backend,
+        "count_text_tokens",
+        completion_text,
+        fallback=_fallback_text_token_count(completion_text),
+    )
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": completion_tokens,
+        "total_tokens": prompt_tokens + completion_tokens,
+    }
+
+
+def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
+    method: Any = getattr(backend, method_name, None)
+    if callable(method):
+        try:
+            return max(0, int(method(value)))
+        except Exception:
+            pass
+    return max(0, int(fallback))
+
+
+def _fallback_message_token_count(messages: list[dict]) -> int:
+    text = " ".join(
+        str(message.get("content", ""))
+        for message in messages
+        if isinstance(message, dict)
+    )
+    return _fallback_text_token_count(text)
+
+
+def _fallback_text_token_count(text: str) -> int:
+    parts = text.split()
+    if parts:
+        return len(parts)
+    return 1 if text else 0
+
+
 class TorchNodeServer:
     """HTTP server backed by a HuggingFace causal language model shard."""