swf-testbed/ai_memory/record.py

0001 #!/usr/bin/env python3
0002 """
0003 Record AI dialogue exchanges to the swf-testbed database.
0004
0005 Called by Claude Code hooks (UserPromptSubmit, Stop) to persist
0006 dialogue for cross-session context.
0007
0008 Environment:
0009     SWF_DIALOGUE_TURNS: Must be set and > 0 to enable recording
0010     SWF_MONITOR_HTTP_URL: Monitor API URL (default: http://pandaserver02.sdcc.bnl.gov/swf-monitor)
0011 """
0012
0013 import json
0014 import os
0015 import sys
0016 import getpass
0017 from typing import Optional
0018
0019
0020 def get_turns_setting():
0021     # type: () -> int
0022     try:
0023         return int(os.getenv('SWF_DIALOGUE_TURNS', '0'))
0024     except ValueError:
0025         return 0
0026
0027
0028 def record_via_api(username, session_id, role, content,
0029                    namespace=None, project_path=None):
0030     # type: (str, str, str, str, str, str) -> bool
0031     import urllib.request
0032     import urllib.error
0033
0034     base_url = os.getenv('SWF_MONITOR_HTTP_URL', 'http://pandaserver02.sdcc.bnl.gov/swf-monitor')
0035     url = "{}/api/ai-memory/record/".format(base_url.rstrip('/'))
0036
0037     payload = {
0038         "username": username,
0039         "session_id": session_id,
0040         "role": role,
0041         "content": content,
0042         "namespace": namespace,
0043         "project_path": project_path,
0044     }
0045
0046     try:
0047         req = urllib.request.Request(
0048             url,
0049             data=json.dumps(payload).encode('utf-8'),
0050             headers={'Content-Type': 'application/json'},
0051             method='POST'
0052         )
0053         with urllib.request.urlopen(req, timeout=5) as resp:
0054             return resp.status == 200
0055     except (urllib.error.URLError, urllib.error.HTTPError, OSError):
0056         return False
0057
0058
0059 def extract_assistant_response(transcript_path):
0060     # type: (str) -> Optional[str]
0061     if not transcript_path or not os.path.exists(transcript_path):
0062         return None
0063
0064     try:
0065         with open(transcript_path, 'r') as f:
0066             lines = f.readlines()
0067
0068         for line in reversed(lines):
0069             if not line.strip():
0070                 continue
0071             try:
0072                 entry = json.loads(line)
0073                 if entry.get('role') == 'assistant':
0074                     content = entry.get('content', '')
0075                     if isinstance(content, list):
0076                         texts = [c.get('text', '') for c in content if c.get('type') == 'text']
0077                         return '\n'.join(texts)
0078                     return str(content)
0079             except json.JSONDecodeError:
0080                 continue
0081     except Exception:
0082         pass
0083
0084     return None
0085
0086
0087 def get_namespace(cwd):
0088     # type: (str) -> Optional[str]
0089     testbed_toml = os.path.join(cwd, 'workflows', 'testbed.toml')
0090     if not os.path.exists(testbed_toml):
0091         return None
0092     try:
0093         try:
0094             import tomllib
0095             with open(testbed_toml, 'rb') as f:
0096                 data = tomllib.load(f)
0097             return data.get('testbed', {}).get('namespace')
0098         except ImportError:
0099             with open(testbed_toml, 'r') as f:
0100                 for line in f:
0101                     line = line.strip()
0102                     if line.startswith('namespace'):
0103                         parts = line.split('=', 1)
0104                         if len(parts) == 2:
0105                             return parts[1].strip().strip('"').strip("'")
0106     except Exception:
0107         pass
0108     return None
0109
0110
0111 def main():
0112     turns = get_turns_setting()
0113     if turns <= 0:
0114         sys.exit(0)
0115
0116     try:
0117         hook_input = json.load(sys.stdin)
0118     except json.JSONDecodeError:
0119         sys.exit(0)
0120
0121     event = hook_input.get('hook_event_name', '')
0122     session_id = hook_input.get('session_id', 'unknown')
0123     cwd = hook_input.get('cwd', '')
0124     username = getpass.getuser()
0125     namespace = get_namespace(cwd)
0126
0127     if event == 'UserPromptSubmit':
0128         prompt = hook_input.get('prompt', '')
0129         if prompt:
0130             record_via_api(
0131                 username=username,
0132                 session_id=session_id,
0133                 role='user',
0134                 content=prompt,
0135                 namespace=namespace,
0136                 project_path=cwd,
0137             )
0138
0139     elif event == 'Stop':
0140         transcript_path = hook_input.get('transcript_path')
0141         response = extract_assistant_response(transcript_path)
0142         if response:
0143             record_via_api(
0144                 username=username,
0145                 session_id=session_id,
0146                 role='assistant',
0147                 content=response,
0148                 namespace=namespace,
0149                 project_path=cwd,
0150             )
0151
0152     sys.exit(0)
0153
0154
0155 if __name__ == '__main__':
0156     main()