Clean up and follow auto import pattern

This commit is contained in:
Brandon Hancock
2025-01-23 15:23:12 -05:00
parent 141ff864f2
commit bcb72a9305
3 changed files with 12 additions and 42 deletions

View File

@@ -1,13 +1,18 @@
from typing import Any, Dict, Optional, Type from typing import TYPE_CHECKING, Any, Dict, Optional, Type
from crewai.tools import BaseTool from crewai.tools import BaseTool
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
# Type checking import if TYPE_CHECKING:
from firecrawl import FirecrawlApp
try: try:
from firecrawl import FirecrawlApp from firecrawl import FirecrawlApp
FIRECRAWL_AVAILABLE = True
except ImportError: except ImportError:
FirecrawlApp = Any FIRECRAWL_AVAILABLE = False
class FirecrawlSearchToolSchema(BaseModel): class FirecrawlSearchToolSchema(BaseModel):
@@ -51,9 +56,10 @@ class FirecrawlSearchTool(BaseTool):
def _initialize_firecrawl(self) -> None: def _initialize_firecrawl(self) -> None:
try: try:
from firecrawl import FirecrawlApp # type: ignore if FIRECRAWL_AVAILABLE:
self._firecrawl = FirecrawlApp(api_key=self.api_key)
self.firecrawl = FirecrawlApp(api_key=self.api_key) else:
raise ImportError
except ImportError: except ImportError:
import click import click

View File

@@ -75,11 +75,6 @@ class PatronusLocalEvaluatorTool(BaseTool):
try: try:
subprocess.run(["uv", "add", "patronus"], check=True) subprocess.run(["uv", "add", "patronus"], check=True)
global patronus # Needed to re-import patronus after installation
import patronus # noqa
global PYPATRONUS_AVAILABLE
PYPATRONUS_AVAILABLE = True
self.client = patronus_client self.client = patronus_client
self._generate_description() self._generate_description()
print( print(
@@ -96,12 +91,6 @@ class PatronusLocalEvaluatorTool(BaseTool):
self, self,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if not PYPATRONUS_AVAILABLE:
raise ImportError(
"The 'patronus' package is not installed. "
"Please install it by running `uv add patronus` to use PatronusLocalEvaluatorTool."
)
evaluated_model_input = kwargs.get("evaluated_model_input") evaluated_model_input = kwargs.get("evaluated_model_input")
evaluated_model_output = kwargs.get("evaluated_model_output") evaluated_model_output = kwargs.get("evaluated_model_output")
evaluated_model_retrieved_context = kwargs.get( evaluated_model_retrieved_context = kwargs.get(

View File

@@ -133,13 +133,7 @@ class SnowflakeSearchTool(BaseTool):
], ],
check=True, check=True,
) )
global snowflake, default_backend, serialization # Needed to re-import after installation
import snowflake.connector # noqa
from cryptography.hazmat.backends import default_backend # noqa
from cryptography.hazmat.primitives import serialization # noqa
global SNOWFLAKE_AVAILABLE
SNOWFLAKE_AVAILABLE = True
self._connection_pool = [] self._connection_pool = []
self._pool_lock = asyncio.Lock() self._pool_lock = asyncio.Lock()
self._thread_pool = ThreadPoolExecutor(max_workers=self.pool_size) self._thread_pool = ThreadPoolExecutor(max_workers=self.pool_size)
@@ -153,13 +147,6 @@ class SnowflakeSearchTool(BaseTool):
async def _get_connection(self) -> "SnowflakeConnection": async def _get_connection(self) -> "SnowflakeConnection":
"""Get a connection from the pool or create a new one.""" """Get a connection from the pool or create a new one."""
if not SNOWFLAKE_AVAILABLE:
raise ImportError(
"The 'snowflake-connector-python' package is not installed. "
"Please install it by running `uv add cryptography snowflake-connector-python snowflake-sqlalchemy` "
"to use SnowflakeSearchTool."
)
async with self._pool_lock: async with self._pool_lock:
if not self._connection_pool: if not self._connection_pool:
conn = await asyncio.get_event_loop().run_in_executor( conn = await asyncio.get_event_loop().run_in_executor(
@@ -199,12 +186,6 @@ class SnowflakeSearchTool(BaseTool):
self, query: str, timeout: int = 300 self, query: str, timeout: int = 300
) -> List[Dict[str, Any]]: ) -> List[Dict[str, Any]]:
"""Execute a query with retries and return results.""" """Execute a query with retries and return results."""
if not SNOWFLAKE_AVAILABLE:
raise ImportError(
"The 'snowflake-connector-python' package is not installed. "
"Please install it by running `uv add cryptography snowflake-connector-python snowflake-sqlalchemy` "
"to use SnowflakeSearchTool."
)
if self.enable_caching: if self.enable_caching:
cache_key = self._get_cache_key(query, timeout) cache_key = self._get_cache_key(query, timeout)
@@ -249,12 +230,6 @@ class SnowflakeSearchTool(BaseTool):
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
"""Execute the search query.""" """Execute the search query."""
if not SNOWFLAKE_AVAILABLE:
raise ImportError(
"The 'snowflake-connector-python' package is not installed. "
"Please install it by running `uv add cryptography snowflake-connector-python snowflake-sqlalchemy` "
"to use SnowflakeSearchTool."
)
try: try:
# Override database/schema if provided # Override database/schema if provided