mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
Clean up and follow auto import pattern
This commit is contained in:
@@ -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 pydantic import BaseModel, ConfigDict, Field, PrivateAttr
|
||||
|
||||
# Type checking import
|
||||
if TYPE_CHECKING:
|
||||
from firecrawl import FirecrawlApp
|
||||
|
||||
|
||||
try:
|
||||
from firecrawl import FirecrawlApp
|
||||
|
||||
FIRECRAWL_AVAILABLE = True
|
||||
except ImportError:
|
||||
FirecrawlApp = Any
|
||||
FIRECRAWL_AVAILABLE = False
|
||||
|
||||
|
||||
class FirecrawlSearchToolSchema(BaseModel):
|
||||
@@ -51,9 +56,10 @@ class FirecrawlSearchTool(BaseTool):
|
||||
|
||||
def _initialize_firecrawl(self) -> None:
|
||||
try:
|
||||
from firecrawl import FirecrawlApp # type: ignore
|
||||
|
||||
self.firecrawl = FirecrawlApp(api_key=self.api_key)
|
||||
if FIRECRAWL_AVAILABLE:
|
||||
self._firecrawl = FirecrawlApp(api_key=self.api_key)
|
||||
else:
|
||||
raise ImportError
|
||||
except ImportError:
|
||||
import click
|
||||
|
||||
|
||||
@@ -75,11 +75,6 @@ class PatronusLocalEvaluatorTool(BaseTool):
|
||||
|
||||
try:
|
||||
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._generate_description()
|
||||
print(
|
||||
@@ -96,12 +91,6 @@ class PatronusLocalEvaluatorTool(BaseTool):
|
||||
self,
|
||||
**kwargs: 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_output = kwargs.get("evaluated_model_output")
|
||||
evaluated_model_retrieved_context = kwargs.get(
|
||||
|
||||
@@ -133,13 +133,7 @@ class SnowflakeSearchTool(BaseTool):
|
||||
],
|
||||
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._pool_lock = asyncio.Lock()
|
||||
self._thread_pool = ThreadPoolExecutor(max_workers=self.pool_size)
|
||||
@@ -153,13 +147,6 @@ class SnowflakeSearchTool(BaseTool):
|
||||
|
||||
async def _get_connection(self) -> "SnowflakeConnection":
|
||||
"""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:
|
||||
if not self._connection_pool:
|
||||
conn = await asyncio.get_event_loop().run_in_executor(
|
||||
@@ -199,12 +186,6 @@ class SnowflakeSearchTool(BaseTool):
|
||||
self, query: str, timeout: int = 300
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""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:
|
||||
cache_key = self._get_cache_key(query, timeout)
|
||||
@@ -249,12 +230,6 @@ class SnowflakeSearchTool(BaseTool):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""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:
|
||||
# Override database/schema if provided
|
||||
|
||||
Reference in New Issue
Block a user