Move metrics logging to the MetricsCallbackHandler

mbklein · mbklein · commit 031a46832bbe · 2025-03-21T21:01:57.000Z
diff --git a/chat/src/agent/callbacks/metrics.py b/chat/src/agent/callbacks/metrics.py
@@ -1,51 +1,109 @@
+from datetime import datetime
 from typing import Any, Dict
 from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.outputs import LLMResult
 from langchain_core.messages.tool import ToolMessage
+import boto3
 import json
+import os
 
 class MetricsCallbackHandler(BaseCallbackHandler):
-  def __init__(self, *args, **kwargs):
-    self.accumulator = {}
-    self.answers = []
-    self.artifacts = []
-    super().__init__(*args, **kwargs)
-
-  def on_llm_end(self, response: LLMResult, **kwargs: Dict[str, Any]):
-    if response is None:
-        return
-
-    if not response.generations or not response.generations[0]:
-        return
-        
-    for generation in response.generations[0]:
-      if generation.text != "":
-        self.answers.append(generation.text)
-
-      if not hasattr(generation, 'message') or generation.message is None:
-          continue
-          
-      metadata = getattr(generation.message, 'usage_metadata', None)
-      if metadata is None:
-          continue
-          
-      for k, v in metadata.items():
-          self.accumulator[k] = self.accumulator.get(k, 0) + v
-
-  def on_tool_end(self, output: ToolMessage, **kwargs: Dict[str, Any]):
+    def __init__(self, log_stream = None, *args, extra_data = {}, **kwargs):
+        self.accumulator = {}
+        self.answers = []
+        self.artifacts = []
+        self.log_stream = log_stream
+        self.extra_data = extra_data
+        super().__init__(*args, **kwargs)
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Dict[str, Any]):
+        if response is None:
+            return
+
+        if not response.generations or not response.generations[0]:
+            return
+
+        for generation in response.generations[0]:
+            if generation.text != "":
+                self.answers.append(generation.text)
+
+            if not hasattr(generation, "message") or generation.message is None:
+                continue
+
+            metadata = getattr(generation.message, "usage_metadata", None)
+            if metadata is None:
+                continue
+
+            for k, v in metadata.items():
+                self.accumulator[k] = self.accumulator.get(k, 0) + v
+
+    def on_tool_end(self, output: ToolMessage, **kwargs: Dict[str, Any]):
         content = output.content
         if isinstance(content, str):
             try:
                 content = json.loads(content)
             except json.decoder.JSONDecodeError as e:
-                print(f"Invalid json ({e}) returned from {output.name} tool: {output.content}")
+                print(
+                    f"Invalid json ({e}) returned from {output.name} tool: {output.content}"
+                )
                 return
-                
+
         match output.name:
             case "aggregate":
-                self.artifacts.append({"type": "aggregation", "artifact": content.get("aggregation_result", {})})
+                self.artifacts.append(
+                    {
+                        "type": "aggregation",
+                        "artifact": content.get("aggregation_result", {}),
+                    }
+                )
             case "search":
                 source_urls = [doc.get("api_link") for doc in content]
                 self.artifacts.append({"type": "source_urls", "artifact": source_urls})
             case "summarize":
                 print(output)
+
+    def log_metrics(self):
+        if self.log_stream is None:
+            return
+        
+        log_group = os.getenv("METRICS_LOG_GROUP")
+        if log_group and ensure_log_stream_exists(log_group, self.log_stream):
+            client = log_client()
+            message = {
+                "answer": self.answers,
+                "artifacts": self.artifacts,
+                "token_counts": self.accumulator,
+            }
+            message.update(self.extra_data)
+
+            log_events = [
+                {
+                    "timestamp": timestamp(),
+                    "message": json.dumps(message),
+                }
+            ]
+            client.put_log_events(
+                logGroupName=log_group, logStreamName=self.log_stream, logEvents=log_events
+            )
+
+
+def ensure_log_stream_exists(log_group, log_stream):
+    client = log_client()
+    try:
+        print(
+            client.create_log_stream(logGroupName=log_group, logStreamName=log_stream)
+        )
+        return True
+    except client.exceptions.ResourceAlreadyExistsException:
+        return True
+    except Exception:
+        print(f"Could not create log stream: {log_group}:{log_stream}")
+        return False
+
+
+def log_client():
+    return boto3.client("logs", region_name=os.getenv("AWS_REGION", "us-east-1"))
+
+
+def timestamp():
+    return round(datetime.timestamp(datetime.now()) * 1000)
diff --git a/chat/src/agent/search_agent.py b/chat/src/agent/search_agent.py
@@ -12,6 +12,7 @@
 from core.setup import checkpoint_saver
 from agent.callbacks.socket import SocketCallbackHandler
 from typing import Optional
+import time
 
 DEFAULT_SYSTEM_MESSAGE = """
 Please provide a brief answer to the question using the tools provided. Include specific details from multiple documents that 
@@ -23,7 +24,8 @@
 MAX_RECURSION_LIMIT = 16
 
 class SearchWorkflow:
-    def __init__(self, model: BaseModel, system_message: str):
+    def __init__(self, model: BaseModel, system_message: str, metrics = None):
+        self.metrics = metrics
         self.model = model
         self.summarization_model = ChatBedrock(model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", streaming=False)
         self.system_message = system_message
@@ -57,15 +59,25 @@ def summarize(self, state: MessagesState):
         It is extremely important that you return only the valid, parsable summarized 
         JSON with no additional text or explanation, no markdown code fencing, and all 
         unnecessary whitespace removed.
+        
+        Prioritize speed over comprehensiveness.
 
         {last_message.content}
         """
+
         config = {
-            "callbacks": [], 
+            "callbacks": [self.metrics] if self.metrics else [], 
             "metadata": {"source": "summarize"}
         }
+        
+        start_time = time.time()
+        
         summary = self.summarization_model.invoke([HumanMessage(content=summary_prompt)], config=config)
-        print(f'Condensed {len(last_message.content)} bytes to {len(summary.content)} bytes via summarization')
+        
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print(f'Condensed {len(last_message.content)} bytes to {len(summary.content)} bytes in {elapsed_time:.2f} seconds')
+        
         last_message.content = summary.content
 
         return {"messages": messages}
@@ -81,6 +93,7 @@ def __init__(
         self,
         model: BaseModel,
         *,
+        metrics = None,
         system_message: str = DEFAULT_SYSTEM_MESSAGE,
         **kwargs
     ):
@@ -92,7 +105,7 @@ def __init__(
         except NotImplementedError:
             pass
 
-        self.workflow_logic = SearchWorkflow(model=model, system_message=system_message)
+        self.workflow_logic = SearchWorkflow(model=model, system_message=system_message, metrics=metrics)
 
         # Define a new graph
         workflow = StateGraph(MessagesState)
diff --git a/chat/src/agent/tools.py b/chat/src/agent/tools.py
@@ -1,7 +1,5 @@
 import json
 
-from langchain_core.language_models.chat_models import BaseModel
-from langchain_core.messages import HumanMessage
 from langchain_core.tools import tool
 from core.setup import opensearch_vector_store
 from typing import List
@@ -29,7 +27,8 @@ def filter_results(results):
     Filters out the embeddings from the results
     """
     filtered = []
-    for doc in results:
+    for result in results:
+        doc = result.metadata
         if 'embedding' in doc:
             doc.pop('embedding')
         filtered.append(doc)
@@ -101,33 +100,3 @@ def retrieve_documents(doc_ids: List[str]):
         return filter_results(response)
     except Exception as e:
         return {"error": str(e)}
-    
-@tool(response_format="content")
-def summarize(content, model: BaseModel):
-    """
-    Summarize content. If content is a list of documents, each document will
-    be replaced with a summary to reduce the amount of content passed to the agent's
-    model at each turn. Otherwise, the content will be summarized as a whole.
-    
-    Args:
-        content: The content to summarize.
-        model (BaseModel): The summarization model to use.
-    
-    Returns:
-        A new list of documents, pared down.
-    """
-
-    summary_prompt = f"""
-    Summarize the following content. If the content is a list of documents
-    with IDs, replace each document with a new dict with the shape
-    {'id': doc.id, 'title': doc.title 'content': summary}, where summary is a 
-    concise but semantically meaningful summary of the document content for the
-    agent to use on subsequent turns. Otherwise, produce a summary of the content
-    as a whole.
-    
-    {content}
-    """
-    print(f"Summarizing content: {content}")
-    summary = model.invoke([HumanMessage(content=summary_prompt)])
-    print(f"Summarized content: {summary.content}")
-    return summary.content
diff --git a/chat/src/handlers.py b/chat/src/handlers.py
@@ -1,9 +1,6 @@
-import boto3
 import json
 import logging
-import os
 from core.secrets import load_secrets
-from datetime import datetime
 from core.event_config import EventConfig
 from honeybadger import honeybadger
 from agent.search_agent import SearchAgent
@@ -63,59 +60,26 @@ def chat(event, context):
         config.socket.send({"type": "error", "message": "Question cannot be blank"})
         return {"statusCode": 400, "body": "Question cannot be blank"}
 
-    metrics = MetricsCallbackHandler()
+    log_info = {
+        "is_dev_team": config.api_token.is_dev_team(),
+        "is_superuser": config.api_token.is_superuser(),
+        "k": config.k,
+        "model": config.model,
+        "question": config.question,
+        "ref": config.ref,
+    }
+    metrics = MetricsCallbackHandler(context.log_stream_name, extra_data=log_info)
     callbacks = [SocketCallbackHandler(config.socket, config.ref), metrics]
     model = chat_model(model=config.model, streaming=config.stream_response)
-    search_agent = SearchAgent(model=model)
-
+    search_agent = SearchAgent(model=model, metrics=metrics)
+    
     try:
         search_agent.invoke(config.question, config.ref, forget=config.forget, docs=config.docs, callbacks=callbacks)
-        log_metrics(context, metrics, config)
+        metrics.log_metrics()
     except Exception as e:
         error_response = {"type": "error", "message": "An unexpected error occurred. Please try again later."}
         if config.socket:
             config.socket.send(error_response)
         raise e
 
     return {"statusCode": 200}
-
-
-def log_metrics(context, metrics, config):
-    log_group = os.getenv("METRICS_LOG_GROUP")
-    log_stream = context.log_stream_name
-    if log_group and ensure_log_stream_exists(log_group, log_stream):
-        client = log_client()
-        log_events = [{
-            "timestamp": timestamp(), 
-            "message": json.dumps({
-                "answer": metrics.answers,
-                "is_dev_team": config.api_token.is_dev_team(),
-                "is_superuser": config.api_token.is_superuser(),
-                "k": config.k,
-                "model": config.model,
-                "question": config.question,
-                "ref": config.ref,
-                "artifacts": metrics.artifacts,
-                "token_counts": metrics.accumulator,
-            })
-        }]
-        client.put_log_events(
-                logGroupName=log_group, logStreamName=log_stream, logEvents=log_events
-            )
-    
-def ensure_log_stream_exists(log_group, log_stream):
-    client = log_client()
-    try:
-        print(client.create_log_stream(logGroupName=log_group, logStreamName=log_stream))
-        return True
-    except client.exceptions.ResourceAlreadyExistsException:
-        return True
-    except Exception:
-        print(f"Could not create log stream: {log_group}:{log_stream}")
-        return False
-
-def log_client():
-    return boto3.client("logs", region_name=os.getenv("AWS_REGION", "us-east-1"))
-
-def timestamp():
-    return round(datetime.timestamp(datetime.now()) * 1000)
diff --git a/chat/template.yaml b/chat/template.yaml
@@ -53,7 +53,6 @@ Resources:
   ChatWebSocket:
     Type: AWS::ApiGatewayV2::Api
     Properties:
-      Name: ChatWebSocket
       ProtocolType: WEBSOCKET
       RouteSelectionExpression: "$request.body.message"
   CheckpointBucket:
diff --git a/chat/test/agent/test_tools.py b/chat/test/agent/test_tools.py
@@ -33,11 +33,17 @@ def test_discover_fields(self, mock_opensearch):
 
     @patch('agent.tools.opensearch_vector_store')
     def test_search(self, mock_opensearch):
-        mock_results = [{"id": "doc1", "text": "example result"}]
+        class MockDoc:
+            def __init__(self, content, metadata):
+                self.content = content
+                self.metadata = metadata
+
+        expected_results = [{"id": "doc1", "text": "example result"}]
+        mock_results = [MockDoc(content=doc["id"], metadata=doc) for doc in expected_results]
         mock_opensearch.return_value.similarity_search.return_value = mock_results
         
         response = search.invoke("test query")
-        self.assertEqual(response, mock_results)
+        self.assertEqual(response, expected_results)
 
     @patch('agent.tools.opensearch_vector_store')
     def test_aggregate(self, mock_opensearch):
