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

dao_ai/middleware/core.py

@@ -7,15 +7,17 @@ from fully qualified function names.
 
 from typing import Any, Callable
 
+from langchain.agents.middleware import AgentMiddleware
 from loguru import logger
 
+from dao_ai.state import AgentState, Context
 from dao_ai.utils import load_function
 
 
 def create_factory_middleware(
     function_name: str,
     args: dict[str, Any] | None = None,
-) -> Any:
+) -> AgentMiddleware[AgentState, Context]:
     """
     Create middleware from a factory function.
 
@@ -33,14 +35,14 @@ def create_factory_middleware(
         args: Arguments to pass to the factory function
 
     Returns:
-        A middleware instance returned by the factory function
+        An AgentMiddleware instance returned by the factory function
 
     Raises:
         ImportError: If the function cannot be loaded
 
     Example:
         # Factory function in my_module.py:
-        def create_custom_middleware(threshold: float = 0.5) -> AgentMiddleware:
+        def create_custom_middleware(threshold: float = 0.5) -> AgentMiddleware[AgentState, Context]:
             return MyCustomMiddleware(threshold=threshold)
 
         # Usage:
@@ -52,10 +54,14 @@ def create_factory_middleware(
     if args is None:
         args = {}
 
-    logger.debug(f"Creating factory middleware: {function_name} with args: {args}")
+    logger.trace("Creating factory middleware", function_name=function_name, args=args)
 
-    factory: Callable[..., Any] = load_function(function_name=function_name)
-    middleware: Any = factory(**args)
+    factory: Callable[..., AgentMiddleware[AgentState, Context]] = load_function(
+        function_name=function_name
+    )
+    middleware: AgentMiddleware[AgentState, Context] = factory(**args)
 
-    logger.debug(f"Created middleware from factory: {type(middleware).__name__}")
+    logger.trace(
+        "Created middleware from factory", middleware_type=type(middleware).__name__
+    )
     return middleware

dao_ai/middleware/guardrails.py

@@ -117,29 +117,29 @@ class GuardrailMiddleware(AgentMiddleware[AgentState, Context]):
 
         # Skip evaluation if the AI message has tool calls (not the final response yet)
         if ai_message.tool_calls:
-            logger.debug(
-                f"Guardrail '{self.guardrail_name}' skipping evaluation - "
-                "AI message contains tool calls, waiting for final response"
+            logger.trace(
+                "Guardrail skipping evaluation - AI message contains tool calls",
+                guardrail_name=self.guardrail_name,
             )
             return None
 
         # Skip evaluation if the AI message has no content to evaluate
         if not ai_message.content:
-            logger.debug(
-                f"Guardrail '{self.guardrail_name}' skipping evaluation - "
-                "AI message has no content"
+            logger.trace(
+                "Guardrail skipping evaluation - AI message has no content",
+                guardrail_name=self.guardrail_name,
             )
             return None
 
-        logger.debug(f"Evaluating response with guardrail '{self.guardrail_name}'")
-
         # Extract text content from messages (handles both string and structured content)
         human_content = _extract_text_content(human_message)
         ai_content = _extract_text_content(ai_message)
 
         logger.debug(
-            f"Guardrail '{self.guardrail_name}' evaluating: "
-            f"input_length={len(human_content)}, output_length={len(ai_content)}"
+            "Evaluating response with guardrail",
+            guardrail_name=self.guardrail_name,
+            input_length=len(human_content),
+            output_length=len(ai_content),
         )
 
         evaluator = create_llm_as_judge(
@@ -150,8 +150,11 @@ class GuardrailMiddleware(AgentMiddleware[AgentState, Context]):
         eval_result = evaluator(inputs=human_content, outputs=ai_content)
 
         if eval_result["score"]:
-            logger.debug(f"Response approved by guardrail '{self.guardrail_name}'")
-            logger.debug(f"Judge's comment: {eval_result['comment']}")
+            logger.debug(
+                "Response approved by guardrail",
+                guardrail_name=self.guardrail_name,
+                comment=eval_result["comment"],
+            )
             self._retry_count = 0
             return None
         else:
@@ -160,10 +163,12 @@ class GuardrailMiddleware(AgentMiddleware[AgentState, Context]):
 
             if self._retry_count >= self.num_retries:
                 logger.warning(
-                    f"Guardrail '{self.guardrail_name}' failed - max retries reached "
-                    f"({self._retry_count}/{self.num_retries})"
+                    "Guardrail failed - max retries reached",
+                    guardrail_name=self.guardrail_name,
+                    retry_count=self._retry_count,
+                    max_retries=self.num_retries,
+                    critique=comment,
                 )
-                logger.warning(f"Final judge's critique: {comment}")
                 self._retry_count = 0
 
                 # Add system message to inform user of guardrail failure
@@ -177,10 +182,12 @@ class GuardrailMiddleware(AgentMiddleware[AgentState, Context]):
                 return {"messages": [AIMessage(content=failure_message)]}
 
             logger.warning(
-                f"Guardrail '{self.guardrail_name}' requested improvements "
-                f"(retry {self._retry_count}/{self.num_retries})"
+                "Guardrail requested improvements",
+                guardrail_name=self.guardrail_name,
+                retry=self._retry_count,
+                max_retries=self.num_retries,
+                critique=comment,
             )
-            logger.warning(f"Judge's critique: {comment}")
 
             content: str = "\n".join([str(human_message.content), comment])
             return {"messages": [HumanMessage(content=content)]}
@@ -250,9 +257,7 @@ class ContentFilterMiddleware(AgentMiddleware[AgentState, Context]):
 
         for keyword in self.banned_keywords:
             if keyword in content:
-                logger.warning(
-                    f"Content filter blocked response containing '{keyword}'"
-                )
+                logger.warning("Content filter blocked response", keyword=keyword)
                 # Modify the last message content
                 last_message.content = self.block_message
                 return None
@@ -347,7 +352,7 @@ def create_guardrail_middleware(
             num_retries=2,
         )
     """
-    logger.debug(f"Creating guardrail middleware: {name}")
+    logger.trace("Creating guardrail middleware", guardrail_name=name)
     return GuardrailMiddleware(
         name=name,
         model=model,
@@ -379,8 +384,8 @@ def create_content_filter_middleware(
             block_message="I cannot discuss sensitive credentials.",
         )
     """
-    logger.debug(
-        f"Creating content filter middleware with {len(banned_keywords)} keywords"
+    logger.trace(
+        "Creating content filter middleware", keywords_count=len(banned_keywords)
     )
     return ContentFilterMiddleware(
         banned_keywords=banned_keywords,
@@ -411,5 +416,5 @@ def create_safety_guardrail_middleware(
             safety_model=ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct"),
         )
     """
-    logger.debug("Creating safety guardrail middleware")
+    logger.trace("Creating safety guardrail middleware")
     return SafetyGuardrailMiddleware(safety_model=safety_model)

dao_ai/middleware/human_in_the_loop.py

@@ -103,7 +103,10 @@ def _config_to_interrupt_on_entry(
             interrupt_entry["description"] = config.review_prompt
         return interrupt_entry
 
-    logger.warning(f"Unknown HITL config type: {type(config)}, defaulting to True")
+    logger.warning(
+        "Unknown HITL config type, defaulting to True",
+        config_type=type(config).__name__,
+    )
     return True
 
 
@@ -152,8 +155,9 @@ def create_human_in_the_loop_middleware(
             normalized_interrupt_on[tool_name] = config
 
     logger.debug(
-        f"Creating HITL middleware for {len(normalized_interrupt_on)} tools: "
-        f"{list(normalized_interrupt_on.keys())}"
+        "Creating HITL middleware",
+        tools_count=len(normalized_interrupt_on),
+        tools=list(normalized_interrupt_on.keys()),
     )
 
     return HumanInTheLoopMiddleware(
@@ -216,10 +220,10 @@ def create_hitl_middleware_from_tool_models(
             tool_name: str | None = getattr(func_tool, "name", None)
             if tool_name:
                 interrupt_on[tool_name] = hitl_config
-                logger.debug(f"Tool '{tool_name}' configured for HITL")
+                logger.trace("Tool configured for HITL", tool_name=tool_name)
 
     if not interrupt_on:
-        logger.debug("No tools require HITL - returning None")
+        logger.trace("No tools require HITL - returning None")
         return None
 
     return create_human_in_the_loop_middleware(

dao_ai/middleware/message_validation.py

@@ -50,7 +50,7 @@ class MessageValidationMiddleware(AgentMiddleware[AgentState, Context]):
         try:
             return self.validate(state, runtime)
         except ValueError as e:
-            logger.error(f"Message validation failed: {e}")
+            logger.error("Message validation failed", error=str(e))
             return {
                 "is_valid": False,
                 "message_error": str(e),
@@ -93,20 +93,28 @@ class UserIdValidationMiddleware(MessageValidationMiddleware):
         self, state: AgentState, runtime: Runtime[Context]
     ) -> dict[str, Any] | None:
         """Validate user_id is present and properly formatted."""
-        logger.debug("Executing user_id validation")
+        logger.trace("Executing user_id validation")
 
         context: Context = runtime.context or Context()
         user_id: str | None = context.user_id
 
         if not user_id:
-            logger.error("User ID is required but not provided in the configuration.")
+            logger.error("User ID is required but not provided in configuration")
 
             thread_val = context.thread_id or "<your_thread_id>"
+            # Get extra fields from context (excluding user_id and thread_id)
+            context_dict = context.model_dump()
+            extra_fields = {
+                k: v
+                for k, v in context_dict.items()
+                if k not in {"user_id", "thread_id"} and v is not None
+            }
+
             corrected_config: dict[str, Any] = {
                 "configurable": {
                     "thread_id": thread_val,
                     "user_id": "<your_user_id>",
-                    **context.custom,
+                    **extra_fields,
                 },
                 "session": {
                     "conversation_id": thread_val,
@@ -138,15 +146,23 @@ Please update your configuration and try again.
             raise ValueError(error_message)
 
         if "." in user_id:
-            logger.error(f"User ID '{user_id}' contains invalid character '.'")
+            logger.error("User ID contains invalid character '.'", user_id=user_id)
 
             corrected_user_id = user_id.replace(".", "_")
             thread_val = context.thread_id or "<your_thread_id>"
+            # Get extra fields from context (excluding user_id and thread_id)
+            context_dict = context.model_dump()
+            extra_fields = {
+                k: v
+                for k, v in context_dict.items()
+                if k not in {"user_id", "thread_id"} and v is not None
+            }
+
             corrected_config: dict[str, Any] = {
                 "configurable": {
                     "thread_id": thread_val,
                     "user_id": corrected_user_id,
-                    **context.custom,
+                    **extra_fields,
                 },
                 "session": {
                     "conversation_id": thread_val,
@@ -183,19 +199,27 @@ class ThreadIdValidationMiddleware(MessageValidationMiddleware):
         self, state: AgentState, runtime: Runtime[Context]
     ) -> dict[str, Any] | None:
         """Validate thread_id/conversation_id is present."""
-        logger.debug("Executing thread_id/conversation_id validation")
+        logger.trace("Executing thread_id/conversation_id validation")
 
         context: Context = runtime.context or Context()
         thread_id: str | None = context.thread_id
 
         if not thread_id:
-            logger.error("Thread ID / Conversation ID is required but not provided.")
+            logger.error("Thread ID / Conversation ID is required but not provided")
+
+            # Get extra fields from context (excluding user_id and thread_id)
+            context_dict = context.model_dump()
+            extra_fields = {
+                k: v
+                for k, v in context_dict.items()
+                if k not in {"user_id", "thread_id"} and v is not None
+            }
 
             corrected_config: dict[str, Any] = {
                 "configurable": {
                     "thread_id": "<your_thread_id>",
                     "user_id": context.user_id or "<your_user_id>",
-                    **context.custom,
+                    **extra_fields,
                 },
                 "session": {
                     "conversation_id": "<your_thread_id>",
@@ -269,7 +293,7 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
     Middleware that validates the presence of required custom fields.
 
     This is a generic validation middleware that can check for multiple
-    required fields in context.custom.
+    required fields in the context object.
 
     Fields are defined in the `fields` list. Each field can have:
     - name: The field name (required)
@@ -312,7 +336,7 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
                 <field_name>: <example_value>
             session: {}
         """
-        logger.debug("Executing custom field validation")
+        logger.trace("Executing custom field validation")
 
         context: Context = runtime.context or Context()
 
@@ -320,7 +344,7 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
         missing_fields: list[RequiredField] = []
         for field in self.fields:
             if field.is_required:
-                field_value: Any = context.custom.get(field.name)
+                field_value: Any = getattr(context, field.name, None)
                 if field_value is None:
                     missing_fields.append(field)
 
@@ -329,7 +353,7 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
 
         # Log the missing fields
         missing_names = [f.name for f in missing_fields]
-        logger.error(f"Required fields missing: {', '.join(missing_names)}")
+        logger.error("Required fields missing", fields=missing_names)
 
         # Build the configurable dict preserving provided values
         # and using example_value for missing required fields
@@ -344,9 +368,11 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
         else:
             configurable["user_id"] = "<your_user_id>"
 
-        # Add all values the user already provided in custom
-        for k, v in context.custom.items():
-            configurable[k] = v
+        # Add all extra values the user already provided
+        context_dict = context.model_dump()
+        for k, v in context_dict.items():
+            if k not in {"user_id", "thread_id"} and v is not None:
+                configurable[k] = v
 
         # Then add our defined fields (provided values take precedence)
         for field in self.fields:
@@ -380,8 +406,16 @@ class CustomFieldValidationMiddleware(MessageValidationMiddleware):
         field_descriptions: list[str] = [
             "- **thread_id**: Thread identifier (required in configurable)",
             "- **conversation_id**: Alias of thread_id (in session)",
-            "- **user_id**: Your unique user identifier (required)",
         ]
+
+        # Add user_id if not in custom fields
+        has_user_id_field = any(f.name == "user_id" for f in self.fields)
+        if not has_user_id_field:
+            field_descriptions.append(
+                "- **user_id**: Your unique user identifier (required)"
+            )
+
+        # Add custom field descriptions
         for field in self.fields:
             required_text = "(required)" if field.is_required else "(optional)"
             field_descriptions.append(
@@ -427,22 +461,22 @@ class FilterLastHumanMessageMiddleware(AgentMiddleware[AgentState, Context]):
         self, state: AgentState, runtime: Runtime[Context]
     ) -> dict[str, Any] | None:
         """Filter messages to keep only the last human message."""
-        logger.debug("Executing filter_last_human_message middleware")
+        logger.trace("Executing filter_last_human_message middleware")
 
         messages: list[BaseMessage] = state.get("messages", [])
 
         if not messages:
-            logger.debug("No messages found in state")
+            logger.trace("No messages found in state")
             return None
 
         last_message: HumanMessage | None = last_human_message(messages)
 
         if last_message is None:
-            logger.debug("No human messages found in state")
+            logger.trace("No human messages found in state")
             return {"messages": []}
 
-        logger.debug(
-            f"Filtered {len(messages)} messages down to 1 (last human message)"
+        logger.trace(
+            "Filtered messages to last human message", original_count=len(messages)
         )
 
         removed_messages = [
@@ -472,7 +506,7 @@ def create_user_id_validation_middleware() -> UserIdValidationMiddleware:
     Example:
         middleware = create_user_id_validation_middleware()
     """
-    logger.debug("Creating user_id validation middleware")
+    logger.trace("Creating user_id validation middleware")
     return UserIdValidationMiddleware()
 
 
@@ -489,7 +523,7 @@ def create_thread_id_validation_middleware() -> ThreadIdValidationMiddleware:
     Example:
         middleware = create_thread_id_validation_middleware()
     """
-    logger.debug("Creating thread_id validation middleware")
+    logger.trace("Creating thread_id validation middleware")
     return ThreadIdValidationMiddleware()
 
 
@@ -500,7 +534,7 @@ def create_custom_field_validation_middleware(
     Create a CustomFieldValidationMiddleware instance.
 
     Factory function for creating middleware that validates the presence
-    of required custom fields in context.custom.
+    of required custom fields in the context object.
 
     Each field in the list should have:
     - name: The field name (required)
@@ -530,9 +564,7 @@ def create_custom_field_validation_middleware(
         )
     """
     field_names = [f.get("name", "unknown") for f in fields]
-    logger.debug(
-        f"Creating custom field validation middleware for fields: {field_names}"
-    )
+    logger.trace("Creating custom field validation middleware", fields=field_names)
     return CustomFieldValidationMiddleware(fields=fields)
 
 
@@ -550,5 +582,5 @@ def create_filter_last_human_message_middleware() -> FilterLastHumanMessageMiddl
     Example:
         middleware = create_filter_last_human_message_middleware()
     """
-    logger.debug("Creating filter_last_human_message middleware")
+    logger.trace("Creating filter_last_human_message middleware")
     return FilterLastHumanMessageMiddleware()

dao_ai/middleware/summarization.py

@@ -76,15 +76,19 @@ class LoggingSummarizationMiddleware(SummarizationMiddleware):
         summarized_count = original_message_count - preserved_count
 
         logger.info(
-            f"Conversation summarized: "
-            f"BEFORE: {original_message_count} messages (~{original_token_count:,} tokens) → "
-            f"AFTER: {new_message_count} messages (~{new_token_count:,} tokens) | "
-            f"{summarized_count} messages condensed into 1 summary"
+            "Conversation summarized",
+            before_messages=original_message_count,
+            before_tokens=original_token_count,
+            after_messages=new_message_count,
+            after_tokens=new_token_count,
+            summarized_messages=summarized_count,
         )
         logger.debug(
-            f"Summarization details: trigger={self.trigger}, keep={self.keep}, "
-            f"preserved_messages={preserved_count}, "
-            f"token_reduction={original_token_count - new_token_count:,}"
+            "Summarization details",
+            trigger=self.trigger,
+            keep=self.keep,
+            preserved_messages=preserved_count,
+            token_reduction=original_token_count - new_token_count,
         )
 
     def _is_remove_message(self, msg: Any) -> bool:
@@ -160,9 +164,10 @@ def create_summarization_middleware(
         middleware = create_summarization_middleware(chat_history)
     """
     logger.debug(
-        f"Creating summarization middleware with max_tokens: {chat_history.max_tokens}, "
-        f"max_tokens_before_summary: {chat_history.max_tokens_before_summary}, "
-        f"max_messages_before_summary: {chat_history.max_messages_before_summary}"
+        "Creating summarization middleware",
+        max_tokens=chat_history.max_tokens,
+        max_tokens_before_summary=chat_history.max_tokens_before_summary,
+        max_messages_before_summary=chat_history.max_messages_before_summary,
     )
 
     # Get the LLM model
@@ -183,7 +188,7 @@ def create_summarization_middleware(
     # Default to keeping enough for context
     keep: Tuple[str, int] = ("tokens", chat_history.max_tokens)
 
-    logger.info(f"Summarization middleware configured: trigger={trigger}, keep={keep}")
+    logger.info("Summarization middleware configured", trigger=trigger, keep=keep)
 
     return LoggingSummarizationMiddleware(
         model=model,