This commit is contained in:
prestonrasmussen 2025-09-12 10:33:48 -04:00
parent 427f917614
commit 2f8b25acf7
2 changed files with 25 additions and 9 deletions

View file

@ -81,7 +81,13 @@ class IsPresidentOf(BaseModel):
async def main(use_bulk: bool = False): async def main(use_bulk: bool = False):
setup_logging() setup_logging()
graph_driver = Neo4jDriver( graph_driver = Neo4jDriver(
neo4j_uri, neo4j_user, neo4j_password, aoss_host=aoss_host, aoss_port=aoss_port neo4j_uri,
neo4j_user,
neo4j_password,
aoss_host=aoss_host,
aoss_port=aoss_port,
region='us-west-2',
service='es',
) )
# client = Graphiti( # client = Graphiti(
# neo4j_uri, # neo4j_uri,

View file

@ -28,7 +28,13 @@ logger = logging.getLogger(__name__)
try: try:
import boto3 import boto3
from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection from opensearchpy import (
AWSV4SignerAuth,
OpenSearch,
RequestsHttpConnection,
Urllib3AWSV4SignerAuth,
Urllib3HttpConnection,
)
_HAS_OPENSEARCH = True _HAS_OPENSEARCH = True
except ImportError: except ImportError:
@ -50,6 +56,8 @@ class Neo4jDriver(GraphDriver):
database: str = 'neo4j', database: str = 'neo4j',
aoss_host: str | None = None, aoss_host: str | None = None,
aoss_port: int | None = None, aoss_port: int | None = None,
region: str | None = None,
service: str | None = None,
): ):
super().__init__() super().__init__()
self.client = AsyncGraphDatabase.driver( self.client = AsyncGraphDatabase.driver(
@ -61,15 +69,17 @@ class Neo4jDriver(GraphDriver):
self.aoss_client = None self.aoss_client = None
if aoss_host and aoss_port and boto3 is not None: if aoss_host and aoss_port and boto3 is not None:
try: try:
session = boto3.Session() region = region
self.aoss_client = OpenSearch( # type: ignore service = service
hosts=[{'host': aoss_host, 'port': aoss_port, 'scheme': 'https'}], credentials = boto3.Session(profile_name='zep-development').get_credentials()
http_auth=Urllib3AWSV4SignerAuth( auth = AWSV4SignerAuth(credentials, region, service)
session.get_credentials(), session.region_name, 'aoss'
), self.aoss_client = OpenSearch(
hosts=[{'host': aoss_host, 'port': aoss_port}],
http_auth=auth,
use_ssl=True, use_ssl=True,
verify_certs=True, verify_certs=True,
connection_class=Urllib3HttpConnection, connection_class=RequestsHttpConnection,
pool_maxsize=20, pool_maxsize=20,
) # type: ignore ) # type: ignore
except Exception as e: except Exception as e: