dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/memory/databricks.py

@@ -59,23 +59,23 @@ class AsyncDatabricksCheckpointSaver(DatabricksCheckpointSaver):
     async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
         """Async version of get_tuple."""
         thread_id = config.get("configurable", {}).get("thread_id", "unknown")
-        logger.debug(f"aget_tuple: Fetching checkpoint for thread_id={thread_id}")
+        logger.trace("Fetching checkpoint", thread_id=thread_id, method="aget_tuple")
         result = await asyncio.to_thread(self.get_tuple, config)
         if result:
-            logger.debug(f"aget_tuple: Found checkpoint for thread_id={thread_id}")
+            logger.trace("Checkpoint found", thread_id=thread_id)
         else:
-            logger.debug(f"aget_tuple: No checkpoint found for thread_id={thread_id}")
+            logger.trace("No checkpoint found", thread_id=thread_id)
         return result
 
     async def aget(self, config: RunnableConfig) -> Checkpoint | None:
         """Async version of get."""
         thread_id = config.get("configurable", {}).get("thread_id", "unknown")
-        logger.debug(f"aget: Fetching checkpoint for thread_id={thread_id}")
+        logger.trace("Fetching checkpoint", thread_id=thread_id, method="aget")
         result = await asyncio.to_thread(self.get, config)
         if result:
-            logger.debug(f"aget: Found checkpoint for thread_id={thread_id}")
+            logger.trace("Checkpoint found", thread_id=thread_id)
         else:
-            logger.debug(f"aget: No checkpoint found for thread_id={thread_id}")
+            logger.trace("No checkpoint found", thread_id=thread_id)
         return result
 
     async def aput(
@@ -88,13 +88,15 @@ class AsyncDatabricksCheckpointSaver(DatabricksCheckpointSaver):
         """Async version of put."""
         thread_id = config.get("configurable", {}).get("thread_id", "unknown")
         checkpoint_id = checkpoint.get("id", "unknown")
-        logger.debug(
-            f"aput: Saving checkpoint id={checkpoint_id} for thread_id={thread_id}"
+        logger.trace(
+            "Saving checkpoint", checkpoint_id=checkpoint_id, thread_id=thread_id
         )
         result = await asyncio.to_thread(
             self.put, config, checkpoint, metadata, new_versions
         )
-        logger.debug(f"aput: Checkpoint saved for thread_id={thread_id}")
+        logger.trace(
+            "Checkpoint saved", thread_id=thread_id, checkpoint_id=checkpoint_id
+        )
         return result
 
     async def aput_writes(
@@ -106,12 +108,14 @@ class AsyncDatabricksCheckpointSaver(DatabricksCheckpointSaver):
     ) -> None:
         """Async version of put_writes."""
         thread_id = config.get("configurable", {}).get("thread_id", "unknown")
-        logger.debug(
-            f"aput_writes: Saving {len(writes)} writes for thread_id={thread_id}, "
-            f"task_id={task_id}"
+        logger.trace(
+            "Saving checkpoint writes",
+            writes_count=len(writes),
+            thread_id=thread_id,
+            task_id=task_id,
         )
         await asyncio.to_thread(self.put_writes, config, writes, task_id, task_path)
-        logger.debug(f"aput_writes: Writes saved for thread_id={thread_id}")
+        logger.trace("Checkpoint writes saved", thread_id=thread_id, task_id=task_id)
 
     async def alist(
         self,
@@ -127,22 +131,20 @@ class AsyncDatabricksCheckpointSaver(DatabricksCheckpointSaver):
             if config
             else "all"
         )
-        logger.debug(
-            f"alist: Listing checkpoints for thread_id={thread_id}, limit={limit}"
-        )
+        logger.trace("Listing checkpoints", thread_id=thread_id, limit=limit)
         # Get all items from sync iterator in a thread
         items = await asyncio.to_thread(
             lambda: list(self.list(config, filter=filter, before=before, limit=limit))
         )
-        logger.debug(f"alist: Found {len(items)} checkpoints for thread_id={thread_id}")
+        logger.debug("Checkpoints listed", thread_id=thread_id, count=len(items))
         for item in items:
             yield item
 
     async def adelete_thread(self, thread_id: str) -> None:
         """Async version of delete_thread."""
-        logger.debug(f"adelete_thread: Deleting thread_id={thread_id}")
+        logger.trace("Deleting thread", thread_id=thread_id)
         await asyncio.to_thread(self.delete_thread, thread_id)
-        logger.debug(f"adelete_thread: Thread deleted thread_id={thread_id}")
+        logger.debug("Thread deleted", thread_id=thread_id)
 
 
 class AsyncDatabricksStore(DatabricksStore):
@@ -156,9 +158,9 @@ class AsyncDatabricksStore(DatabricksStore):
     async def abatch(self, ops: Iterable[Op]) -> list[Result]:
         """Async version of batch."""
         ops_list = list(ops)
-        logger.debug(f"abatch: Executing {len(ops_list)} operations")
+        logger.trace("Executing batch operations", operations_count=len(ops_list))
         result = await asyncio.to_thread(self.batch, ops_list)
-        logger.debug(f"abatch: Completed {len(result)} operations")
+        logger.debug("Batch operations completed", operations_count=len(result))
         return result
 
     async def aget(
@@ -170,14 +172,14 @@ class AsyncDatabricksStore(DatabricksStore):
     ) -> Item | None:
         """Async version of get."""
         ns_str = "/".join(namespace)
-        logger.debug(f"aget: Fetching key={key} from namespace={ns_str}")
+        logger.trace("Fetching store item", key=key, namespace=ns_str)
         result = await asyncio.to_thread(
             partial(self.get, namespace, key, refresh_ttl=refresh_ttl)
         )
         if result:
-            logger.debug(f"aget: Found item key={key} in namespace={ns_str}")
+            logger.trace("Store item found", key=key, namespace=ns_str)
         else:
-            logger.debug(f"aget: No item found key={key} in namespace={ns_str}")
+            logger.trace("Store item not found", key=key, namespace=ns_str)
         return result
 
     async def aput(
@@ -191,7 +193,7 @@ class AsyncDatabricksStore(DatabricksStore):
     ) -> None:
         """Async version of put."""
         ns_str = "/".join(namespace)
-        logger.debug(f"aput: Storing key={key} in namespace={ns_str}")
+        logger.trace("Storing item", key=key, namespace=ns_str, has_ttl=ttl is not None)
         # Handle the ttl parameter - only pass if explicitly provided
         if ttl is not None:
             await asyncio.to_thread(
@@ -199,14 +201,14 @@ class AsyncDatabricksStore(DatabricksStore):
             )
         else:
             await asyncio.to_thread(partial(self.put, namespace, key, value, index))
-        logger.debug(f"aput: Stored key={key} in namespace={ns_str}")
+        logger.trace("Item stored", key=key, namespace=ns_str)
 
     async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
         """Async version of delete."""
         ns_str = "/".join(namespace)
-        logger.debug(f"adelete: Deleting key={key} from namespace={ns_str}")
+        logger.trace("Deleting item", key=key, namespace=ns_str)
         await asyncio.to_thread(self.delete, namespace, key)
-        logger.debug(f"adelete: Deleted key={key} from namespace={ns_str}")
+        logger.trace("Item deleted", key=key, namespace=ns_str)
 
     async def asearch(
         self,
@@ -221,8 +223,8 @@ class AsyncDatabricksStore(DatabricksStore):
     ) -> list[SearchItem]:
         """Async version of search."""
         ns_str = "/".join(namespace_prefix)
-        logger.debug(
-            f"asearch: Searching namespace_prefix={ns_str}, query={query}, limit={limit}"
+        logger.trace(
+            "Searching store", namespace_prefix=ns_str, query=query, limit=limit
         )
         result = await asyncio.to_thread(
             partial(
@@ -235,7 +237,9 @@ class AsyncDatabricksStore(DatabricksStore):
                 refresh_ttl=refresh_ttl,
             )
         )
-        logger.debug(f"asearch: Found {len(result)} items in namespace_prefix={ns_str}")
+        logger.debug(
+            "Store search completed", namespace_prefix=ns_str, results_count=len(result)
+        )
         return result
 
     async def alist_namespaces(
@@ -249,9 +253,7 @@ class AsyncDatabricksStore(DatabricksStore):
     ) -> list[tuple[str, ...]]:
         """Async version of list_namespaces."""
         prefix_str = "/".join(prefix) if prefix else "all"
-        logger.debug(
-            f"alist_namespaces: Listing namespaces prefix={prefix_str}, limit={limit}"
-        )
+        logger.trace("Listing namespaces", prefix=prefix_str, limit=limit)
         result = await asyncio.to_thread(
             partial(
                 self.list_namespaces,
@@ -262,7 +264,7 @@ class AsyncDatabricksStore(DatabricksStore):
                 offset=offset,
             )
         )
-        logger.debug(f"alist_namespaces: Found {len(result)} namespaces")
+        logger.debug("Namespaces listed", count=len(result))
         return result
 
 
@@ -297,7 +299,7 @@ class DatabricksCheckpointerManager(CheckpointManagerBase):
             workspace_client = database.workspace_client
 
             logger.debug(
-                f"Creating AsyncDatabricksCheckpointSaver for instance: {instance_name}"
+                "Creating Databricks checkpointer", instance_name=instance_name
             )
 
             checkpointer = AsyncDatabricksCheckpointSaver(
@@ -306,10 +308,10 @@ class DatabricksCheckpointerManager(CheckpointManagerBase):
             )
 
             # Setup the checkpointer (creates necessary tables if needed)
-            logger.debug(f"Setting up checkpoint tables for instance: {instance_name}")
+            logger.debug("Setting up checkpoint tables", instance_name=instance_name)
             checkpointer.setup()
-            logger.debug(
-                f"Checkpoint tables setup complete for instance: {instance_name}"
+            logger.success(
+                "Databricks checkpointer initialized", instance_name=instance_name
             )
 
             self._checkpointer = checkpointer
@@ -360,12 +362,18 @@ class DatabricksStoreManager(StoreManagerBase):
                 embedding_dims = self.store_model.dims
 
                 logger.debug(
-                    f"Configuring embeddings: endpoint={embedding_endpoint}, dims={embedding_dims}"
+                    "Configuring store embeddings",
+                    endpoint=embedding_endpoint,
+                    dimensions=embedding_dims,
                 )
 
                 embeddings = DatabricksEmbeddings(endpoint=embedding_endpoint)
 
-            logger.debug(f"Creating AsyncDatabricksStore for instance: {instance_name}")
+            logger.debug(
+                "Creating Databricks store",
+                instance_name=instance_name,
+                embeddings_enabled=embeddings is not None,
+            )
 
             store = AsyncDatabricksStore(
                 instance_name=instance_name,
@@ -376,6 +384,11 @@ class DatabricksStoreManager(StoreManagerBase):
 
             # Setup the store (creates necessary tables if needed)
             store.setup()
+            logger.success(
+                "Databricks store initialized",
+                instance_name=instance_name,
+                embeddings_enabled=embeddings is not None,
+            )
             self._store = store
 
         return self._store

dao_ai/memory/postgres.py

@@ -28,9 +28,10 @@ def _create_pool(
     kwargs: dict,
 ) -> ConnectionPool:
     """Create a connection pool using the provided connection parameters."""
-    logger.debug(
-        f"Connection params for {database_name}: {', '.join(k + '=' + (str(v) if k != 'password' else '***') for k, v in connection_params.items())}"
-    )
+    safe_params = {
+        k: (str(v) if k != "password" else "***") for k, v in connection_params.items()
+    }
+    logger.debug("Creating connection pool", database=database_name, **safe_params)
 
     # Merge connection_params into kwargs for psycopg
     connection_kwargs = kwargs | connection_params
@@ -43,7 +44,11 @@ def _create_pool(
         kwargs=connection_kwargs,
     )
     pool.open(wait=True, timeout=timeout_seconds)
-    logger.info(f"Successfully connected to {database_name}")
+    logger.success(
+        "PostgreSQL connection pool created",
+        database=database_name,
+        pool_size=max_pool_size,
+    )
     return pool
 
 
@@ -55,8 +60,11 @@ async def _create_async_pool(
     kwargs: dict,
 ) -> AsyncConnectionPool:
     """Create an async connection pool using the provided connection parameters."""
+    safe_params = {
+        k: (str(v) if k != "password" else "***") for k, v in connection_params.items()
+    }
     logger.debug(
-        f"Connection params for {database_name}: {', '.join(k + '=' + (str(v) if k != 'password' else '***') for k, v in connection_params.items())}"
+        "Creating async connection pool", database=database_name, **safe_params
     )
 
     # Merge connection_params into kwargs for psycopg
@@ -69,7 +77,11 @@ async def _create_async_pool(
         kwargs=connection_kwargs,
     )
     await pool.open(wait=True, timeout=timeout_seconds)
-    logger.info(f"Successfully connected to {database_name}")
+    logger.success(
+        "Async PostgreSQL connection pool created",
+        database=database_name,
+        pool_size=max_pool_size,
+    )
     return pool
 
 
@@ -84,10 +96,12 @@ class AsyncPostgresPoolManager:
 
         async with cls._lock:
             if connection_key in cls._pools:
-                logger.debug(f"Reusing existing PostgreSQL pool for {database.name}")
+                logger.trace(
+                    "Reusing existing async PostgreSQL pool", database=database.name
+                )
                 return cls._pools[connection_key]
 
-            logger.debug(f"Creating new PostgreSQL pool for {database.name}")
+            logger.debug("Creating new async PostgreSQL pool", database=database.name)
 
             kwargs: dict[str, Any] = {
                 "row_factory": dict_row,
@@ -114,7 +128,7 @@ class AsyncPostgresPoolManager:
             if connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 await pool.close()
-                logger.debug(f"Closed PostgreSQL pool for {database.name}")
+                logger.debug("Async PostgreSQL pool closed", database=database.name)
 
     @classmethod
     async def close_all_pools(cls):
@@ -123,17 +137,21 @@ class AsyncPostgresPoolManager:
                 try:
                     # Use a short timeout to avoid blocking on pool closure
                     await asyncio.wait_for(pool.close(), timeout=2.0)
-                    logger.debug(f"Closed PostgreSQL pool: {connection_key}")
+                    logger.debug("Async PostgreSQL pool closed", pool=connection_key)
                 except asyncio.TimeoutError:
                     logger.warning(
-                        f"Timeout closing pool {connection_key}, forcing closure"
+                        "Timeout closing async pool, forcing closure",
+                        pool=connection_key,
                     )
                 except asyncio.CancelledError:
                     logger.warning(
-                        f"Pool closure cancelled for {connection_key} (shutdown in progress)"
+                        "Async pool closure cancelled (shutdown in progress)",
+                        pool=connection_key,
                     )
                 except Exception as e:
-                    logger.error(f"Error closing pool {connection_key}: {e}")
+                    logger.error(
+                        "Error closing async pool", pool=connection_key, error=str(e)
+                    )
             cls._pools.clear()
 
 
@@ -181,12 +199,16 @@ class AsyncPostgresStoreManager(StoreManagerBase):
             await self._store.setup()
 
             self._setup_complete = True
-            logger.debug(
-                f"PostgresStore initialized successfully for {self.store_model.name}"
+            logger.success(
+                "Async PostgreSQL store initialized", store=self.store_model.name
             )
 
         except Exception as e:
-            logger.error(f"Error setting up PostgresStore: {e}")
+            logger.error(
+                "Error setting up async PostgreSQL store",
+                store=self.store_model.name,
+                error=str(e),
+            )
             raise
 
 
@@ -244,12 +266,17 @@ class AsyncPostgresCheckpointerManager(CheckpointManagerBase):
             await self._checkpointer.setup()
 
             self._setup_complete = True
-            logger.debug(
-                f"PostgresSaver initialized successfully for {self.checkpointer_model.name}"
+            logger.success(
+                "Async PostgreSQL checkpointer initialized",
+                checkpointer=self.checkpointer_model.name,
             )
 
         except Exception as e:
-            logger.error(f"Error setting up PostgresSaver: {e}")
+            logger.error(
+                "Error setting up async PostgreSQL checkpointer",
+                checkpointer=self.checkpointer_model.name,
+                error=str(e),
+            )
             raise
 
 
@@ -269,10 +296,10 @@ class PostgresPoolManager:
 
         with cls._lock:
             if connection_key in cls._pools:
-                logger.debug(f"Reusing existing PostgreSQL pool for {database.name}")
+                logger.trace("Reusing existing PostgreSQL pool", database=database.name)
                 return cls._pools[connection_key]
 
-            logger.debug(f"Creating new PostgreSQL pool for {database.name}")
+            logger.debug("Creating new PostgreSQL pool", database=database.name)
 
             kwargs: dict[str, Any] = {
                 "row_factory": dict_row,
@@ -299,7 +326,7 @@ class PostgresPoolManager:
             if connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 pool.close()
-                logger.debug(f"Closed PostgreSQL pool for {database.name}")
+                logger.debug("PostgreSQL pool closed", database=database.name)
 
     @classmethod
     def close_all_pools(cls):
@@ -307,9 +334,13 @@ class PostgresPoolManager:
             for connection_key, pool in cls._pools.items():
                 try:
                     pool.close()
-                    logger.debug(f"Closed PostgreSQL pool: {connection_key}")
+                    logger.debug("PostgreSQL pool closed", pool=connection_key)
                 except Exception as e:
-                    logger.error(f"Error closing pool {connection_key}: {e}")
+                    logger.error(
+                        "Error closing PostgreSQL pool",
+                        pool=connection_key,
+                        error=str(e),
+                    )
             cls._pools.clear()
 
 
@@ -349,12 +380,14 @@ class PostgresStoreManager(StoreManagerBase):
             self._store.setup()
 
             self._setup_complete = True
-            logger.debug(
-                f"PostgresStore initialized successfully for {self.store_model.name}"
-            )
+            logger.success("PostgreSQL store initialized", store=self.store_model.name)
 
         except Exception as e:
-            logger.error(f"Error setting up PostgresStore: {e}")
+            logger.error(
+                "Error setting up PostgreSQL store",
+                store=self.store_model.name,
+                error=str(e),
+            )
             raise
 
 
@@ -400,21 +433,28 @@ class PostgresCheckpointerManager(CheckpointManagerBase):
             self._checkpointer.setup()
 
             self._setup_complete = True
-            logger.debug(
-                f"PostgresSaver initialized successfully for {self.checkpointer_model.name}"
+            logger.success(
+                "PostgreSQL checkpointer initialized",
+                checkpointer=self.checkpointer_model.name,
             )
 
         except Exception as e:
-            logger.error(f"Error setting up PostgresSaver: {e}")
+            logger.error(
+                "Error setting up PostgreSQL checkpointer",
+                checkpointer=self.checkpointer_model.name,
+                error=str(e),
+            )
             raise
 
 
 def _shutdown_pools() -> None:
     try:
         PostgresPoolManager.close_all_pools()
-        logger.debug("Successfully closed all synchronous PostgreSQL pools")
+        logger.debug("All synchronous PostgreSQL pools closed during shutdown")
     except Exception as e:
-        logger.error(f"Error closing synchronous PostgreSQL pools during shutdown: {e}")
+        logger.error(
+            "Error closing synchronous PostgreSQL pools during shutdown", error=str(e)
+        )
 
 
 def _shutdown_async_pools() -> None:
@@ -434,15 +474,16 @@ def _shutdown_async_pools() -> None:
                     loop = asyncio.new_event_loop()
                     asyncio.set_event_loop(loop)
                 loop.run_until_complete(AsyncPostgresPoolManager.close_all_pools())
-                logger.debug("Successfully closed all asynchronous PostgreSQL pools")
+                logger.debug("All asynchronous PostgreSQL pools closed during shutdown")
             except Exception as inner_e:
                 # If all else fails, just log the error
                 logger.warning(
-                    f"Could not close async pools cleanly during shutdown: {inner_e}"
+                    "Could not close async pools cleanly during shutdown",
+                    error=str(inner_e),
                 )
     except Exception as e:
         logger.error(
-            f"Error closing asynchronous PostgreSQL pools during shutdown: {e}"
+            "Error closing asynchronous PostgreSQL pools during shutdown", error=str(e)
         )
 
 

dao_ai/middleware/assertions.py

@@ -367,20 +367,27 @@ class AssertMiddleware(AgentMiddleware[AgentState, Context]):
             "runtime": runtime,
         }
 
-        logger.debug(f"Evaluating Assert constraint '{self.constraint.name}'")
+        logger.trace(
+            "Evaluating Assert constraint", constraint_name=self.constraint.name
+        )
 
         result = self.constraint.evaluate(response, context)
 
         if result.passed:
-            logger.debug(f"Assert constraint '{self.constraint.name}' passed")
+            logger.trace(
+                "Assert constraint passed", constraint_name=self.constraint.name
+            )
             self._retry_count = 0
             return None
 
         # Constraint failed
         self._retry_count += 1
         logger.warning(
-            f"Assert constraint '{self.constraint.name}' failed "
-            f"(attempt {self._retry_count}/{self.max_retries}): {result.feedback}"
+            "Assert constraint failed",
+            constraint_name=self.constraint.name,
+            attempt=self._retry_count,
+            max_retries=self.max_retries,
+            feedback=result.feedback,
         )
 
         if self._retry_count >= self.max_retries:
@@ -396,7 +403,8 @@ class AssertMiddleware(AgentMiddleware[AgentState, Context]):
                 return None
             else:  # "pass"
                 logger.warning(
-                    f"Assert constraint '{self.constraint.name}' failed but passing through"
+                    "Assert constraint failed but passing through",
+                    constraint_name=self.constraint.name,
                 )
                 return None
 
@@ -475,25 +483,38 @@ class SuggestMiddleware(AgentMiddleware[AgentState, Context]):
             "runtime": runtime,
         }
 
-        logger.debug(f"Evaluating Suggest constraint '{self.constraint.name}'")
+        logger.trace(
+            "Evaluating Suggest constraint", constraint_name=self.constraint.name
+        )
 
         result = self.constraint.evaluate(response, context)
 
         if result.passed:
-            logger.debug(f"Suggest constraint '{self.constraint.name}' passed")
+            logger.trace(
+                "Suggest constraint passed", constraint_name=self.constraint.name
+            )
             self._has_retried = False
             return None
 
         # Log feedback based on configured level
-        log_msg = (
-            f"Suggest constraint '{self.constraint.name}' feedback: {result.feedback}"
-        )
         if self.log_level == "warning":
-            logger.warning(log_msg)
+            logger.warning(
+                "Suggest constraint feedback",
+                constraint_name=self.constraint.name,
+                feedback=result.feedback,
+            )
         elif self.log_level == "info":
-            logger.info(log_msg)
+            logger.info(
+                "Suggest constraint feedback",
+                constraint_name=self.constraint.name,
+                feedback=result.feedback,
+            )
         else:
-            logger.debug(log_msg)
+            logger.debug(
+                "Suggest constraint feedback",
+                constraint_name=self.constraint.name,
+                feedback=result.feedback,
+            )
 
         # Optionally request one improvement
         if self.allow_one_retry and not self._has_retried:
@@ -593,8 +614,11 @@ class RefineMiddleware(AgentMiddleware[AgentState, Context]):
         self._iteration += 1
 
         logger.debug(
-            f"Refine iteration {self._iteration}/{self.max_iterations}: "
-            f"score={score:.3f}, threshold={self.threshold}"
+            "Refine iteration",
+            iteration=self._iteration,
+            max_iterations=self.max_iterations,
+            score=f"{score:.3f}",
+            threshold=self.threshold,
         )
 
         # Track best response
@@ -604,13 +628,17 @@ class RefineMiddleware(AgentMiddleware[AgentState, Context]):
 
         # Check if we should stop
         if score >= self.threshold:
-            logger.debug(f"Refine threshold reached: {score:.3f} >= {self.threshold}")
+            logger.debug(
+                "Refine threshold reached",
+                score=f"{score:.3f}",
+                threshold=self.threshold,
+            )
             self._reset()
             return None
 
         if self._iteration >= self.max_iterations:
             logger.debug(
-                f"Refine max iterations reached. Best score: {self._best_score:.3f}"
+                "Refine max iterations reached", best_score=f"{self._best_score:.3f}"
             )
             # Use best response if tracking
             if self.select_best and self._best_response: