From 832ee4323d4fe60a34738abbb09dbfb415828b83 Mon Sep 17 00:00:00 2001
From: Tymec <tymek1rt@hotmail.com>
Date: Wed, 12 Apr 2023 04:08:42 +0200
Subject: [PATCH] API: add endpoint for counting tokens (#1051)

---
 extensions/api/script.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/extensions/api/script.py b/extensions/api/script.py
index 44769027..458f60e7 100644
--- a/extensions/api/script.py
+++ b/extensions/api/script.py
@@ -81,6 +81,19 @@ class Handler(BaseHTTPRequestHandler):
                 }]
             })
             self.wfile.write(response.encode('utf-8'))
+        elif self.path == '/api/v1/token-count':
+            # Not compatible with KoboldAI api
+            self.send_response(200)
+            self.send_header('Content-Type', 'application/json')
+            self.end_headers()
+
+            tokens = encode(body['prompt'])[0]
+            response = json.dumps({
+                'results': [{
+                    'tokens': len(tokens)
+                }]
+            })
+            self.wfile.write(response.encode('utf-8'))
         else:
             self.send_error(404)