diff --git a/chat/src/agent/agent_handler.py b/chat/src/agent/agent_handler.py new file mode 100644 index 0000000..91fb6d2 --- /dev/null +++ b/chat/src/agent/agent_handler.py @@ -0,0 +1,12 @@ +from typing import Any, Dict, List + +from langchain_core.callbacks import BaseCallbackHandler +from langchain_core.messages import BaseMessage +from langchain_core.outputs import LLMResult + + +class AgentHandler(BaseCallbackHandler): + def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any: + print(f"on_tool_start (A tool is starting): {serialized, input_str}") + +callbacks = [AgentHandler()] \ No newline at end of file diff --git a/chat/src/agent/tools.py b/chat/src/agent/tools.py index 4bfba66..a936045 100644 --- a/chat/src/agent/tools.py +++ b/chat/src/agent/tools.py @@ -6,7 +6,7 @@ @tool(response_format="content_and_artifact") def search(query: str): """Perform a semantic search of Northwestern University Library digital collections. When answering a search query, ground your answer in the context of the results with references to the document's metadata.""" - query_results = opensearch_vector_store.similarity_search(query, size=20) + query_results = opensearch_vector_store().similarity_search(query, size=20) return json.dumps(query_results, default=str), query_results @tool(response_format="content_and_artifact") @@ -29,7 +29,7 @@ def aggregate(aggregation_query: str): - Number of works by work type: work_type """ try: - response = opensearch_vector_store.aggregations_search(aggregation_query) + response = opensearch_vector_store().aggregations_search(aggregation_query) return json.dumps(response, default=str), response except Exception as e: return json.dumps({"error": str(e)}), None diff --git a/chat/src/handlers/chat.py b/chat/src/handlers/chat.py index 4250008..2442cff 100644 --- a/chat/src/handlers/chat.py +++ b/chat/src/handlers/chat.py @@ -8,6 +8,7 @@ from honeybadger import honeybadger from agent.search_agent import search_agent from langchain_core.messages import HumanMessage +from agent.agent_handler import callbacks honeybadger.configure() logging.getLogger("honeybadger").addHandler(logging.StreamHandler()) @@ -61,23 +62,20 @@ def handler(event, context): config.socket.send({"type": "error", "message": "Question cannot be blank"}) return {"statusCode": 400, "body": "Question cannot be blank"} + 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): + log_client = boto3.client("logs") + log_events = [{"timestamp": timestamp(), "message": "Hello world"}] + log_client.put_log_events( + logGroupName=log_group, logStreamName=log_stream, logEvents=log_events + ) + response = search_agent.invoke( {"messages": [HumanMessage(content=config.question)]}, - config={"configurable": {"thread_id": config.ref}}, + config={"configurable": {"thread_id": config.ref}, "callbacks": callbacks}, ) - 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): - log_client = boto3.client('logs') - log_events = [ - { - 'timestamp': timestamp(), - 'message': json.dumps(response) - } - ] - log_client.put_log_events(logGroupName=log_group, logStreamName=log_stream, logEvents=log_events) - return {"statusCode": 200} diff --git a/chat/template.yaml b/chat/template.yaml index 6001950..94e23c7 100644 --- a/chat/template.yaml +++ b/chat/template.yaml @@ -208,6 +208,7 @@ Resources: HONEYBADGER_REVISION: !Ref HoneybadgerRevision METRICS_LOG_GROUP: !Ref ChatMetricsLog SECRETS_PATH: !Ref SecretsPath + NO_COLOR: 1 Policies: - !Ref SecretsPolicy - Statement: @@ -230,37 +231,37 @@ Resources: Resource: !Sub "${ChatMetricsLog.Arn}:*" #* Metadata: #* BuildMethod: nodejs20.x - ChatSyncFunction: - Type: AWS::Serverless::Function - Properties: - CodeUri: ./src - Runtime: python3.12 - Architectures: - - x86_64 - #* Layers: - #* - !Ref ChatDependencies - MemorySize: 1024 - Handler: handlers/chat_sync.handler - Timeout: 300 - Environment: - Variables: - API_TOKEN_NAME: !Ref ApiTokenName - ENV_PREFIX: !Ref EnvironmentPrefix - HONEYBADGER_API_KEY: !Ref HoneybadgerApiKey - HONEYBADGER_ENVIRONMENT: !Ref HoneybadgerEnv - HONEYBADGER_REVISION: !Ref HoneybadgerRevision - METRICS_LOG_GROUP: !Ref ChatMetricsLog - SECRETS_PATH: !Ref SecretsPath - FunctionUrlConfig: - AuthType: NONE - Policies: - - !Ref SecretsPolicy - - Statement: - - Effect: Allow - Action: - - 'es:ESHttpGet' - - 'es:ESHttpPost' - Resource: '*' + # ChatSyncFunction: + # Type: AWS::Serverless::Function + # Properties: + # CodeUri: ./src + # Runtime: python3.12 + # Architectures: + # - x86_64 + # #* Layers: + # #* - !Ref ChatDependencies + # MemorySize: 1024 + # Handler: handlers/chat_sync.handler + # Timeout: 300 + # Environment: + # Variables: + # API_TOKEN_NAME: !Ref ApiTokenName + # ENV_PREFIX: !Ref EnvironmentPrefix + # HONEYBADGER_API_KEY: !Ref HoneybadgerApiKey + # HONEYBADGER_ENVIRONMENT: !Ref HoneybadgerEnv + # HONEYBADGER_REVISION: !Ref HoneybadgerRevision + # METRICS_LOG_GROUP: !Ref ChatMetricsLog + # SECRETS_PATH: !Ref SecretsPath + # FunctionUrlConfig: + # AuthType: NONE + # Policies: + # - !Ref SecretsPolicy + # - Statement: + # - Effect: Allow + # Action: + # - 'es:ESHttpGet' + # - 'es:ESHttpPost' + # Resource: '*' # - Statement: # - Effect: Allow # Action: