Merge branch 'main' of github.com:OpenGradient/bluefin-agent

Author: Adam Balogh

.env.example

@@ -17,5 +17,6 @@ export FIREBASE_CLIENT_EMAIL=''
 export FIREBASE_CLIENT_ID=''
 export FIREBASE_CLIENT_X509_CERT_URL=''
 
-# To skip firebase token authentication on endpoints. DO NOT ENABLE IN PRODUCTION.
-export SKIP_FIREBASE_TOKEN_AUTH=false
+# Skip token authentication on endpoints.
+export SKIP_TOKEN_AUTH_HEADER=''
+export SKIP_TOKEN_AUTH_KEY=''

.github/workflows/deploy.yml

@@ -50,7 +50,7 @@ jobs:
 
       - name: AWS LINKS TO VIEW DEPLOYMENT
         run: |
-          echo "View deployment progress here - https://us-east-2.console.aws.amazon.com/ecs/v2/clusters/Agents/services/two-ligma-agent-api/tasks?region=us-east-2"
+          echo "View deployment progress here - https://us-east-2.console.aws.amazon.com/ecs/v2/clusters/Agents/services/two-ligma-api-service/tasks?region=us-east-2"
       # - name: Wait for deployment to complete
       #   run: |
       #     echo "Waiting for ECS deployment to complete..."

server/auth.py

@@ -4,7 +4,7 @@
 from typing import Callable, TypeVar, cast
 from flask import abort, request, jsonify, g
 from .logging import logger
-from .config import SKIP_FIREBASE_TOKEN_AUTH
+from .config import SKIP_TOKEN_AUTH_HEADER, SKIP_TOKEN_AUTH_KEY
 
 class FirebaseIDTokenData(BaseModel):
     uid: str
@@ -48,8 +48,10 @@ def protected_route():
     """
     @wraps(f)
     def decorated_function(*args, **kwargs):
-        if SKIP_FIREBASE_TOKEN_AUTH:
-            return f(*args, **kwargs)
+        if SKIP_TOKEN_AUTH_HEADER and SKIP_TOKEN_AUTH_KEY:
+            skip_auth_header = request.headers.get(SKIP_TOKEN_AUTH_HEADER)
+            if skip_auth_header and skip_auth_header == SKIP_TOKEN_AUTH_KEY:
+                return f(*args, **kwargs)
 
         auth_header = request.headers.get('Authorization')
         if not auth_header or not auth_header.startswith('Bearer '):

server/config.py

@@ -1,8 +1,8 @@
 import os
 import logging
-from typing import Literal
 
-SKIP_FIREBASE_TOKEN_AUTH = bool(os.getenv("SKIP_FIREBASE_TOKEN_AUTH", False))
+SKIP_TOKEN_AUTH_HEADER = os.getenv("SKIP_TOKEN_AUTH_HEADER")
+SKIP_TOKEN_AUTH_KEY = os.getenv("SKIP_TOKEN_AUTH_KEY")
 
 # See if we are running in subnet mode
 SUBNET_MODE = os.getenv("subnet_mode", "false").lower() == "true"

server/firebase.py

@@ -1,6 +1,7 @@
 import os
 import firebase_admin  # type: ignore[import-untyped]
 from firebase_admin import auth  # noqa: F401
+from dotenv import dotenv_values
 
 def validate_firebase_env_vars():
     """
@@ -24,6 +25,13 @@ def validate_firebase_env_vars():
     # Get and process the private key to handle newlines
     private_key = os.environ.get("FIREBASE_PRIVATE_KEY").replace("\\n", "\n")
 
+    # Sometimes, when loading multiline env variables from a .env file,
+    # Only the -----BEGIN PRIVATE KEY----- is loaded (\n causes issues).
+    # This handles this case.
+    if '\n' not in private_key:
+        config = dotenv_values()
+        private_key = config["FIREBASE_PRIVATE_KEY"]
+
     return (
         os.environ.get("FIREBASE_PROJECT_ID"),
         os.environ.get("FIREBASE_PRIVATE_KEY_ID"),

subnet/subnet_methods.py

@@ -41,9 +41,12 @@ def make_request(input_data: Dict[str, Any], endpoint: str) -> requests.Response
     )
 
 
+import time
+
+
 def subnet_evaluation(quant_query: QuantQuery, quant_response: QuantResponse) -> float:
     """
-    Evaluate the subnet miner query based on the provided QuantQuery and QuantResponse.
+    Evaluate the subnet miner query based on the provided QuantQuery and QuantResponse, with up to 3 retries on failure.
 
     Args:
         quant_query (QuantQuery): The query object containing the query string and metadata.
@@ -56,31 +59,39 @@ def subnet_evaluation(quant_query: QuantQuery, quant_response: QuantResponse) ->
     if evaluation_model is None:
         evaluation_model = create_evaluation_model()
 
-    try:
-        template = env.get_template("evaluation_prompt.txt")
-        prompt = template.render(
-            user_prompt=quant_query.query,
-            agent_answer="No response provided" if quant_response is None else quant_response.response)
-
-        # Format messages properly for ChatOpenAI
-        messages = [{"role": "user", "content": prompt}]
-        response = evaluation_model.invoke(messages)
-
-        # Parse the response
-        answer = response.content if hasattr(response, "content") else response["content"]
-
-        # Find ```json{{...}}``` in the answer
-        match = re.search(r"```json\s*({.*})\s*```", answer, re.DOTALL)
-        if not match:
-            logging.error(f"Could not find JSON in model response: {answer}")
-            return 0.0
-        json_str = match.group(1)
-        score = json.loads(json_str)["score"]
-        # Normalize the score to be between 0 and 1
-        return float(score) / 50
-    except Exception as e:
-        logging.error(f"subnet_evaluation error: {e}")
-        return 0.0
+    retries = 3
+    delay = 3.0
+    last_exception = None
+    for attempt in range(1, retries + 1):
+        try:
+            template = env.get_template("evaluation_prompt.txt")
+            prompt = template.render(
+                user_prompt=quant_query.query,
+                agent_answer="No response provided" if quant_response is None else quant_response.response)
+
+            # Format messages properly for ChatOpenAI
+            messages = [{"role": "user", "content": prompt}]
+            response = evaluation_model.invoke(messages)
+
+            # Parse the response
+            answer = response.content if hasattr(response, "content") else response["content"]
+
+            # Find ```json{{...}}``` in the answer
+            match = re.search(r"```json\s*({.*})\s*```", answer, re.DOTALL)
+            if not match:
+                logging.error(f"Could not find JSON in model response: {answer}")
+                return 0.0
+            json_str = match.group(1)
+            score = json.loads(json_str)["score"]
+            # Normalize the score to be between 0 and 1
+            return float(score) / 50
+        except Exception as e:
+            last_exception = e
+            logging.error(f"subnet_evaluation attempt {attempt} failed: {e}")
+            if attempt < retries:
+                time.sleep(delay)
+    logging.error(f"subnet_evaluation failed after {retries} attempts: {last_exception}")
+    return 0.0
 
 
 def subnet_query(quant_query: QuantQuery) -> QuantResponse:

testclient/client.py

@@ -36,7 +36,7 @@ def main():
         tokens = agent_output.get("tokens", [])
 
         # print results
-        print(f"Two-Ligma: {answer}")
+        print(f"BitQuant: {answer}")
         if pools:
             print(f"Pools: {pools}")
         if tokens: