refactor: simplify internal usage of gateway_logger (#88)

Snowflake-Labs · Dec 10, 2024 · c16e088 · c16e088
1 parent 99945ff
commit c16e088
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 29 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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):