cursor support

Author: acer-king

cursor/app/core/middleware.py

@@ -1,7 +1,29 @@
-import time
-from fastapi import HTTPException
+from fastapi import FastAPI, Depends, HTTPException, Request
+from cursor.app.core.config import config
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse
+# Your predefined valid API keys
+VALID_API_KEYS = {config.api_key}
 
 
-async def verify_api_key_rate_limit(config, api_key):
-    # NOTE: abit dangerous but very useful
-    pass
+class APIKeyMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        # Get the API key from the `Authorization` header
+        if request.method == "OPTIONS":
+            return await call_next(request)
+
+        if not request.headers.get("Authorization"):
+            return JSONResponse(
+                {"detail": "Invalid or missing API Key"}, status_code=401
+            )
+
+        api_key = request.headers.get("Authorization").split(" ")[1]
+
+        # Validate the API key
+        if not api_key or api_key not in VALID_API_KEYS:
+            return JSONResponse(
+                {"detail": "Invalid or missing API Key"}, status_code=401
+            )
+
+        # Proceed to the next middleware or route handler
+        return await call_next(request)

validators/weight_setter.py

@@ -10,6 +10,8 @@
 from substrateinterface import SubstrateInterface
 from functools import partial
 from typing import Tuple
+
+from fastapi import HTTPException
 import bittensor as bt
 from bittensor import StreamingSynapse
 
@@ -24,7 +26,10 @@
 from validators.task_manager import TaskMgr
 from cortext.dendrite import CortexDendrite
 from cortext.axon import CortexAxon
-from fastapi import HTTPException
+from cursor.app.endpoints.text import chat
+from cursor.app.endpoints.generic import models
+from cursor.app.core.middleware import APIKeyMiddleware
+
 
 scoring_organic_timeout = 60
 NUM_INTERVALS_PER_CYCLE = 10
@@ -104,7 +109,6 @@ def __init__(self, config, cache: QueryResponseCache, loop=None):
         # organic_thread.start()
         self.loop.create_task(self.consume_organic_queries())
 
-
     def start_axon_server(self):
         asyncio.run(self.consume_organic_queries())
 
@@ -557,11 +561,23 @@ async def consume_organic_queries(self):
             forward_fn=self.embeddings,
             blacklist_fn=self.blacklist_embeddings,
         )
+        self.cursor_setup()
         self.axon.serve(netuid=self.netuid, subtensor=self.subtensor)
         print(f"axon: {self.axon}")
         self.axon.start()
         bt.logging.info(f"Running validator on uid: {self.my_uid}")
 
+    def cursor_setup(self):
+        self.axon.router.add_api_route(
+            "/v1/chat/completions",
+            chat,
+            methods=["POST", "OPTIONS"],
+            tags=["StreamPrompting"],
+            response_model=None
+        )
+        self.axon.router.add_api_route("/v1/models", models, methods=["GET"], tags=["Text"], response_model=None)
+        self.axon.app.add_middleware(APIKeyMiddleware)
+
     def get_scoring_tasks_from_query_responses(self, queries_to_process):
 
         grouped_query_resps = defaultdict(list)
@@ -617,7 +633,6 @@ async def process_queries_from_database(self):
             queries_to_process = self.query_database.copy()
             self.query_database.clear()
 
-
             self.synthetic_task_done = False
             bt.logging.info("start scoring process")
             start_time = time.time()