From 4b4cba893bbdfd971fd3f98feef53bd26517c4b9 Mon Sep 17 00:00:00 2001 From: Karen Shaw Date: Thu, 29 Aug 2024 19:07:41 +0000 Subject: [PATCH] Log dev team flag to metrics --- chat/src/event_config.py | 4 ++++ chat/src/handlers/chat.py | 6 +++--- chat/src/handlers/chat_sync.py | 2 -- chat/src/helpers/apitoken.py | 4 ++++ chat/src/helpers/metrics.py | 1 + chat/test/fixtures/apitoken.py | 3 +++ chat/test/helpers/test_apitoken.py | 10 ++++++++-- node/src/api/api-token.js | 7 ++++++- node/src/environment.js | 5 +++++ template.yaml | 4 ++++ 10 files changed, 38 insertions(+), 8 deletions(-) diff --git a/chat/src/event_config.py b/chat/src/event_config.py index 3b1e8ae7..28e09348 100644 --- a/chat/src/event_config.py +++ b/chat/src/event_config.py @@ -46,7 +46,9 @@ class EventConfig: deployment_name: str = field(init=False) document_prompt: ChatPromptTemplate = field(init=False) event: dict = field(default_factory=dict) + is_dev_team: bool = field(init=False) is_logged_in: bool = field(init=False) + is_superuser: bool = field(init=False) k: int = field(init=False) max_tokens: int = field(init=False) openai_api_version: str = field(init=False) @@ -70,7 +72,9 @@ def __post_init__(self): self.azure_resource_name = self._get_azure_resource_name() self.debug_mode = self._is_debug_mode_enabled() self.deployment_name = self._get_deployment_name() + self.is_dev_team = self.api_token.is_dev_team() self.is_logged_in = self.api_token.is_logged_in() + self.is_superuser = self.api_token.is_superuser() self.k = self._get_k() self.max_tokens = min(self.payload.get("max_tokens", MAX_TOKENS), MAX_TOKENS) self.openai_api_version = self._get_openai_api_version() diff --git a/chat/src/handlers/chat.py b/chat/src/handlers/chat.py index 3630bfab..b1c1f0e2 100644 --- a/chat/src/handlers/chat.py +++ b/chat/src/handlers/chat.py @@ -12,8 +12,8 @@ RESPONSE_TYPES = { "base": ["answer", "ref"], - "debug": ["answer", "attributes", "azure_endpoint", "deployment_name", "is_superuser", "k", "openai_api_version", "prompt", "question", "ref", "temperature", "text_key", "token_counts"], - "log": ["answer", "deployment_name", "is_superuser", "k", "openai_api_version", "prompt", "question", "ref", "size", "source_documents", "temperature", "token_counts"], + "debug": ["answer", "attributes", "azure_endpoint", "deployment_name", "is_dev_team", "is_superuser", "k", "openai_api_version", "prompt", "question", "ref", "temperature", "text_key", "token_counts"], + "log": ["answer", "deployment_name", "is_superuser", "k", "openai_api_version", "prompt", "question", "ref", "size", "source_documents", "temperature", "token_counts", "is_dev_team"], "error": ["question", "error", "source_documents"] } @@ -22,7 +22,7 @@ def handler(event, context): socket = event.get('socket', None) config.setup_websocket(socket) - if not config.is_logged_in: + if not (config.is_logged_in or config.is_superuser): config.socket.send({"type": "error", "message": "Unauthorized"}) return {"statusCode": 401, "body": "Unauthorized"} diff --git a/chat/src/handlers/chat_sync.py b/chat/src/handlers/chat_sync.py index 8166870e..98c512ce 100644 --- a/chat/src/handlers/chat_sync.py +++ b/chat/src/handlers/chat_sync.py @@ -17,8 +17,6 @@ } def handler(event, context): - print(f'Event: {event}') - config = HTTPEventConfig(event) if not config.is_logged_in: diff --git a/chat/src/helpers/apitoken.py b/chat/src/helpers/apitoken.py index 46c97263..0cd03e29 100644 --- a/chat/src/helpers/apitoken.py +++ b/chat/src/helpers/apitoken.py @@ -13,6 +13,7 @@ def empty_token(cls): "iat": time, "entitlements": [], "isLoggedIn": False, + "isDevTeam": False, } def __init__(self, signed_token=None): @@ -33,3 +34,6 @@ def is_logged_in(self): def is_superuser(self): return self.token.get("isSuperUser", False) + + def is_dev_team(self): + return self.token.get("isDevTeam", False) diff --git a/chat/src/helpers/metrics.py b/chat/src/helpers/metrics.py index 9610eac7..f00abc00 100644 --- a/chat/src/helpers/metrics.py +++ b/chat/src/helpers/metrics.py @@ -8,6 +8,7 @@ def debug_response(config, response, original_question): "attributes": config.attributes, "azure_endpoint": config.azure_endpoint, "deployment_name": config.deployment_name, + "is_dev_team": config.api_token.is_dev_team(), "is_superuser": config.api_token.is_superuser(), "k": config.k, "openai_api_version": config.openai_api_version, diff --git a/chat/test/fixtures/apitoken.py b/chat/test/fixtures/apitoken.py index 08691856..5097a56c 100644 --- a/chat/test/fixtures/apitoken.py +++ b/chat/test/fixtures/apitoken.py @@ -6,3 +6,6 @@ TEST_TOKEN = ('eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjQ4NDM1ODY2MDYxNjUs' 'ImlhdCI6MTY4Nzg5MTM2OSwiZW50aXRsZW1lbnRzIjpbXSwiaXNMb2dnZWRJbiI6d' 'HJ1ZSwic3ViIjoidGVzdFVzZXIifQ.vIZag1pHE1YyrxsKKlakXX_44ckAvkg7xWOoA_w4x58') +DEV_TEAM_TOKEN = ('eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjI1MjY1OTQ2MDcsInN1YiI6ImFiYzEyMyIsImVudGl0bGVtZW50cyI6W10sImlhdCI6MTcyNDk1MTI1NiwiaXNMb2dnZWRJbiI6dHJ1ZSwiaXNTdXBlclVzZXIiOmZhbHNlLCJpc3MiOiJtZWFkb3ciLCJpc0RldlRlYW0iOnRydWV9.wLDnGywoLSYMPte05tO4vPOou_QY7g_RJTeYUN_8o3I' + 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjI1MjY1OTQ2MDcsInN1YiI6ImFiYzEyMyIsImVudGl0bGVtZW50cyI6W10sImlhdCI6MTcyNDk1MTg4MCwiaXNMb2dnZWRJbiI6dHJ1ZSwiaXNTdXBlclVzZXIiOmZhbHNlLCJpc3MiOiJtZWFkb3ciLCJpc0RldlRlYW0iOnRydWV9.enO4CjujymoE0ow0zQcDyoTdtSAT93Jn-svSk6d5Urk' + 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjI1MjY1OTQ2MDcsInN1YiI6ImFiYzEyMyIsImVudGl0bGVtZW50cyI6W10sImlhdCI6MTcyNDk1MTg5MiwiaXNMb2dnZWRJbiI6dHJ1ZSwiaXNTdXBlclVzZXIiOmZhbHNlLCJpc3MiOiJtZWFkb3ciLCJpc0RldlRlYW0iOnRydWV9.4zwXnGyFkSM4zhpnk0xwNlnKzxDMDdl1YyBXMMkHK1k') diff --git a/chat/test/helpers/test_apitoken.py b/chat/test/helpers/test_apitoken.py index a330f56a..e23a1646 100644 --- a/chat/test/helpers/test_apitoken.py +++ b/chat/test/helpers/test_apitoken.py @@ -1,15 +1,16 @@ # ruff: noqa: E402 import os import sys + sys.path.append('./src') from helpers.apitoken import ApiToken -from test.fixtures.apitoken import SUPER_TOKEN, TEST_SECRET, TEST_TOKEN +from test.fixtures.apitoken import DEV_TEAM_TOKEN, SUPER_TOKEN, TEST_SECRET, TEST_TOKEN from unittest import mock, TestCase - +@mock.patch.dict(os.environ, {"DEV_TEAM_NET_IDS": "abc123"}) @mock.patch.dict(os.environ, {"API_TOKEN_SECRET": TEST_SECRET}) class TestFunction(TestCase): def test_empty_token(self): @@ -29,6 +30,11 @@ def test_superuser_token(self): self.assertTrue(subject.is_logged_in()) self.assertTrue(subject.is_superuser()) + def test_devteam_token(self): + subject = ApiToken(DEV_TEAM_TOKEN) + self.assertIsInstance(subject, ApiToken) + self.assertTrue(subject.is_dev_team()) + def test_invalid_token(self): subject = ApiToken("INVALID_TOKEN") self.assertIsInstance(subject, ApiToken) diff --git a/node/src/api/api-token.js b/node/src/api/api-token.js index 13f3e1c0..1edb3385 100644 --- a/node/src/api/api-token.js +++ b/node/src/api/api-token.js @@ -1,4 +1,4 @@ -const { apiTokenSecret, dcApiEndpoint } = require("../environment"); +const { apiTokenSecret, dcApiEndpoint, devTeamNetIds } = require("../environment"); const jwt = require("jsonwebtoken"); function emptyToken() { @@ -34,6 +34,7 @@ class ApiToken { name: user?.displayName?.[0], email: user?.mail, isLoggedIn: !!user, + isDevTeam: !!user && user?.uid && devTeamNetIds().includes(user?.uid), }; return this.update(); @@ -101,6 +102,10 @@ class ApiToken { return this.token.entitlements.has(entitlement); } + isDevTeam() { + return this.token.isDevTeam; + } + isLoggedIn() { return this.token.isLoggedIn; } diff --git a/node/src/environment.js b/node/src/environment.js index 44958569..cf330b1c 100644 --- a/node/src/environment.js +++ b/node/src/environment.js @@ -40,6 +40,10 @@ function dcUrl() { return process.env.DC_URL; } +function devTeamNetIds() { + return process.env.DEV_TEAM_NET_IDS.split(","); +} + function openSearchEndpoint() { return process.env.OPENSEARCH_ENDPOINT; } @@ -61,6 +65,7 @@ module.exports = { appInfo, dcApiEndpoint, dcUrl, + devTeamNetIds, openSearchEndpoint, prefix, region, diff --git a/template.yaml b/template.yaml index 67fd7ad6..9801c329 100644 --- a/template.yaml +++ b/template.yaml @@ -26,6 +26,7 @@ Globals: API_TOKEN_SECRET: !Ref ApiTokenSecret DC_API_ENDPOINT: !Ref DcApiEndpoint DC_URL: !Ref DcUrl + DEV_TEAM_NET_IDS: !Ref DevTeamNetIds OPENSEARCH_ENDPOINT: !Ref OpenSearchEndpoint ENV_PREFIX: !Ref EnvironmentPrefix HONEYBADGER_API_KEY: !Ref HoneybadgerApiKey @@ -63,6 +64,9 @@ Parameters: DcUrl: Type: String Description: URL of Digital Collections website + DevTeamNetIds: + Type: String + Description: Northwestern NetIDs of the development team FfmpegLayer: Type: String Description: "FFMPEG Lambda Layer ARN"