From c16e088079eabd5d52757badbaadda391154a73b Mon Sep 17 00:00:00 2001 From: Tyler White Date: Tue, 10 Dec 2024 16:15:11 -0500 Subject: [PATCH] refactor: simplify internal usage of gateway_logger (#88) --- agent_gateway/gateway/gateway.py | 9 ++++----- agent_gateway/gateway/planner.py | 5 ++--- agent_gateway/gateway/task_processor.py | 7 +++---- agent_gateway/tools/logger.py | 8 ++------ agent_gateway/tools/snowflake_tools.py | 19 ++++++++----------- 5 files changed, 19 insertions(+), 29 deletions(-) diff --git a/agent_gateway/gateway/gateway.py b/agent_gateway/gateway/gateway.py index ebc704f..629a25a 100644 --- a/agent_gateway/gateway/gateway.py +++ b/agent_gateway/gateway/gateway.py @@ -12,7 +12,6 @@ import asyncio import json -import logging import re import threading from collections.abc import Sequence @@ -37,7 +36,7 @@ class AgentGatewayError(Exception): def __init__(self, message): self.message = message - gateway_logger.log(logging.ERROR, self.message) + gateway_logger.log("ERROR", self.message) super().__init__(self.message) @@ -201,7 +200,7 @@ def __init__( # callbacks self.planner_callback = None self.executor_callback = None - gateway_logger.log(logging.INFO, "Cortex gateway successfully initialized") + gateway_logger.log("INFO", "Cortex gateway successfully initialized") @property def input_keys(self) -> List[str]: @@ -328,8 +327,8 @@ async def fuse( response = await self.agent.arun(prompt) raw_answer = cast(str, response) - gateway_logger.log(logging.DEBUG, "Question: \n", input_query, block=True) - gateway_logger.log(logging.DEBUG, "Raw Answer: \n", raw_answer, block=True) + gateway_logger.log("DEBUG", "Question: \n", input_query, block=True) + gateway_logger.log("DEBUG", "Raw Answer: \n", raw_answer, block=True) thought, answer, is_replan = self._parse_fusion_output(raw_answer) if is_final: # If final, we don't need to replan diff --git a/agent_gateway/gateway/planner.py b/agent_gateway/gateway/planner.py index 529b2f1..9009b8e 100644 --- a/agent_gateway/gateway/planner.py +++ b/agent_gateway/gateway/planner.py @@ -14,7 +14,6 @@ import asyncio import json -import logging import re from collections.abc import Sequence from typing import Any, Optional, Union @@ -39,7 +38,7 @@ class AgentGatewayError(Exception): def __init__(self, message): self.message = message - gateway_logger.log(logging.ERROR, message) + gateway_logger.log("ERROR", message) super().__init__(self.message) @@ -297,7 +296,7 @@ def _parse_snowflake_response(self, data_str): if "content" in choices["delta"].keys(): completion += choices["delta"]["content"] - gateway_logger.log(logging.DEBUG, f"LLM Generated Plan:\n{completion}") + gateway_logger.log("DEBUG", f"LLM Generated Plan:\n{completion}") return completion async def plan(self, inputs: dict, is_replan: bool, **kwargs: Any): diff --git a/agent_gateway/gateway/task_processor.py b/agent_gateway/gateway/task_processor.py index 2941b45..4e9d149 100644 --- a/agent_gateway/gateway/task_processor.py +++ b/agent_gateway/gateway/task_processor.py @@ -13,7 +13,6 @@ from __future__ import annotations import asyncio -import logging from collections.abc import Collection from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional @@ -27,7 +26,7 @@ class AgentGatewayError(Exception): def __init__(self, message): self.message = message - gateway_logger.log(logging.ERROR, self.message) + gateway_logger.log("ERROR", self.message) super().__init__(self.message) @@ -73,10 +72,10 @@ class Task: is_fuse: bool = False async def __call__(self) -> Any: - gateway_logger.log(logging.INFO, f"running {self.name} task") + gateway_logger.log("INFO", f"running {self.name} task") try: x = await self.tool(*self.args) - gateway_logger.log(logging.DEBUG, "task successfully completed") + gateway_logger.log("DEBUG", "task successfully completed") return x except SnowflakeError as e: return f"Unexpected error during Cortex Gateway Tool request: {str(e)}" diff --git a/agent_gateway/tools/logger.py b/agent_gateway/tools/logger.py index ff4b3e3..6404ed8 100644 --- a/agent_gateway/tools/logger.py +++ b/agent_gateway/tools/logger.py @@ -57,6 +57,8 @@ def init(self): self.logger.addHandler(self.file_handler) def log(self, level, *args, block=False, **kwargs): + if isinstance(level, str): + level = getattr(logging, level.upper()) if LOGGING_ENABLED: if block: self.logger.log(level, "=" * 80) @@ -71,9 +73,3 @@ def log(self, level, *args, block=False, **kwargs): gateway_logger = Logger() - -# The updated log function - - -def log(level, *args, block=False, **kwargs): - gateway_logger.log(level, *args, block=block, **kwargs) diff --git a/agent_gateway/tools/snowflake_tools.py b/agent_gateway/tools/snowflake_tools.py index 4bf19ed..2faeec8 100644 --- a/agent_gateway/tools/snowflake_tools.py +++ b/agent_gateway/tools/snowflake_tools.py @@ -13,7 +13,6 @@ import asyncio import inspect import json -import logging import re from typing import Any, Type, Union @@ -34,7 +33,7 @@ class SnowflakeError(Exception): def __init__(self, message): self.message = message - gateway_logger.log(logging.ERROR, message) + gateway_logger.log("ERROR", message) super().__init__(self.message) @@ -78,17 +77,17 @@ def __init__( self.k = k self.retrieval_columns = retrieval_columns self.service_name = service_name - gateway_logger.log(logging.INFO, "Cortex Search Tool successfully initialized") + gateway_logger.log("INFO", "Cortex Search Tool successfully initialized") def __call__(self, question) -> Any: return self.asearch(question) async def asearch(self, query): - gateway_logger.log(logging.DEBUG, f"Cortex Search Query:{query}") + gateway_logger.log("DEBUG", f"Cortex Search Query:{query}") headers, url, data = self._prepare_request(query=query) response_text = await post_cortex_request(url=url, headers=headers, data=data) response_json = json.loads(response_text) - gateway_logger.log(logging.DEBUG, f"Cortex Search Response:{response_json}") + gateway_logger.log("DEBUG", f"Cortex Search Response:{response_json}") try: return response_json["results"] except Exception: @@ -218,22 +217,20 @@ def __init__( self.FILE = semantic_model self.STAGE = stage - gateway_logger.log(logging.INFO, "Cortex Analyst Tool successfully initialized") + gateway_logger.log("INFO", "Cortex Analyst Tool successfully initialized") def __call__(self, prompt) -> Any: return self.asearch(query=prompt) async def asearch(self, query): - gateway_logger.log(logging.DEBUG, f"Cortex Analyst Prompt:{query}") + gateway_logger.log("DEBUG", f"Cortex Analyst Prompt:{query}") url, headers, data = self._prepare_analyst_request(prompt=query) response_text = await post_cortex_request(url=url, headers=headers, data=data) json_response = json.loads(response_text) - gateway_logger.log( - logging.DEBUG, f"Cortex Analyst Raw Response:{json_response}" - ) + gateway_logger.log("DEBUG", f"Cortex Analyst Raw Response:{json_response}") try: query_response = self._process_analyst_message( @@ -316,7 +313,7 @@ def __init__(self, python_func, tool_description, output_description) -> None: name=python_func.__name__, func=python_callable, description=desc ) self.python_callable = python_func - gateway_logger.log(logging.INFO, "Python Tool successfully initialized") + gateway_logger.log("INFO", "Python Tool successfully initialized") def asyncify(self, sync_func): async def async_func(*args, **kwargs):