Skip to content

Commit 9524f0b

Browse files
dbirmandbirman
committed
feat: logging for /chat, fix bug in upgrade endpoint
1 parent 87d1d62 commit 9524f0b

6 files changed

Lines changed: 425 additions & 15 deletions

File tree

‎README.md‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,12 @@ print(response.json())
6565

6666
`POST /chat` — Ask a natural-language question about the metadata store. The agent can query records, look up schema info, and summarize results.
6767

68+
Query parameters:
69+
70+
| Parameter | Type | Required | Description |
71+
|---|---|---|---|
72+
| `id` | string | no | Optional caller identifier logged with each request (e.g. a username or app name) |
73+
6874
Request body:
6975

7076
| Field | Type | Required | Description |
@@ -83,9 +89,26 @@ Response body:
8389

8490
Rate-limited per IP (default: 10/min, 200/day).
8591

92+
Every successful request is appended as a JSON Lines record to a daily log file in S3:
93+
`s3://aind-scratch-data/aind-metadata-viz-logs/chat_log_{YYYY-MM-DD}.json`
94+
95+
Each log record contains:
96+
97+
| Field | Description |
98+
|---|---|
99+
| `timestamp` | ISO-8601 UTC timestamp |
100+
| `requester_id` | Value of the `?id=` query parameter, or `null` |
101+
| `ip` | Client IP address |
102+
| `message` | User message |
103+
| `response` | Agent response |
104+
| `stop_reason` | Agent stop reason |
105+
| `iterations` | Number of agent reasoning steps |
106+
| `tool_call_count` | Number of tool calls made |
107+
86108
```python
87109
response = requests.post(
88110
"https://metadata-portal.allenneuraldynamics.org/chat",
111+
params={"id": "my-app"},
89112
json={"message": "How many SmartSPIM assets are there?"},
90113
)
91114
print(response.json()["response"])
@@ -247,6 +270,12 @@ print(response.json())
247270

248271
`POST /chat` — Ask a natural-language question about the metadata store. The agent can query records, look up schema info, and summarize results.
249272

273+
Query parameters:
274+
275+
| Parameter | Type | Required | Description |
276+
|---|---|---|---|
277+
| `id` | string | no | Optional caller identifier logged with each request (e.g. a username or app name) |
278+
250279
Request body:
251280

252281
| Field | Type | Required | Description |
@@ -265,9 +294,26 @@ Response body:
265294

266295
Rate-limited per IP (default: 10/min, 200/day).
267296

297+
Every successful request is appended as a JSON Lines record to a daily log file in S3:
298+
`s3://aind-scratch-data/aind-metadata-viz-logs/chat_log_{YYYY-MM-DD}.json`
299+
300+
Each log record contains:
301+
302+
| Field | Description |
303+
|---|---|
304+
| `timestamp` | ISO-8601 UTC timestamp |
305+
| `requester_id` | Value of the `?id=` query parameter, or `null` |
306+
| `ip` | Client IP address |
307+
| `message` | User message |
308+
| `response` | Agent response |
309+
| `stop_reason` | Agent stop reason |
310+
| `iterations` | Number of agent reasoning steps |
311+
| `tool_call_count` | Number of tool calls made |
312+
268313
```python
269314
response = requests.post(
270315
"https://metadata-portal.allenneuraldynamics.org/chat",
316+
params={"id": "my-app"},
271317
json={"message": "How many SmartSPIM assets are there?"},
272318
)
273319
print(response.json()["response"])

‎scripts/test_log_integration.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Integration test for chat endpoint logging to S3.
4+
5+
Sends a chat request with a known ?id= value, then reads the daily S3 log
6+
file and verifies that the record was written with the correct fields.
7+
8+
Usage:
9+
python test_log_integration.py --env local
10+
python test_log_integration.py --env prod
11+
12+
Requires:
13+
- A running server (local or prod).
14+
- AWS credentials with s3:GetObject / s3:PutObject on
15+
s3://aind-scratch-data/aind-metadata-viz-logs/*.
16+
"""
17+
18+
from __future__ import annotations
19+
20+
import json
21+
import sys
22+
import uuid
23+
from datetime import datetime, timezone
24+
25+
import boto3
26+
import requests
27+
28+
from test_config import parse_test_args
29+
30+
_S3_BUCKET = "aind-scratch-data"
31+
_S3_PREFIX = "aind-metadata-viz-logs"
32+
33+
34+
def check(label: str, condition: bool, detail: str = "") -> bool:
35+
status = "OK " if condition else "FAIL"
36+
print(f"[{status}] {label}" + (f": {detail}" if detail else ""))
37+
return condition
38+
39+
40+
def _log_key(date_str: str) -> str:
41+
return f"{_S3_PREFIX}/chat_log_{date_str}.json"
42+
43+
44+
def _read_log_lines(date_str: str) -> list[dict]:
45+
s3 = boto3.client("s3")
46+
try:
47+
resp = s3.get_object(Bucket=_S3_BUCKET, Key=_log_key(date_str))
48+
raw = resp["Body"].read().decode()
49+
except s3.exceptions.NoSuchKey:
50+
return []
51+
except Exception as exc:
52+
print(f"[WARN] Could not read S3 log: {exc}")
53+
return []
54+
records = []
55+
for line in raw.splitlines():
56+
line = line.strip()
57+
if line:
58+
try:
59+
records.append(json.loads(line))
60+
except json.JSONDecodeError:
61+
pass
62+
return records
63+
64+
65+
def run_log_tests(base_url: str) -> bool:
66+
all_passed = True
67+
chat_url = f"{base_url}/chat"
68+
requester_id = f"log-integration-test-{uuid.uuid4().hex[:8]}"
69+
today_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d")
70+
71+
print(f"\n--- Chat log integration tests (id={requester_id}) ---")
72+
73+
print("\n[Test 1] POST /chat with ?id= writes a log record to S3")
74+
resp = requests.post(
75+
chat_url,
76+
params={"id": requester_id},
77+
json={"message": "Say the word 'hello' and nothing else."},
78+
timeout=120,
79+
)
80+
passed = check("Status 200", resp.status_code == 200, str(resp.status_code))
81+
if not passed:
82+
print(f" Response: {resp.text[:200]}")
83+
all_passed = False
84+
return all_passed
85+
86+
agent_response = resp.json().get("response", "")
87+
passed &= check("response non-empty", bool(agent_response))
88+
all_passed &= passed
89+
90+
print("\n[Test 2] Log record exists in S3 with correct fields")
91+
records = _read_log_lines(today_utc)
92+
matching = [r for r in records if r.get("requester_id") == requester_id]
93+
passed = check(
94+
"at least one log record with requester_id",
95+
len(matching) >= 1,
96+
f"{len(matching)} found",
97+
)
98+
all_passed &= passed
99+
100+
if matching:
101+
record = matching[-1]
102+
all_passed &= check(
103+
"timestamp present",
104+
bool(record.get("timestamp")),
105+
record.get("timestamp", ""),
106+
)
107+
all_passed &= check(
108+
"message matches",
109+
record.get("message") == "Say the word 'hello' and nothing else.",
110+
)
111+
all_passed &= check(
112+
"response matches agent response",
113+
record.get("response") == agent_response,
114+
)
115+
all_passed &= check(
116+
"stop_reason present", bool(record.get("stop_reason"))
117+
)
118+
all_passed &= check(
119+
"iterations is int", isinstance(record.get("iterations"), int)
120+
)
121+
all_passed &= check(
122+
"tool_call_count is int",
123+
isinstance(record.get("tool_call_count"), int),
124+
)
125+
126+
print("\n[Test 3] POST /chat without ?id= logs requester_id as null")
127+
null_id_marker = f"null-id-test-{uuid.uuid4().hex[:8]}"
128+
resp2 = requests.post(
129+
chat_url,
130+
json={"message": f"Reply with the exact phrase: {null_id_marker}"},
131+
timeout=120,
132+
)
133+
passed = check("Status 200", resp2.status_code == 200, str(resp2.status_code))
134+
if passed:
135+
records2 = _read_log_lines(today_utc)
136+
null_id_records = [
137+
r
138+
for r in records2
139+
if r.get("requester_id") is None
140+
and null_id_marker in r.get("message", "")
141+
]
142+
all_passed &= check(
143+
"record with null requester_id written",
144+
len(null_id_records) >= 1,
145+
f"{len(null_id_records)} found",
146+
)
147+
all_passed &= passed
148+
149+
return all_passed
150+
151+
152+
def main() -> None:
153+
args = parse_test_args()
154+
base_url = (
155+
"https://metadata-portal.allenneuraldynamics-test.org"
156+
if args.env == "prod"
157+
else "http://localhost:5006"
158+
)
159+
160+
print(f"Testing chat logging against: {base_url}")
161+
print(f"S3 log bucket: s3://{_S3_BUCKET}/{_S3_PREFIX}/")
162+
print("=" * 80)
163+
164+
ok = run_log_tests(base_url)
165+
166+
print("\n" + "=" * 80)
167+
if ok:
168+
print("All log integration tests passed.")
169+
else:
170+
print("Some log integration tests failed.")
171+
sys.exit(1)
172+
173+
174+
if __name__ == "__main__":
175+
main()

‎src/aind_metadata_viz/chat/handlers.py‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44

55
import logging
66
import os
7+
from typing import Optional
78

8-
from fastapi import APIRouter, Request
9+
from fastapi import APIRouter, Query, Request
910
from fastapi.responses import JSONResponse
1011

1112
from .agent import result_to_dict, run_agent
13+
from .log import append_chat_log
1214
from .ratelimit import RateLimiter, client_ip
1315

1416
logger = logging.getLogger(__name__)
@@ -65,7 +67,10 @@ def _validate(data) -> tuple[str | None, dict | None]:
6567

6668

6769
@chat_router.post("/chat")
68-
async def chat_endpoint(request: Request):
70+
async def chat_endpoint(
71+
request: Request,
72+
id: Optional[str] = Query(default=None, description="Optional caller identifier"),
73+
):
6974
ip = client_ip(
7075
request.headers, request.client.host if request.client else None
7176
)
@@ -109,4 +114,13 @@ async def chat_endpoint(request: Request):
109114
result.iterations,
110115
len(result.tool_calls),
111116
)
117+
append_chat_log(
118+
message=payload["message"],
119+
response=result_to_dict(result).get("response", ""),
120+
stop_reason=result.stop_reason,
121+
iterations=result.iterations,
122+
tool_call_count=len(result.tool_calls),
123+
ip=ip,
124+
requester_id=id,
125+
)
112126
return JSONResponse(content=result_to_dict(result))

‎src/aind_metadata_viz/chat/log.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
"""S3-backed chat request/response logger.
2+
3+
Appends one JSON Lines record per chat request to a daily log file in S3:
4+
s3://aind-scratch-data/aind-metadata-viz-logs/chat_log_{YYYY-MM-DD}.json
5+
6+
Each record contains:
7+
timestamp – ISO-8601 UTC
8+
requester_id – optional caller-supplied identifier (from ?id= query param)
9+
ip – client IP (may be None)
10+
message – the incoming user message
11+
response – the agent's reply text
12+
stop_reason
13+
iterations
14+
tool_call_count
15+
"""
16+
17+
from __future__ import annotations
18+
19+
import json
20+
import logging
21+
from datetime import datetime, timezone
22+
from typing import Optional
23+
24+
import boto3
25+
from botocore.exceptions import ClientError
26+
27+
logger = logging.getLogger(__name__)
28+
29+
_S3_BUCKET = "aind-scratch-data"
30+
_S3_PREFIX = "aind-metadata-viz-logs"
31+
32+
33+
def _s3():
34+
return boto3.client("s3")
35+
36+
37+
def _log_key(date_str: str) -> str:
38+
return f"{_S3_PREFIX}/chat_log_{date_str}.json"
39+
40+
41+
def _today_utc() -> str:
42+
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
43+
44+
45+
def _get_existing(key: str) -> bytes:
46+
try:
47+
resp = _s3().get_object(Bucket=_S3_BUCKET, Key=key)
48+
body = resp["Body"].read()
49+
if not body.endswith(b"\n"):
50+
body += b"\n"
51+
return body
52+
except ClientError as exc:
53+
if exc.response["Error"]["Code"] in ("NoSuchKey", "404"):
54+
return b""
55+
raise
56+
57+
58+
def append_chat_log(
59+
*,
60+
message: str,
61+
response: str,
62+
stop_reason: str,
63+
iterations: int,
64+
tool_call_count: int,
65+
ip: Optional[str],
66+
requester_id: Optional[str],
67+
) -> None:
68+
"""Append a single chat record to today's S3 log file."""
69+
record = {
70+
"timestamp": datetime.now(timezone.utc).isoformat(),
71+
"requester_id": requester_id,
72+
"ip": ip,
73+
"message": message,
74+
"response": response,
75+
"stop_reason": stop_reason,
76+
"iterations": iterations,
77+
"tool_call_count": tool_call_count,
78+
}
79+
line = (json.dumps(record) + "\n").encode()
80+
81+
key = _log_key(_today_utc())
82+
try:
83+
existing = _get_existing(key)
84+
updated = existing + line
85+
_s3().put_object(
86+
Bucket=_S3_BUCKET,
87+
Key=key,
88+
Body=updated,
89+
ContentType="application/x-ndjson",
90+
)
91+
except Exception:
92+
logger.exception("Failed to write chat log to S3 key=%s", key)

0 commit comments

Comments
 (0)