mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 08:38:30 +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 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
|
||||||
|
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user