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

dao_ai/chat_models.py

@@ -1,204 +0,0 @@
-import json
-from typing import Any, Iterator, Optional, Sequence
-
-from databricks_langchain import ChatDatabricks
-from langchain_core.callbacks import CallbackManagerForLLMRun
-from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
-from loguru import logger
-
-
-class ChatDatabricksFiltered(ChatDatabricks):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def _preprocess_messages(
-        self, messages: Sequence[BaseMessage]
-    ) -> Sequence[BaseMessage]:
-        logger.debug(f"Preprocessing {len(messages)} messages for filtering")
-
-        logger.trace(
-            f"Original messages:\n{json.dumps([msg.model_dump() for msg in messages], indent=2)}"
-        )
-
-        # Diagnostic logging to understand what types of messages we're getting
-        message_types = {}
-        remove_message_count = 0
-        empty_content_count = 0
-
-        for msg in messages:
-            msg_type = msg.__class__.__name__
-            message_types[msg_type] = message_types.get(msg_type, 0) + 1
-
-            if msg_type == "RemoveMessage":
-                remove_message_count += 1
-            elif hasattr(msg, "content") and (msg.content == "" or msg.content is None):
-                empty_content_count += 1
-
-        logger.debug(f"Message type breakdown: {message_types}")
-        logger.debug(
-            f"RemoveMessage count: {remove_message_count}, Empty content count: {empty_content_count}"
-        )
-
-        filtered_messages = []
-        for i, msg in enumerate(messages):
-            # First, filter out RemoveMessage objects completely - they're LangGraph-specific
-            # and should never be sent to an LLM
-            if hasattr(msg, "__class__") and msg.__class__.__name__ == "RemoveMessage":
-                logger.debug(f"Filtering out RemoveMessage at index {i}")
-                continue
-
-            # Be very conservative with filtering - only filter out messages that are:
-            # 1. Have empty or None content AND
-            # 2. Are not tool-related messages AND
-            # 3. Don't break tool_use/tool_result pairing
-            # 4. Are not the only remaining message (to avoid filtering everything)
-            has_empty_content = hasattr(msg, "content") and (
-                msg.content == "" or msg.content is None
-            )
-
-            # Check if this message has tool calls (non-empty list)
-            has_tool_calls = (
-                hasattr(msg, "tool_calls")
-                and msg.tool_calls
-                and len(msg.tool_calls) > 0
-            )
-
-            # Check if this is a tool result message
-            is_tool_result = hasattr(msg, "tool_call_id") or isinstance(
-                msg, ToolMessage
-            )
-
-            # Check if the previous message had tool calls (this message might be a tool result)
-            prev_had_tool_calls = False
-            if i > 0:
-                prev_msg = messages[i - 1]
-                prev_had_tool_calls = (
-                    hasattr(prev_msg, "tool_calls")
-                    and prev_msg.tool_calls
-                    and len(prev_msg.tool_calls) > 0
-                )
-
-            # Check if the next message is a tool result (this message might be a tool use)
-            next_is_tool_result = False
-            if i < len(messages) - 1:
-                next_msg = messages[i + 1]
-                next_is_tool_result = hasattr(next_msg, "tool_call_id") or isinstance(
-                    next_msg, ToolMessage
-                )
-
-            # Special handling for empty AIMessages - they might be placeholders or incomplete responses
-            # Don't filter them if they're the only AI response or seem important to the conversation flow
-            is_empty_ai_message = has_empty_content and isinstance(msg, AIMessage)
-
-            # Only filter out messages with empty content that are definitely not needed
-            should_filter = (
-                has_empty_content
-                and not has_tool_calls
-                and not is_tool_result
-                and not prev_had_tool_calls  # Don't filter if previous message had tool calls
-                and not next_is_tool_result  # Don't filter if next message is a tool result
-                and not (
-                    is_empty_ai_message and len(messages) <= 2
-                )  # Don't filter empty AI messages in short conversations
-            )
-
-            if should_filter:
-                logger.debug(f"Filtering out message at index {i}: {msg.model_dump()}")
-                continue
-            else:
-                filtered_messages.append(msg)
-
-        logger.debug(
-            f"Filtered {len(messages)} messages down to {len(filtered_messages)} messages"
-        )
-
-        # Log diagnostic information if all messages were filtered out
-        if len(filtered_messages) == 0:
-            logger.warning(
-                f"All {len(messages)} messages were filtered out! This indicates a problem with the conversation state."
-            )
-            logger.debug(f"Original message types: {message_types}")
-
-            if remove_message_count == len(messages):
-                logger.warning(
-                    "All messages were RemoveMessage objects - this suggests a bug in summarization logic"
-                )
-            elif empty_content_count > 0:
-                logger.debug(f"{empty_content_count} messages had empty content")
-
-        return filtered_messages
-
-    def _postprocess_message(self, message: BaseMessage) -> BaseMessage:
-        return message
-
-    def _generate(
-        self,
-        messages: Sequence[BaseMessage],
-        stop: Optional[Sequence[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        """Override _generate to apply message preprocessing and postprocessing."""
-        # Apply message preprocessing
-        processed_messages: Sequence[BaseMessage] = self._preprocess_messages(messages)
-
-        if len(processed_messages) == 0:
-            logger.error(
-                "All messages were filtered out during preprocessing. This indicates a serious issue with the conversation state."
-            )
-            empty_generation = ChatGeneration(
-                message=AIMessage(content="", id="empty-response")
-            )
-            return ChatResult(generations=[empty_generation])
-
-        logger.trace(
-            f"Processed messages:\n{json.dumps([msg.model_dump() for msg in processed_messages], indent=2)}"
-        )
-
-        result: ChatResult = super()._generate(
-            processed_messages, stop, run_manager, **kwargs
-        )
-
-        if result.generations:
-            for generation in result.generations:
-                if isinstance(generation, ChatGeneration) and generation.message:
-                    generation.message = self._postprocess_message(generation.message)
-
-        return result
-
-    def _stream(
-        self,
-        messages: Sequence[BaseMessage],
-        stop: Optional[Sequence[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[ChatGeneration]:
-        """Override _stream to apply message preprocessing and postprocessing."""
-        # Apply message preprocessing
-        processed_messages: Sequence[BaseMessage] = self._preprocess_messages(messages)
-
-        # Handle the edge case where all messages were filtered out
-        if len(processed_messages) == 0:
-            logger.error(
-                "All messages were filtered out during preprocessing. This indicates a serious issue with the conversation state."
-            )
-            # Return an empty streaming result without calling the underlying API
-            # This prevents API errors while making the issue visible through an empty response
-            empty_chunk = ChatGenerationChunk(
-                message=AIMessage(content="", id="empty-response")
-            )
-            yield empty_chunk
-            return
-
-        logger.trace(
-            f"Processed messages:\n{json.dumps([msg.model_dump() for msg in processed_messages], indent=2)}"
-        )
-
-        # Call the parent ChatDatabricks implementation
-        for chunk in super()._stream(processed_messages, stop, run_manager, **kwargs):
-            chunk: ChatGenerationChunk
-            # Apply message postprocessing to each chunk
-            if isinstance(chunk, ChatGeneration) and chunk.message:
-                chunk.message = self._postprocess_message(chunk.message)
-            yield chunk

dao_ai/guardrails.py

@@ -1,112 +0,0 @@
-from typing import Any, Literal, Optional, Type
-
-from langchain_core.language_models import LanguageModelLike
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
-from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.base import RunnableLike
-from langgraph.graph import END, START, MessagesState, StateGraph
-from langgraph.graph.state import CompiledStateGraph
-from langgraph.managed import RemainingSteps
-from loguru import logger
-from openevals.llm import create_llm_as_judge
-
-from dao_ai.config import GuardrailModel
-from dao_ai.messages import last_ai_message, last_human_message
-from dao_ai.state import SharedState
-
-
-class MessagesWithSteps(MessagesState):
-    guardrails_remaining_steps: RemainingSteps
-
-
-def end_or_reflect(state: MessagesWithSteps) -> Literal[END, "graph"]:
-    if state["guardrails_remaining_steps"] < 2:
-        return END
-    if len(state["messages"]) == 0:
-        return END
-    last_message = state["messages"][-1]
-    if isinstance(last_message, HumanMessage):
-        return "graph"
-    else:
-        return END
-
-
-def create_reflection_graph(
-    graph: CompiledStateGraph,
-    reflection: CompiledStateGraph,
-    state_schema: Optional[Type[Any]] = None,
-    config_schema: Optional[Type[Any]] = None,
-) -> StateGraph:
-    logger.debug("Creating reflection graph")
-    _state_schema = state_schema or graph.builder.schema
-
-    if "guardrails_remaining_steps" in _state_schema.__annotations__:
-        raise ValueError(
-            "Has key 'guardrails_remaining_steps' in state_schema, this shadows a built in key"
-        )
-
-    if "messages" not in _state_schema.__annotations__:
-        raise ValueError("Missing required key 'messages' in state_schema")
-
-    class StateSchema(_state_schema):
-        guardrails_remaining_steps: RemainingSteps
-
-    rgraph = StateGraph(StateSchema, config_schema=config_schema)
-    rgraph.add_node("graph", graph)
-    rgraph.add_node("reflection", reflection)
-    rgraph.add_edge(START, "graph")
-    rgraph.add_edge("graph", "reflection")
-    rgraph.add_conditional_edges("reflection", end_or_reflect)
-    return rgraph
-
-
-def with_guardrails(
-    graph: CompiledStateGraph, guardrail: CompiledStateGraph
-) -> CompiledStateGraph:
-    logger.debug("Creating graph with guardrails")
-    return create_reflection_graph(
-        graph, guardrail, state_schema=SharedState, config_schema=RunnableConfig
-    ).compile()
-
-
-def judge_node(guardrails: GuardrailModel) -> RunnableLike:
-    def judge(state: SharedState, config: RunnableConfig) -> dict[str, BaseMessage]:
-        llm: LanguageModelLike = guardrails.model.as_chat_model()
-
-        evaluator = create_llm_as_judge(
-            prompt=guardrails.prompt,
-            judge=llm,
-        )
-
-        ai_message: AIMessage = last_ai_message(state["messages"])
-        human_message: HumanMessage = last_human_message(state["messages"])
-
-        logger.debug(f"Evaluating response: {ai_message.content}")
-        eval_result = evaluator(
-            inputs=human_message.content, outputs=ai_message.content
-        )
-
-        if eval_result["score"]:
-            logger.debug("Response approved by judge")
-            logger.debug(f"Judge's comment: {eval_result['comment']}")
-            return
-        else:
-            # Otherwise, return the judge's critique as a new user message
-            logger.warning("Judge requested improvements")
-            comment: str = eval_result["comment"]
-            logger.warning(f"Judge's critique: {comment}")
-            content: str = "\n".join([human_message.content, comment])
-            return {"messages": [HumanMessage(content=content)]}
-
-    return judge
-
-
-def reflection_guardrail(guardrails: GuardrailModel) -> CompiledStateGraph:
-    judge: CompiledStateGraph = (
-        StateGraph(SharedState, config_schema=RunnableConfig)
-        .add_node("judge", judge_node(guardrails=guardrails))
-        .add_edge(START, "judge")
-        .add_edge("judge", END)
-        .compile()
-    )
-    return judge

dao_ai/tools/human_in_the_loop.py

@@ -1,100 +0,0 @@
-from typing import Any, Optional
-
-from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.base import RunnableLike
-from langchain_core.tools import BaseTool
-from langchain_core.tools import tool as create_tool
-from langgraph.prebuilt.interrupt import HumanInterrupt, HumanInterruptConfig
-from langgraph.types import interrupt
-from loguru import logger
-
-from dao_ai.config import (
-    BaseFunctionModel,
-    HumanInTheLoopModel,
-)
-
-
-def add_human_in_the_loop(
-    tool: RunnableLike,
-    *,
-    interrupt_config: HumanInterruptConfig | None = None,
-    review_prompt: Optional[str] = "Please review the tool call",
-) -> BaseTool:
-    """
-    Wrap a tool with human-in-the-loop functionality.
-    This function takes a tool (either a callable or a BaseTool instance) and wraps it
-    with a human-in-the-loop mechanism. When the tool is invoked, it will first
-    request human review before executing the tool's logic. The human can choose to
-    accept, edit the input, or provide a custom response.
-
-    Args:
-        tool (Callable[..., Any] | BaseTool): _description_
-        interrupt_config (HumanInterruptConfig | None, optional): _description_. Defaults to None.
-
-    Raises:
-        ValueError: _description_
-
-    Returns:
-        BaseTool: _description_
-    """
-    if not isinstance(tool, BaseTool):
-        tool = create_tool(tool)
-
-    if interrupt_config is None:
-        interrupt_config = {
-            "allow_accept": True,
-            "allow_edit": True,
-            "allow_respond": True,
-        }
-
-    logger.debug(f"Wrapping tool {tool} with human-in-the-loop functionality")
-
-    @create_tool(tool.name, description=tool.description, args_schema=tool.args_schema)
-    async def call_tool_with_interrupt(config: RunnableConfig, **tool_input) -> Any:
-        logger.debug(f"call_tool_with_interrupt: {tool.name} with input: {tool_input}")
-        request: HumanInterrupt = {
-            "action_request": {
-                "action": tool.name,
-                "args": tool_input,
-            },
-            "config": interrupt_config,
-            "description": review_prompt,
-        }
-
-        logger.debug(f"Human interrupt request: {request}")
-        response: dict[str, Any] = interrupt([request])[0]
-        logger.debug(f"Human interrupt response: {response}")
-
-        if response["type"] == "accept":
-            tool_response = await tool.ainvoke(tool_input, config=config)
-        elif response["type"] == "edit":
-            tool_input = response["args"]["args"]
-            tool_response = await tool.ainvoke(tool_input, config=config)
-        elif response["type"] == "response":
-            user_feedback = response["args"]
-            tool_response = user_feedback
-        else:
-            raise ValueError(f"Unknown interrupt response type: {response['type']}")
-
-        return tool_response
-
-    return call_tool_with_interrupt
-
-
-def as_human_in_the_loop(
-    tool: RunnableLike, function: BaseFunctionModel | str
-) -> RunnableLike:
-    if isinstance(function, BaseFunctionModel):
-        human_in_the_loop: HumanInTheLoopModel | None = function.human_in_the_loop
-        if human_in_the_loop:
-            # Get tool name safely - handle RunnableBinding objects
-            tool_name = getattr(tool, "name", None) or getattr(
-                getattr(tool, "bound", None), "name", "unknown_tool"
-            )
-            logger.debug(f"Adding human-in-the-loop to tool: {tool_name}")
-            tool = add_human_in_the_loop(
-                tool=tool,
-                interrupt_config=human_in_the_loop.interupt_config,
-                review_prompt=human_in_the_loop.review_prompt,
-            )
-    return tool