dao-ai 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

dao_ai/cli.py

@@ -3,6 +3,7 @@ import json
 import os
 import subprocess
 import sys
+import traceback
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
 from typing import Optional, Sequence
@@ -388,7 +389,9 @@ def handle_chat_command(options: Namespace) -> None:
 
                 except Exception as e:
                     print(f"\n❌ Error during streaming: {e}")
+                    print(f"Stack trace:\n{traceback.format_exc()}")
                     logger.error(f"Streaming error: {e}")
+                    logger.error(f"Stack trace: {traceback.format_exc()}")
 
             except EOFError:
                 # Handle Ctrl-D

dao_ai/config.py

@@ -43,8 +43,6 @@ from mlflow.models.resources import (
 )
 from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator
 
-from dao_ai.chat_models import ChatDatabricksFiltered
-
 
 class HasValue(ABC):
     @abstractmethod
@@ -275,9 +273,14 @@ class LLMModel(BaseModel, IsDatabricksResource):
         # chat_client: LanguageModelLike = self.as_open_ai_client()
 
         # Create ChatDatabricksWrapper instance directly
+        from dao_ai.chat_models import ChatDatabricksFiltered
+
         chat_client: LanguageModelLike = ChatDatabricksFiltered(
             model=self.name, temperature=self.temperature, max_tokens=self.max_tokens
         )
+        # chat_client: LanguageModelLike = ChatDatabricks(
+        #     model=self.name, temperature=self.temperature, max_tokens=self.max_tokens
+        # )
 
         fallbacks: Sequence[LanguageModelLike] = []
         for fallback in self.fallbacks:
@@ -1001,7 +1004,15 @@ class ChatHistoryModel(BaseModel):
     max_tokens: int = 256
     max_tokens_before_summary: Optional[int] = None
     max_messages_before_summary: Optional[int] = None
-    max_summary_tokens: Optional[int] = None
+    max_summary_tokens: int = 255
+
+    @model_validator(mode="after")
+    def validate_max_summary_tokens(self) -> "ChatHistoryModel":
+        if self.max_summary_tokens >= self.max_tokens:
+            raise ValueError(
+                f"max_summary_tokens ({self.max_summary_tokens}) must be less than max_tokens ({self.max_tokens})"
+            )
+        return self
 
 
 class AppModel(BaseModel):
@@ -1057,12 +1068,19 @@ class AppModel(BaseModel):
         return self
 
 
+class GuidelineModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str
+    guidelines: list[str]
+
+
 class EvaluationModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     model: LLMModel
     table: TableModel
     num_evals: int
     custom_inputs: dict[str, Any] = Field(default_factory=dict)
+    guidelines: list[GuidelineModel] = Field(default_factory=list)
 
 
 class DatasetFormat(str, Enum):

dao_ai/graph.py

@@ -3,13 +3,10 @@ from typing import Sequence
 from langchain_core.language_models import LanguageModelLike
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
-from langgraph.cache.base import BaseCache
-from langgraph.cache.memory import InMemoryCache
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import END, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.store.base import BaseStore
-from langgraph.types import CachePolicy
 from langgraph_supervisor import create_handoff_tool as supervisor_handoff_tool
 from langgraph_supervisor import create_supervisor
 from langgraph_swarm import create_handoff_tool as swarm_handoff_tool
@@ -29,7 +26,7 @@ from dao_ai.nodes import (
     message_hook_node,
 )
 from dao_ai.prompts import make_prompt
-from dao_ai.state import IncomingState, OutgoingState, SharedState
+from dao_ai.state import Context, IncomingState, OutgoingState, SharedState
 
 
 def route_message(state: SharedState) -> str:
@@ -113,9 +110,6 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
         checkpointer = orchestration.memory.checkpointer.as_checkpointer()
         logger.debug(f"Using checkpointer: {checkpointer}")
 
-    cache: BaseCache = None
-    cache = InMemoryCache()
-
     prompt: str = supervisor.prompt
 
     model: LanguageModelLike = supervisor.model.as_chat_model()
@@ -127,22 +121,29 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
         tools=tools,
         state_schema=SharedState,
         config_schema=RunnableConfig,
+        output_mode="last_message",
+        add_handoff_messages=False,
+        add_handoff_back_messages=False,
+        context_schema=Context,
+        # output_mode="full",
+        # add_handoff_messages=True,
+        # add_handoff_back_messages=True,
     )
 
-    supervisor_node: CompiledStateGraph = supervisor_workflow.compile()
+    supervisor_node: CompiledStateGraph = supervisor_workflow.compile(
+        checkpointer=checkpointer, store=store
+    )
 
     workflow: StateGraph = StateGraph(
         SharedState,
-        config_schema=RunnableConfig,
         input=IncomingState,
         output=OutgoingState,
+        context_schema=Context,
     )
 
     workflow.add_node("message_hook", message_hook_node(config=config))
 
-    workflow.add_node(
-        "orchestration", supervisor_node, cache_policy=CachePolicy(ttl=60)
-    )
+    workflow.add_node("orchestration", supervisor_node)
     workflow.add_conditional_edges(
         "message_hook",
         route_message,
@@ -153,7 +154,7 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
     )
     workflow.set_entry_point("message_hook")
 
-    return workflow.compile(checkpointer=checkpointer, store=store, cache=cache)
+    return workflow.compile(checkpointer=checkpointer, store=store)
 
 
 def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
@@ -172,6 +173,16 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
     orchestration: OrchestrationModel = config.app.orchestration
     swarm: SwarmModel = orchestration.swarm
 
+    store: BaseStore = None
+    if orchestration.memory and orchestration.memory.store:
+        store = orchestration.memory.store.as_store()
+        logger.debug(f"Using memory store: {store}")
+
+    checkpointer: BaseCheckpointSaver = None
+    if orchestration.memory and orchestration.memory.checkpointer:
+        checkpointer = orchestration.memory.checkpointer.as_checkpointer()
+        logger.debug(f"Using checkpointer: {checkpointer}")
+
     default_agent: AgentModel = swarm.default_agent
     if isinstance(default_agent, AgentModel):
         default_agent = default_agent.name
@@ -180,20 +191,22 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
         agents=agents,
         default_active_agent=default_agent,
         state_schema=SharedState,
-        config_schema=RunnableConfig,
+        context_schema=Context,
     )
 
-    swarm_node: CompiledStateGraph = swarm_workflow.compile()
+    swarm_node: CompiledStateGraph = swarm_workflow.compile(
+        checkpointer=checkpointer, store=store
+    )
 
     workflow: StateGraph = StateGraph(
         SharedState,
-        config_schema=RunnableConfig,
         input=IncomingState,
         output=OutgoingState,
+        context_schema=Context,
     )
 
     workflow.add_node("message_hook", message_hook_node(config=config))
-    workflow.add_node("orchestration", swarm_node, cache_policy=CachePolicy(ttl=60))
+    workflow.add_node("orchestration", swarm_node)
 
     workflow.add_conditional_edges(
         "message_hook",
@@ -206,20 +219,7 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
 
     workflow.set_entry_point("message_hook")
 
-    store: BaseStore = None
-    if orchestration.memory and orchestration.memory.store:
-        store = orchestration.memory.store.as_store()
-        logger.debug(f"Using memory store: {store}")
-
-    checkpointer: BaseCheckpointSaver = None
-    if orchestration.memory and orchestration.memory.checkpointer:
-        checkpointer = orchestration.memory.checkpointer.as_checkpointer()
-        logger.debug(f"Using checkpointer: {checkpointer}")
-
-    cache: BaseCache = None
-    cache = InMemoryCache()
-
-    return workflow.compile(checkpointer=checkpointer, store=store, cache=cache)
+    return workflow.compile(checkpointer=checkpointer, store=store)
 
 
 def create_dao_ai_graph(config: AppConfig) -> CompiledStateGraph:

dao_ai/hooks/__init__.py

@@ -1,5 +1,6 @@
 from dao_ai.hooks.core import (
     create_hooks,
+    filter_last_human_message_hook,
     null_hook,
     null_initialization_hook,
     null_shutdown_hook,
@@ -14,4 +15,5 @@ __all__ = [
     "null_shutdown_hook",
     "require_thread_id_hook",
     "require_user_id_hook",
+    "filter_last_human_message_hook",
 ]

dao_ai/hooks/core.py

@@ -1,9 +1,13 @@
 import json
 from typing import Any, Callable, Sequence
 
+from langchain_core.messages import BaseMessage, HumanMessage, RemoveMessage
+from langgraph.runtime import Runtime
 from loguru import logger
 
 from dao_ai.config import AppConfig, FunctionHook, PythonFunctionModel
+from dao_ai.messages import last_human_message
+from dao_ai.state import Context
 
 
 def create_hooks(
@@ -23,7 +27,7 @@ def create_hooks(
     return hooks
 
 
-def null_hook(state: dict[str, Any], config: dict[str, Any]) -> dict[str, Any]:
+def null_hook(state: dict[str, Any], runtime: Runtime[Context]) -> dict[str, Any]:
     logger.debug("Executing null hook")
     return {}
 
@@ -36,19 +40,79 @@ def null_shutdown_hook(config: AppConfig) -> None:
     logger.debug("Executing null shutdown hook")
 
 
+def filter_last_human_message_hook(
+    state: dict[str, Any], runtime: Runtime[Context]
+) -> dict[str, Any]:
+    """
+    Filter messages to keep only the last human message.
+
+    This hook removes all messages except for the most recent human message,
+    which can be useful for scenarios where you want to process only the
+    latest user input without conversation history.
+
+    Args:
+        state: The current state containing messages
+        runtime: The runtime context (unused in this hook)
+
+    Returns:
+        Updated state with filtered messages
+    """
+    logger.debug("Executing filter_last_human_message hook")
+
+    messages: list[BaseMessage] = state.get("messages", [])
+
+    if not messages:
+        logger.debug("No messages found in state")
+        return state
+
+    # Use the helper function to find the last human message
+    last_message: HumanMessage = last_human_message(messages)
+
+    if last_message is None:
+        logger.debug("No human messages found in state")
+        # Return empty messages if no human message found
+        updated_state = state.copy()
+        updated_state["messages"] = []
+        return updated_state
+
+    logger.debug(f"Filtered {len(messages)} messages down to 1 (last human message)")
+
+    removed_messages: Sequence[BaseMessage] = [
+        RemoveMessage(id=message.id)
+        for message in messages
+        if message.id != last_message.id
+    ]
+
+    updated_state: dict[str, Sequence[BaseMessage]] = {"messages": removed_messages}
+
+    return updated_state
+
+
 def require_user_id_hook(
-    state: dict[str, Any], config: dict[str, Any]
+    state: dict[str, Any], runtime: Runtime[Context]
 ) -> dict[str, Any]:
     logger.debug("Executing user_id validation hook")
 
-    config = config.get("custom_inputs", config)
+    context: Context = runtime.context or Context()
 
-    configurable: dict[str, Any] = config.get("configurable", {})
+    user_id: str | None = context.user_id
 
-    if "user_id" not in configurable or not configurable["user_id"]:
+    if not user_id:
         logger.error("User ID is required but not provided in the configuration.")
 
-        error_message = """
+        # Create corrected configuration using any provided context parameters
+        corrected_config = {
+            "configurable": {
+                "thread_id": context.thread_id or "1",
+                "user_id": "my_user_id",
+                "store_num": context.store_num or 87887,
+            }
+        }
+
+        # Format as JSON for copy-paste
+        corrected_config_json = json.dumps(corrected_config, indent=2)
+
+        error_message = f"""
 ## Authentication Required
 
 A **user_id** is required to process your request. Please provide your user ID in the configuration.
@@ -58,25 +122,19 @@ A **user_id** is required to process your request. Please provide your user ID i
 Please include the following JSON in your request configuration:
 
 ```json
-{
-  "configurable": {
-    "thread_id": "1",
-    "user_id": "my_user_id"
-  }
-}
+{corrected_config_json}
 ```
 
 ### Field Descriptions
 - **user_id**: Your unique user identifier (required)
-- **thread_id**: Conversation thread identifier (optional)
+- **thread_id**: Conversation thread identifier (required)
+- **store_num**: Your store number (required)
 
 Please update your configuration and try again.
         """.strip()
 
         raise ValueError(error_message)
 
-    # Validate that user_id doesn't contain dots
-    user_id = configurable["user_id"]
     if "." in user_id:
         logger.error(f"User ID '{user_id}' contains invalid character '.'")
 
@@ -88,9 +146,9 @@ Please update your configuration and try again.
         # Corrected config with fixed user_id
         corrected_config = {
             "configurable": {
-                "thread_id": configurable.get("thread_id", "1"),
+                "thread_id": context.thread_id or "1",
                 "user_id": corrected_user_id,
-                "store_num": configurable.get("store_num", 87887),
+                "store_num": context.store_num or 87887,
             }
         }
 
@@ -116,38 +174,46 @@ Please update your user_id and try again.
 
 
 def require_thread_id_hook(
-    state: dict[str, Any], config: dict[str, Any]
+    state: dict[str, Any], runtime: Runtime[Context]
 ) -> dict[str, Any]:
     logger.debug("Executing thread_id validation hook")
 
-    config = config.get("custom_inputs", config)
+    context: Context = runtime.context or Context()
 
-    configurable: dict[str, Any] = config.get("configurable", {})
+    thread_id: str | None = context.thread_id
 
-    if "thread_id" not in configurable or not configurable["thread_id"]:
+    if not thread_id:
         logger.error("Thread ID is required but not provided in the configuration.")
 
-        error_message = """
+        # Create corrected configuration using any provided context parameters
+        corrected_config = {
+            "configurable": {
+                "thread_id": "1",
+                "user_id": context.user_id or "my_user_id",
+                "store_num": context.store_num or 87887,
+            }
+        }
+
+        # Format as JSON for copy-paste
+        corrected_config_json = json.dumps(corrected_config, indent=2)
+
+        error_message = f"""
 ## Authentication Required
 
-A **thread_id** is required to process your request. Please provide your user ID in the configuration.
+A **thread_id** is required to process your request. Please provide your thread ID in the configuration.
 
 ### Required Configuration Format
 
 Please include the following JSON in your request configuration:
 
 ```json
-{
-  "configurable": {
-    "thread_id": "1",
-    "user_id": "my_user_id"
-  }
-}
+{corrected_config_json}
 ```
 
 ### Field Descriptions
+- **thread_id**: Conversation thread identifier (required)
 - **user_id**: Your unique user identifier (required)
-- **thread_id**: Conversation thread identifier (optional)
+- **store_num**: Your store number (required)
 
 Please update your configuration and try again.
         """.strip()

dao_ai/memory/postgres.py

@@ -4,8 +4,8 @@ import threading
 from typing import Any, Optional
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.checkpoint.postgres import PostgresSaver
-from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.postgres import ShallowPostgresSaver
+from langgraph.checkpoint.postgres.aio import AsyncShallowPostgresSaver
 from langgraph.store.base import BaseStore
 from langgraph.store.postgres import PostgresStore
 from langgraph.store.postgres.aio import AsyncPostgresStore
@@ -141,7 +141,7 @@ class AsyncPostgresCheckpointerManager(CheckpointManagerBase):
     def __init__(self, checkpointer_model: CheckpointerModel):
         self.checkpointer_model = checkpointer_model
         self.pool: Optional[AsyncConnectionPool] = None
-        self._checkpointer: Optional[AsyncPostgresSaver] = None
+        self._checkpointer: Optional[AsyncShallowPostgresSaver] = None
         self._setup_complete = False
 
     def checkpointer(self) -> BaseCheckpointSaver:
@@ -183,7 +183,7 @@ class AsyncPostgresCheckpointerManager(CheckpointManagerBase):
             )
 
             # Create checkpointer with the shared pool
-            self._checkpointer = AsyncPostgresSaver(conn=self.pool)
+            self._checkpointer = AsyncShallowPostgresSaver(conn=self.pool)
             await self._checkpointer.setup()
 
             self._setup_complete = True
@@ -315,7 +315,7 @@ class PostgresCheckpointerManager(CheckpointManagerBase):
     def __init__(self, checkpointer_model: CheckpointerModel):
         self.checkpointer_model = checkpointer_model
         self.pool: Optional[ConnectionPool] = None
-        self._checkpointer: Optional[PostgresSaver] = None
+        self._checkpointer: Optional[ShallowPostgresSaver] = None
         self._setup_complete = False
 
     def checkpointer(self) -> BaseCheckpointSaver:
@@ -345,7 +345,7 @@ class PostgresCheckpointerManager(CheckpointManagerBase):
             self.pool = PostgresPoolManager.get_pool(self.checkpointer_model.database)
 
             # Create checkpointer with the shared pool
-            self._checkpointer = PostgresSaver(conn=self.pool)
+            self._checkpointer = ShallowPostgresSaver(conn=self.pool)
             self._checkpointer.setup()
 
             self._setup_complete = True

dao_ai/messages.py

@@ -78,6 +78,12 @@ def convert_to_langchain_messages(messages: dict[str, Any]) -> Sequence[BaseMess
     return langchain_messages
 
 
+def has_human_message(messages: BaseMessage | Sequence[BaseMessage]) -> bool:
+    if isinstance(messages, BaseMessage):
+        messages = [messages]
+    return any(isinstance(m, HumanMessage) for m in messages)
+
+
 def has_langchain_messages(messages: BaseMessage | Sequence[BaseMessage]) -> bool:
     if isinstance(messages, BaseMessage):
         messages = [messages]

dao_ai/models.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Any, Generator, Optional, Sequence
 
 from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
-from langchain_core.runnables import RunnableConfig
 from langgraph.graph.state import CompiledStateGraph
 from loguru import logger
 from mlflow import MlflowClient
@@ -20,7 +19,7 @@ from mlflow.types.llm import (
 )
 
 from dao_ai.messages import has_langchain_messages, has_mlflow_messages
-from dao_ai.state import SharedState
+from dao_ai.state import Context
 
 
 def get_latest_model_version(model_name: str) -> int:
@@ -63,10 +62,11 @@ class LanggraphChatModel(ChatModel):
 
         request = {"messages": self._convert_messages_to_dict(messages)}
 
-        config: SharedState = self._convert_to_config(params)
+        context: Context = self._convert_to_context(params)
+        custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
         response: dict[str, Sequence[BaseMessage]] = self.graph.invoke(
-            request, config=config
+            request, context=context, config=custom_inputs
         )
         logger.trace(f"response: {response}")
 
@@ -75,12 +75,9 @@ class LanggraphChatModel(ChatModel):
         response_message = ChatMessage(role="assistant", content=last_message.content)
         return ChatCompletionResponse(choices=[ChatChoice(message=response_message)])
 
-    def _convert_to_config(
+    def _convert_to_context(
         self, params: Optional[ChatParams | dict[str, Any]]
-    ) -> RunnableConfig:
-        if not params:
-            return {}
-
+    ) -> Context:
         input_data = params
         if isinstance(params, ChatParams):
             input_data = params.to_dict()
@@ -102,8 +99,8 @@ class LanggraphChatModel(ChatModel):
         if "thread_id" not in configurable:
             configurable["thread_id"] = str(uuid.uuid4())
 
-        agent_config: RunnableConfig = RunnableConfig(**{"configurable": configurable})
-        return agent_config
+        context: Context = Context(**configurable)
+        return context
 
     def predict_stream(
         self, context, messages: list[ChatMessage], params: ChatParams
@@ -114,25 +111,36 @@ class LanggraphChatModel(ChatModel):
 
         request = {"messages": self._convert_messages_to_dict(messages)}
 
-        config: SharedState = self._convert_to_config(params)
+        context: Context = self._convert_to_context(params)
+        custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
-        for message, metadata in self.graph.stream(
-            request, config=config, stream_mode="messages"
+        for nodes, stream_mode, messages_batch in self.graph.stream(
+            request,
+            context=context,
+            config=custom_inputs,
+            stream_mode=["messages", "custom"],
+            subgraphs=True,
         ):
-            logger.trace(f"message_type: {type(message)}, message: {message}")
-            if (
-                isinstance(
-                    message,
-                    (
-                        AIMessageChunk,
-                        AIMessage,
-                    ),
-                )
-                and message.content
-                and metadata["langgraph_node"] not in ["summarization"]
-            ):
-                content = message.content
-                yield self._create_chat_completion_chunk(content)
+            nodes: tuple[str, ...]
+            stream_mode: str
+            messages_batch: Sequence[BaseMessage]
+            logger.trace(
+                f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {messages_batch}"
+            )
+            for message in messages_batch:
+                if (
+                    isinstance(
+                        message,
+                        (
+                            AIMessageChunk,
+                            AIMessage,
+                        ),
+                    )
+                    and message.content
+                    and "summarization" not in nodes
+                ):
+                    content = message.content
+                    yield self._create_chat_completion_chunk(content)
 
     def _create_chat_completion_chunk(self, content: str) -> ChatCompletionChunk:
         return ChatCompletionChunk(
@@ -183,11 +191,37 @@ def _process_langchain_messages_stream(
     if isinstance(app, LanggraphChatModel):
         app = app.graph
 
-    for message, _ in app.stream(
-        {"messages": messages}, config=custom_inputs, stream_mode="messages"
+    logger.debug(f"Processing messages: {messages}, custom_inputs: {custom_inputs}")
+
+    custom_inputs = custom_inputs.get("configurable", custom_inputs or {})
+    context: Context = Context(**custom_inputs)
+
+    for nodes, stream_mode, messages in app.stream(
+        {"messages": messages},
+        context=context,
+        config=custom_inputs,
+        stream_mode=["messages", "custom"],
+        subgraphs=True,
     ):
-        message: AIMessageChunk
-        yield message
+        nodes: tuple[str, ...]
+        stream_mode: str
+        messages: Sequence[BaseMessage]
+        logger.trace(
+            f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {messages}"
+        )
+        for message in messages:
+            if (
+                isinstance(
+                    message,
+                    (
+                        AIMessageChunk,
+                        AIMessage,
+                    ),
+                )
+                and message.content
+                and "summarization" not in nodes
+            ):
+                yield message
 
 
 def _process_mlflow_messages(