refactor OBPClient

Author: nemozak1

.github/instructions/ai.md.instructions.md

@@ -19,6 +19,6 @@ applyTo: '**'
 - Keep documentation as simple as possible
 
 # Personality
-- Avoid ecessive sycophancy and pandering. Be professional and respectful.
+- Avoid ecessive sycophancy and pandering.
 - Please disagree explicitly if you think the user is wrong or there is a better way to do something.
 - Do not always agree with the user.

src/auth/admin_client.py

@@ -71,8 +71,29 @@ async def initialize(
 
         except Exception as e:
             logger.error(f'❌ Failed to initialize admin OBP client: {e}')
+            # Clean up any partially created resources
+            await self._cleanup_partial_init()
             raise ValueError(f'Admin client initialization failed: {e}') from e
 
+    async def _cleanup_partial_init(self) -> None:
+        """Clean up resources from a failed initialization."""
+        if self._client:
+            try:
+                await self._client.close()
+            except Exception as cleanup_error:
+                logger.debug(f'Error during cleanup: {cleanup_error}')
+        
+        if self._auth and hasattr(self._auth, 'async_requests_client'):
+            if self._auth.async_requests_client:
+                try:
+                    await self._auth.async_requests_client.close()
+                except Exception as cleanup_error:
+                    logger.debug(f'Error during cleanup: {cleanup_error}')
+        
+        self._client = None
+        self._auth = None
+        self._initialized = False
+    
     def get_client(self) -> OBPClient:
         """
         Get the singleton admin OBP client instance.
@@ -112,10 +133,16 @@ def is_initialized(self) -> bool:
 
     async def close(self) -> None:
         """Clean up resources during app shutdown."""
+        # Close OBP client session
+        if self._client:
+            await self._client.close()
+            logger.info('🔌 Admin OBP client session closed')
+        
+        # Close auth session
         if self._auth and hasattr(self._auth, 'async_requests_client'):
             if self._auth.async_requests_client:
                 await self._auth.async_requests_client.close()
-                logger.info('🔌 Admin client HTTP session closed')
+                logger.info('🔌 Admin auth HTTP session closed')
 
         self._initialized = False
         self._client = None
@@ -160,7 +187,7 @@ def get_admin_client() -> OBPClient:
         >>> 
         >>> # Later, anywhere in the app
         >>> admin_client = get_admin_client()
-        >>> response = await admin_client.get("/obp/v6.0.0/banks")
+        >>> banks = await admin_client.get("/obp/v6.0.0/banks")
     """
     return _admin_manager.get_client()
 

src/auth/auth.py

@@ -459,17 +459,15 @@ async def create_admin_direct_login_auth(
         version = os.getenv('OBP_API_VERSION', 'v6.0.0')
 
         try:
-            response_json = await obp_client.get(
+            entitlements_response = await obp_client.get(
                 f"/obp/{version}/my/entitlements"
             )
+            entitlements_data = entitlements_response.json()
 
-            if not response_json:
+            if not entitlements_data:
                 logger.warning('Failed to fetch entitlements - received empty response')
                 return admin_auth
 
-            # Parse the JSON string response
-            entitlements_response = json.loads(response_json) if isinstance(response_json, str) else response_json
-            
             # Use provided entitlements or default set
             if required_entitlements is None:
                 required_entitlements = [
@@ -479,7 +477,7 @@ async def create_admin_direct_login_auth(
                     'CanGetSystemLevelDynamicEntities'
                 ]
 
-            all_present, missing = _check_entitlements(entitlements_response, required_entitlements)
+            all_present, missing = _check_entitlements(entitlements_data, required_entitlements)
 
             if not all_present:
                 logger.warning(
@@ -490,5 +488,8 @@ async def create_admin_direct_login_auth(
         except Exception as e:
             logger.error(f'Failed to verify admin entitlements: {e}')
             # Don't raise - auth still works, just couldn't verify entitlements
+        finally:
+            # Always close the temporary client session
+            await obp_client.close()
 
     return admin_auth

src/checkpointer/entities.py

@@ -88,7 +88,7 @@ async def read(self, entity_id: str) -> dict:
         response = await self.client.get(
             path=f"{self.endpoint_url}/{entity_id}"
         )
-        return json.loads(response) if isinstance(response, str) else response
+        return response.json()
 
 opey_checkpoint_entity = {
   "hasPersonalEntity": True,

src/checkpointer/obp_checkpoint_saver.py

@@ -10,18 +10,50 @@
 )
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 
-from src.client import obp_client
+from src.client.obp_client import OBPClient
 from src.checkpointer.entities import OpeyCheckpointEntity, OpeyCheckpointWriteEntity
-
+from src.auth.admin_client import get_admin_client
 class OBPCheckpointSaver(BaseCheckpointSaver):
     """Checkpoint saver using OBP Dynamic Entitites as a storage backend."""
 
     def __init__(self, client: OBPClient):
         super().__init__()
+        if not client:
+            try:
+                client = get_admin_client()
+            except RuntimeError as e:
+                raise ValueError("An OBPClient must be provided or the admin client must be initialized.") from e
+        
         self.client = client
         self.is_setup = False
         self.serde = JsonPlusSerializer()
 
+    async def _check_existing_setup(self) -> bool:
+        """
+        Check if the required Dynamic Entities are already set up in OBP.
+        
+        Returns:
+            bool: True if setup exists, False otherwise.
+        """
+        try:
+            response = await self.client.get("/obp/v6.0.0/management/system-dynamic-entities")
+            response_data = response.json()
+        except Exception as e:
+            raise RuntimeError(f"Error checking existing OBP setup: {e}") from e
+        
+        if not response_data or not isinstance(response_data, dict):
+            return False
+        
+        existing_entities = response_data.get("dynamic_entities", [])
+        required_entities = {OpeyCheckpointEntity.obp_entity_name(), OpeyCheckpointWriteEntity.obp_entity_name()}
+        
+        for entity in existing_entities:
+            entity_name = next(iter(entity))
+            if entity_name in required_entities:
+                required_entities.remove(entity_name)
+        
+        return len(required_entities) == 0    
+    
     def setup(self) -> None:
         """
         Setup the Dynamic Entities in OBP if they do not exist.
@@ -30,6 +62,9 @@ def setup(self) -> None:
         if self.is_setup:
             return
 
+        
+        
+        
     def put(
         self,
         config: RunnableConfig,