diff --git a/graphiti_core/driver/driver.py b/graphiti_core/driver/driver.py index d56524d2..a8d94022 100644 --- a/graphiti_core/driver/driver.py +++ b/graphiti_core/driver/driver.py @@ -34,9 +34,6 @@ except ImportError: helpers = None _HAS_OPENSEARCH = False -if TYPE_CHECKING: - from opensearchpy import OpenSearch, helpers - logger = logging.getLogger(__name__) DEFAULT_SIZE = 10 @@ -166,7 +163,7 @@ class GraphDriver(ABC): '' # Neo4j (default) syntax does not require a prefix for fulltext queries ) _database: str - aoss_client: OpenSearch | None + aoss_client: OpenSearch | None # type: ignore @abstractmethod def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine: diff --git a/graphiti_core/driver/neo4j_driver.py b/graphiti_core/driver/neo4j_driver.py index 7bf5b4a7..f10d176b 100644 --- a/graphiti_core/driver/neo4j_driver.py +++ b/graphiti_core/driver/neo4j_driver.py @@ -62,16 +62,16 @@ class Neo4jDriver(GraphDriver): if aoss_host and aoss_port and boto3 is not None: try: session = boto3.Session() - self.aoss_client = OpenSearch( + self.aoss_client = OpenSearch( # type: ignore hosts=[{'host': aoss_host, 'port': aoss_port}], - http_auth=Urllib3AWSV4SignerAuth( + http_auth=Urllib3AWSV4SignerAuth( # type: ignore session.get_credentials(), session.region_name, 'aoss' ), use_ssl=True, verify_certs=True, connection_class=Urllib3HttpConnection, pool_maxsize=20, - ) + ) # type: ignore except Exception as e: logger.warning(f'Failed to initialize OpenSearch client: {e}') self.aoss_client = None diff --git a/graphiti_core/search/search_filters.py b/graphiti_core/search/search_filters.py index ade829f5..4f4cb8d1 100644 --- a/graphiti_core/search/search_filters.py +++ b/graphiti_core/search/search_filters.py @@ -246,7 +246,7 @@ def build_aoss_node_filters(group_ids: list[str], search_filters: SearchFilters) def build_aoss_edge_filters(group_ids: list[str], search_filters: SearchFilters) -> list[dict]: - filters = [{'terms': {'group_id': group_ids}}] + filters: list[dict] = [{'terms': {'group_id': group_ids}}] if search_filters.edge_types: filters.append({'terms': {'edge_types': search_filters.edge_types}})