From 2b3c959037b8ef2938f629accfb1b6f70d705c16 Mon Sep 17 00:00:00 2001
From: alejandro <alejandro.herrera@snowflake.com>
Date: Thu, 5 Dec 2024 09:35:55 -0500
Subject: [PATCH] refactor: filter cleanup

---
 agent_gateway/tools/snowflake_tools.py | 34 +-------------------------
 1 file changed, 1 insertion(+), 33 deletions(-)

diff --git a/agent_gateway/tools/snowflake_tools.py b/agent_gateway/tools/snowflake_tools.py
index 2d6066e..fe9d988 100644
--- a/agent_gateway/tools/snowflake_tools.py
+++ b/agent_gateway/tools/snowflake_tools.py
@@ -11,14 +11,13 @@
 # limitations under the License.
 
 import asyncio
-import contextlib
 import inspect
 import json
 import logging
 import re
 from typing import Any, Type, Union
 
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel
 from snowflake.connector.connection import SnowflakeConnection
 from snowflake.snowpark import Session
 from snowflake.snowpark.functions import col
@@ -45,8 +44,6 @@ class CortexSearchTool(Tool):
     retrieval_columns: list = []
     service_name: str = ""
     connection: Union[Session, SnowflakeConnection] = None
-    auto_filter: bool = False
-    filter_generator: object = None
 
     def __init__(
         self,
@@ -55,7 +52,6 @@ def __init__(
         data_description,
         retrieval_columns,
         snowflake_connection,
-        auto_filter=False,
         k=5,
     ):
         """Parameters
@@ -185,34 +181,6 @@ def get_min_length(model: Type[BaseModel]):
     return min_length
 
 
-class JSONFilter(BaseModel):
-    answer: str = Field(description="The filter_query in valid JSON format")
-
-    @classmethod
-    def model_validate_json(
-        cls,
-        json_data: str,
-        *,
-        strict: bool | None = None,
-        context: dict[str, Any] | None = None,
-    ):
-        __tracebackhide__ = True
-        try:
-            return cls.__pydantic_validator__.validate_json(
-                json_data, strict=strict, context=context
-            )
-        except ValidationError:
-            min_length = get_min_length(cls)
-            for substring_length in range(len(json_data), min_length - 1, -1):
-                for start in range(len(json_data) - substring_length + 1):
-                    substring = json_data[start : start + substring_length]
-                    with contextlib.suppress(ValidationError):
-                        return cls.__pydantic_validator__.validate_json(
-                            substring, strict=strict, context=context
-                        )
-        raise ValueError("Could not find valid json")
-
-
 class CortexAnalystTool(Tool):
     """""Cortex Analyst tool for use with Snowflake Agent Gateway""" ""