Skip to content

Commit

Permalink
Log dev team flag to metrics
Browse files Browse the repository at this point in the history
  • Loading branch information
kdid committed Aug 29, 2024
1 parent 30b9bbc commit 4b4cba8
Show file tree
Hide file tree
Showing 10 changed files with 38 additions and 8 deletions.
4 changes: 4 additions & 0 deletions chat/src/event_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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()
Expand Down
6 changes: 3 additions & 3 deletions chat/src/handlers/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
}

Expand All @@ -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"}

Expand Down
2 changes: 0 additions & 2 deletions chat/src/handlers/chat_sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,6 @@
}

def handler(event, context):
print(f'Event: {event}')

config = HTTPEventConfig(event)

if not config.is_logged_in:
Expand Down
4 changes: 4 additions & 0 deletions chat/src/helpers/apitoken.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ def empty_token(cls):
"iat": time,
"entitlements": [],
"isLoggedIn": False,
"isDevTeam": False,
}

def __init__(self, signed_token=None):
Expand All @@ -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)
1 change: 1 addition & 0 deletions chat/src/helpers/metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions chat/test/fixtures/apitoken.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
10 changes: 8 additions & 2 deletions chat/test/helpers/test_apitoken.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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)
Expand Down
7 changes: 6 additions & 1 deletion node/src/api/api-token.js
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
const { apiTokenSecret, dcApiEndpoint } = require("../environment");
const { apiTokenSecret, dcApiEndpoint, devTeamNetIds } = require("../environment");
const jwt = require("jsonwebtoken");

function emptyToken() {
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -101,6 +102,10 @@ class ApiToken {
return this.token.entitlements.has(entitlement);
}

isDevTeam() {
return this.token.isDevTeam;
}

isLoggedIn() {
return this.token.isLoggedIn;
}
Expand Down
5 changes: 5 additions & 0 deletions node/src/environment.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -61,6 +65,7 @@ module.exports = {
appInfo,
dcApiEndpoint,
dcUrl,
devTeamNetIds,
openSearchEndpoint,
prefix,
region,
Expand Down
4 changes: 4 additions & 0 deletions template.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand Down

0 comments on commit 4b4cba8

Please sign in to comment.