coco-framework 0.1.0__py3-none-any.whl

coco_graph/prebuilt/tool_validator.py

@@ -0,0 +1,221 @@
+"""This module provides a ValidationNode class that can be used to validate tool calls
+in a coco graph. It applies a pydantic schema to tool_calls in the models' outputs,
+and returns a ToolMessage with the validated content. If the schema is not valid, it
+returns a ToolMessage with the error message. The ValidationNode can be used in a
+StateGraph with a "messages" key. If multiple tool calls are requested, they will be run in parallel.
+"""
+
+from collections.abc import Callable, Sequence
+from typing import (
+    Any,
+    cast,
+)
+
+from coco_core.messages import (
+    AIMessage,
+    AnyMessage,
+    ToolCall,
+    ToolMessage,
+)
+from coco_core.runnables import (
+    RunnableConfig,
+)
+from coco_core.runnables.config import get_executor_for_config
+from coco_core.tools import BaseTool, create_schema_from_function
+from coco_core.utils.pydantic import is_basemodel_subclass
+from coco_graph._internal._runnable import RunnableCallable
+from coco_graph.warnings import CocoGraphDeprecatedSinceV10
+from pydantic import BaseModel, ValidationError
+from pydantic.v1 import BaseModel as BaseModelV1
+from pydantic.v1 import ValidationError as ValidationErrorV1
+from typing_extensions import deprecated
+
+
+def _default_format_error(
+    error: BaseException,
+    call: ToolCall,
+    schema: type[BaseModel] | type[BaseModelV1],
+) -> str:
+    """Default error formatting function."""
+    return f"{repr(error)}\n\nRespond after fixing all validation errors."
+
+
+@deprecated(
+    "ValidationNode is deprecated. Please use `create_agent` from `coco.agents` with custom tool error handling.",
+    category=CocoGraphDeprecatedSinceV10,
+)
+class ValidationNode(RunnableCallable):
+    """A node that validates all tools requests from the last `AIMessage`.
+
+    It can be used either in `StateGraph` with a `'messages'` key.
+
+    !!! note
+
+        This node does not actually **run** the tools, it only validates the tool calls,
+        which is useful for extraction and other use cases where you need to generate
+        structured output that conforms to a complex schema without losing the original
+        messages and tool IDs (for use in multi-turn conversations).
+
+    Returns:
+        (Union[Dict[str, List[ToolMessage]], Sequence[ToolMessage]]): A list of
+            `ToolMessage` objects with the validated content or error messages.
+
+    Example:
+        ```python title="Example usage for re-prompting the model to generate a valid response:"
+        from typing import Literal, Annotated
+        from typing_extensions import TypedDict
+
+        from coco_anthropic import ChatAnthropic
+        from pydantic import BaseModel, field_validator
+
+        from coco_graph.graph import END, START, StateGraph
+        from coco_graph.prebuilt import ValidationNode
+        from coco_graph.graph.message import add_messages
+
+        class SelectNumber(BaseModel):
+            a: int
+
+            @field_validator("a")
+            def a_must_be_meaningful(cls, v):
+                if v != 37:
+                    raise ValueError("Only 37 is allowed")
+                return v
+
+        builder = StateGraph(Annotated[list, add_messages])
+        llm = ChatAnthropic(model="claude-3-5-haiku-latest").bind_tools([SelectNumber])
+        builder.add_node("model", llm)
+        builder.add_node("validation", ValidationNode([SelectNumber]))
+        builder.add_edge(START, "model")
+
+        def should_validate(state: list) -> Literal["validation", "__end__"]:
+            if state[-1].tool_calls:
+                return "validation"
+            return END
+
+        builder.add_conditional_edges("model", should_validate)
+
+        def should_reprompt(state: list) -> Literal["model", "__end__"]:
+            for msg in state[::-1]:
+                # None of the tool calls were errors
+                if msg.type == "ai":
+                    return END
+                if msg.additional_kwargs.get("is_error"):
+                    return "model"
+            return END
+
+        builder.add_conditional_edges("validation", should_reprompt)
+
+        graph = builder.compile()
+        res = graph.invoke(("user", "Select a number, any number"))
+        # Show the retry logic
+        for msg in res:
+            msg.pretty_print()
+        ```
+    """
+
+    def __init__(
+        self,
+        schemas: Sequence[BaseTool | type[BaseModel] | Callable],
+        *,
+        format_error: Callable[[BaseException, ToolCall, type[BaseModel]], str]
+        | None = None,
+        name: str = "validation",
+        tags: list[str] | None = None,
+    ) -> None:
+        """Initialize the ValidationNode.
+
+        Args:
+            schemas: A list of schemas to validate the tool calls with. These can be
+                any of the following:
+                - A pydantic BaseModel class
+                - A BaseTool instance (the args_schema will be used)
+                - A function (a schema will be created from the function signature)
+            format_error: A function that takes an exception, a ToolCall, and a schema
+                and returns a formatted error string. By default, it returns the
+                exception repr and a message to respond after fixing validation errors.
+            name: The name of the node.
+            tags: A list of tags to add to the node.
+        """
+        super().__init__(self._func, None, name=name, tags=tags, trace=False)
+        self._format_error = format_error or _default_format_error
+        self.schemas_by_name: dict[str, type[BaseModel]] = {}
+        for schema in schemas:
+            if isinstance(schema, BaseTool):
+                if schema.args_schema is None:
+                    raise ValueError(
+                        f"Tool {schema.name} does not have an args_schema defined."
+                    )
+                elif not isinstance(
+                    schema.args_schema, type
+                ) or not is_basemodel_subclass(schema.args_schema):
+                    raise ValueError(
+                        "Validation node only works with tools that have a pydantic BaseModel args_schema. "
+                        f"Got {schema.name} with args_schema: {schema.args_schema}."
+                    )
+                self.schemas_by_name[schema.name] = schema.args_schema
+            elif isinstance(schema, type) and issubclass(
+                schema, (BaseModel, BaseModelV1)
+            ):
+                self.schemas_by_name[schema.__name__] = cast(type[BaseModel], schema)
+            elif callable(schema):
+                base_model = create_schema_from_function("Validation", schema)
+                self.schemas_by_name[schema.__name__] = base_model
+            else:
+                raise ValueError(
+                    f"Unsupported input to ValidationNode. Expected BaseModel, tool or function. Got: {type(schema)}."
+                )
+
+    def _get_message(
+        self, input: list[AnyMessage] | dict[str, Any]
+    ) -> tuple[str, AIMessage]:
+        """Extract the last AIMessage from the input."""
+        if isinstance(input, list):
+            output_type = "list"
+            messages: list = input
+        elif messages := input.get("messages", []):
+            output_type = "dict"
+        else:
+            raise ValueError("No message found in input")
+        message: AnyMessage = messages[-1]
+        if not isinstance(message, AIMessage):
+            raise ValueError("Last message is not an AIMessage")
+        return output_type, message
+
+    def _func(
+        self, input: list[AnyMessage] | dict[str, Any], config: RunnableConfig
+    ) -> Any:
+        """Validate and run tool calls synchronously."""
+        output_type, message = self._get_message(input)
+
+        def run_one(call: ToolCall) -> ToolMessage:
+            schema = self.schemas_by_name[call["name"]]
+            try:
+                if issubclass(schema, BaseModel):
+                    output = schema.model_validate(call["args"])
+                    content = output.model_dump_json()
+                elif issubclass(schema, BaseModelV1):
+                    output = schema.validate(call["args"])
+                    content = output.json()
+                else:
+                    raise ValueError(
+                        f"Unsupported schema type: {type(schema)}. Expected BaseModel or BaseModelV1."
+                    )
+                return ToolMessage(
+                    content=content,
+                    name=call["name"],
+                    tool_call_id=cast(str, call["id"]),
+                )
+            except (ValidationError, ValidationErrorV1) as e:
+                return ToolMessage(
+                    content=self._format_error(e, call, schema),
+                    name=call["name"],
+                    tool_call_id=cast(str, call["id"]),
+                    additional_kwargs={"is_error": True},
+                )
+
+        with get_executor_for_config(config) as executor:
+            outputs = [*executor.map(run_one, message.tool_calls)]
+            if output_type == "list":
+                return outputs
+            else:
+                return {"messages": outputs}

coco_graph/pregel/__init__.py

@@ -0,0 +1,3 @@
+from coco_graph.pregel.main import NodeBuilder, Pregel
+
+__all__ = ("Pregel", "NodeBuilder")