dao-ai 0.0.28__py3-none-any.whl → 0.1.2__py3-none-any.whl

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,24 +433,31 @@ 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():
+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():
+def _shutdown_async_pools() -> None:
     try:
         # Try to get the current event loop first
         try:
@@ -434,15 +474,16 @@ def _shutdown_async_pools():
                     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/messages.py

@@ -125,8 +125,6 @@ def has_image(messages: BaseMessage | Sequence[BaseMessage]) -> bool:
                     "image_url",
                 ]:
                     return True
-                if hasattr(item, "type") and item.type in ["image", "image_url"]:
-                    return True
         return False
 
     if isinstance(messages, BaseMessage):
@@ -176,7 +174,9 @@ def last_human_message(messages: Sequence[BaseMessage]) -> Optional[HumanMessage
     Returns:
         The last HumanMessage in the sequence, or None if no human messages found
     """
-    return last_message(messages, lambda m: isinstance(m, HumanMessage))
+    return last_message(
+        messages, lambda m: isinstance(m, HumanMessage) and bool(m.content)
+    )
 
 
 def last_ai_message(messages: Sequence[BaseMessage]) -> Optional[AIMessage]:
@@ -192,7 +192,9 @@ def last_ai_message(messages: Sequence[BaseMessage]) -> Optional[AIMessage]:
     Returns:
         The last AIMessage in the sequence, or None if no AI messages found
     """
-    return last_message(messages, lambda m: isinstance(m, AIMessage))
+    return last_message(
+        messages, lambda m: isinstance(m, AIMessage) and bool(m.content)
+    )
 
 
 def last_tool_message(messages: Sequence[BaseMessage]) -> Optional[ToolMessage]:

dao_ai/middleware/__init__.py

@@ -0,0 +1,125 @@
+# DAO AI Middleware Module
+# This module provides middleware implementations compatible with LangChain v1's create_agent
+
+# Re-export LangChain built-in middleware
+from langchain.agents.middleware import (
+    HumanInTheLoopMiddleware,
+    SummarizationMiddleware,
+    after_agent,
+    after_model,
+    before_agent,
+    before_model,
+    dynamic_prompt,
+    wrap_model_call,
+    wrap_tool_call,
+)
+
+# DSPy-style assertion middleware
+from dao_ai.middleware.assertions import (
+    # Middleware classes
+    AssertMiddleware,
+    # Types
+    Constraint,
+    ConstraintResult,
+    FunctionConstraint,
+    KeywordConstraint,
+    LengthConstraint,
+    LLMConstraint,
+    RefineMiddleware,
+    SuggestMiddleware,
+    # Factory functions
+    create_assert_middleware,
+    create_refine_middleware,
+    create_suggest_middleware,
+)
+from dao_ai.middleware.base import (
+    AgentMiddleware,
+    ModelRequest,
+    ModelResponse,
+)
+from dao_ai.middleware.core import create_factory_middleware
+from dao_ai.middleware.guardrails import (
+    ContentFilterMiddleware,
+    GuardrailMiddleware,
+    SafetyGuardrailMiddleware,
+    create_content_filter_middleware,
+    create_guardrail_middleware,
+    create_safety_guardrail_middleware,
+)
+from dao_ai.middleware.human_in_the_loop import (
+    create_hitl_middleware_from_tool_models,
+    create_human_in_the_loop_middleware,
+)
+from dao_ai.middleware.message_validation import (
+    CustomFieldValidationMiddleware,
+    FilterLastHumanMessageMiddleware,
+    MessageValidationMiddleware,
+    RequiredField,
+    ThreadIdValidationMiddleware,
+    UserIdValidationMiddleware,
+    create_custom_field_validation_middleware,
+    create_filter_last_human_message_middleware,
+    create_thread_id_validation_middleware,
+    create_user_id_validation_middleware,
+)
+from dao_ai.middleware.summarization import (
+    LoggingSummarizationMiddleware,
+    create_summarization_middleware,
+)
+
+__all__ = [
+    # Base class (from LangChain)
+    "AgentMiddleware",
+    # Types
+    "ModelRequest",
+    "ModelResponse",
+    # LangChain decorators
+    "before_agent",
+    "before_model",
+    "after_agent",
+    "after_model",
+    "wrap_model_call",
+    "wrap_tool_call",
+    "dynamic_prompt",
+    # LangChain built-in middleware
+    "SummarizationMiddleware",
+    "LoggingSummarizationMiddleware",
+    "HumanInTheLoopMiddleware",
+    # Core factory function
+    "create_factory_middleware",
+    # DAO AI middleware implementations
+    "GuardrailMiddleware",
+    "ContentFilterMiddleware",
+    "SafetyGuardrailMiddleware",
+    "MessageValidationMiddleware",
+    "UserIdValidationMiddleware",
+    "ThreadIdValidationMiddleware",
+    "CustomFieldValidationMiddleware",
+    "RequiredField",
+    "FilterLastHumanMessageMiddleware",
+    # DSPy-style assertion middleware
+    "Constraint",
+    "ConstraintResult",
+    "FunctionConstraint",
+    "KeywordConstraint",
+    "LengthConstraint",
+    "LLMConstraint",
+    "AssertMiddleware",
+    "SuggestMiddleware",
+    "RefineMiddleware",
+    # DAO AI middleware factory functions
+    "create_guardrail_middleware",
+    "create_content_filter_middleware",
+    "create_safety_guardrail_middleware",
+    "create_user_id_validation_middleware",
+    "create_thread_id_validation_middleware",
+    "create_custom_field_validation_middleware",
+    "create_filter_last_human_message_middleware",
+    "create_summarization_middleware",
+    "create_human_in_the_loop_middleware",
+    "create_hitl_middleware_from_tool_models",
+    # DSPy-style assertion factory functions
+    "create_assert_middleware",
+    "create_suggest_middleware",
+    "create_refine_middleware",
+]