coco-framework 0.1.0__py3-none-any.whl

coco_core/tools/simple.py

@@ -0,0 +1,208 @@
+"""Tool that takes in function or coroutine directly."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from inspect import signature
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from typing_extensions import override
+
+# Cannot move to TYPE_CHECKING as _run/_arun parameter annotations are needed at runtime
+from coco_core.callbacks import (
+    AsyncCallbackManagerForToolRun,  # noqa: TC001
+    CallbackManagerForToolRun,  # noqa: TC001
+)
+from coco_core.runnables import RunnableConfig, run_in_executor
+from coco_core.tools.base import (
+    ArgsSchema,
+    BaseTool,
+    ToolException,
+    _get_runnable_config_param,
+)
+
+if TYPE_CHECKING:
+    from coco_core.messages import ToolCall
+
+
+class Tool(BaseTool):
+    """Tool that takes in function or coroutine directly."""
+
+    description: str = ""
+
+    func: Callable[..., str] | None
+    """The function to run when the tool is called."""
+
+    coroutine: Callable[..., Awaitable[str]] | None = None
+    """The asynchronous version of the function."""
+
+    # --- Runnable ---
+
+    @override
+    async def ainvoke(
+        self,
+        input: str | dict[str, Any] | ToolCall,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        if not self.coroutine:
+            # If the tool does not implement async, fall back to default implementation
+            return await run_in_executor(config, self.invoke, input, config, **kwargs)
+
+        return await super().ainvoke(input, config, **kwargs)
+
+    # --- Tool ---
+
+    @property
+    def args(self) -> dict[str, Any]:
+        """The tool's input arguments.
+
+        Returns:
+            The input arguments for the tool.
+        """
+        if self.args_schema is not None:
+            return super().args
+        # For backwards compatibility, if the function signature is ambiguous,
+        # assume it takes a single string input.
+        return {"tool_input": {"type": "string"}}
+
+    def _to_args_and_kwargs(
+        self, tool_input: str | dict[str, Any], tool_call_id: str | None
+    ) -> tuple[tuple[str, ...], dict[str, Any]]:
+        """Convert tool input to Pydantic model.
+
+        Args:
+            tool_input: The input to the tool.
+            tool_call_id: The ID of the tool call.
+
+        Raises:
+            ToolException: If the tool input is invalid.
+
+        Returns:
+            The Pydantic model args and kwargs.
+        """
+        args, kwargs = super()._to_args_and_kwargs(tool_input, tool_call_id)
+        # For backwards compatibility. The tool must be run with a single input
+        all_args = list(args) + list(kwargs.values())
+        if len(all_args) != 1:
+            msg = (
+                f"""Too many arguments to single-input tool {self.name}.
+                Consider using StructuredTool instead."""
+                f" Args: {all_args}"
+            )
+            raise ToolException(msg)
+        return tuple(all_args), {}
+
+    def _run(
+        self,
+        *args: Any,
+        config: RunnableConfig,
+        run_manager: CallbackManagerForToolRun | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Use the tool.
+
+        Args:
+            *args: Positional arguments to pass to the tool
+            config: Configuration for the run
+            run_manager: Optional callback manager to use for the run
+            **kwargs: Keyword arguments to pass to the tool
+
+        Returns:
+            The result of the tool execution
+        """
+        if self.func:
+            if run_manager and signature(self.func).parameters.get("callbacks"):
+                kwargs["callbacks"] = run_manager.get_child()
+            if config_param := _get_runnable_config_param(self.func):
+                kwargs[config_param] = config
+            return self.func(*args, **kwargs)
+        msg = "Tool does not support sync invocation."
+        raise NotImplementedError(msg)
+
+    async def _arun(
+        self,
+        *args: Any,
+        config: RunnableConfig,
+        run_manager: AsyncCallbackManagerForToolRun | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Use the tool asynchronously.
+
+        Args:
+            *args: Positional arguments to pass to the tool
+            config: Configuration for the run
+            run_manager: Optional callback manager to use for the run
+            **kwargs: Keyword arguments to pass to the tool
+
+        Returns:
+            The result of the tool execution
+        """
+        if self.coroutine:
+            if run_manager and signature(self.coroutine).parameters.get("callbacks"):
+                kwargs["callbacks"] = run_manager.get_child()
+            if config_param := _get_runnable_config_param(self.coroutine):
+                kwargs[config_param] = config
+            return await self.coroutine(*args, **kwargs)
+
+        # NOTE: this code is unreachable since _arun is only called if coroutine is not
+        # None.
+        return await super()._arun(
+            *args, config=config, run_manager=run_manager, **kwargs
+        )
+
+    # TODO: this is for backwards compatibility, remove in future
+    def __init__(
+        self,
+        name: str,
+        func: Callable[..., Any] | None,
+        description: str,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize tool."""
+        super().__init__(name=name, func=func, description=description, **kwargs)
+
+    @classmethod
+    def from_function(
+        cls,
+        func: Callable[..., Any] | None,
+        name: str,  # We keep these required to support backwards compatibility
+        description: str,
+        return_direct: bool = False,  # noqa: FBT001,FBT002
+        args_schema: ArgsSchema | None = None,
+        coroutine: Callable[..., Awaitable[Any]]
+        | None = None,  # This is last for compatibility, but should be after func
+        **kwargs: Any,
+    ) -> Tool:
+        """Initialize tool from a function.
+
+        Args:
+            func: The function to create the tool from.
+            name: The name of the tool.
+            description: The description of the tool.
+            return_direct: Whether to return the output directly.
+            args_schema: The schema of the tool's input arguments.
+            coroutine: The asynchronous version of the function.
+            **kwargs: Additional arguments to pass to the tool.
+
+        Returns:
+            The tool.
+
+        Raises:
+            ValueError: If the function is not provided.
+        """
+        if func is None and coroutine is None:
+            msg = "Function and/or coroutine must be provided"
+            raise ValueError(msg)
+        return cls(
+            name=name,
+            func=func,
+            coroutine=coroutine,
+            description=description,
+            return_direct=return_direct,
+            args_schema=args_schema,
+            **kwargs,
+        )

coco_core/tools/structured.py

@@ -0,0 +1,322 @@
+"""Structured tool."""
+
+from __future__ import annotations
+
+import functools
+import textwrap
+from collections.abc import Awaitable, Callable
+from inspect import signature
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Literal,
+)
+
+from pydantic import Field, PlainSerializer, SkipValidation
+from typing_extensions import override
+
+# Cannot move to TYPE_CHECKING as _run/_arun parameter annotations are needed at runtime
+from coco_core.callbacks import (
+    AsyncCallbackManagerForToolRun,  # noqa: TC001
+    CallbackManagerForToolRun,  # noqa: TC001
+)
+from coco_core.runnables import RunnableConfig, run_in_executor
+from coco_core.tools.base import (
+    _EMPTY_SET,
+    FILTERED_ARGS,
+    ArgsSchema,
+    BaseTool,
+    _get_injected_args_keys_from_signature,
+    _get_runnable_config_param,
+    create_schema_from_function,
+)
+from coco_core.utils.pydantic import is_basemodel_subclass, model_json_schema
+
+if TYPE_CHECKING:
+    from coco_core.messages import ToolCall
+    from coco_core.utils.pydantic import TypeBaseModel
+
+
+def _serialize_as_str(value: Any) -> str:
+    """Stringify a value that has no JSON form.
+
+    A plain function rather than the `str` builtin: Pydantic < 2.9 inspects the
+    serializer's signature, and `inspect.signature` raises for C builtins.
+    """
+    return str(value)
+
+
+@functools.lru_cache(maxsize=256)
+def _model_json_schema(args_schema: TypeBaseModel, _rebuild_token: object) -> Any:
+    """Ask a v1 or v2 model class for its JSON schema, or fall back to its repr.
+
+    Cached because pydantic does not memoize this per class and it dominates the
+    dump -- 1.10 ms against 0.02 ms for an eight-tool payload -- which tracing
+    pays on every run.
+
+    `_rebuild_token` is the class's validator, which `model_rebuild()` replaces.
+    Keying on it retires the entry for a rebuilt schema, and for one whose
+    forward reference was still unresolved when it was first dumped. Mutating a
+    class in place without rebuilding it is not detected.
+    """
+    try:
+        return model_json_schema(args_schema)
+    except Exception:  # a schema holding an arbitrary type has no JSON schema
+        return str(args_schema)
+
+
+def _serialize_args_schema(args_schema: ArgsSchema) -> Any:
+    """Represent a schema class by its JSON schema when dumping to JSON.
+
+    A Pydantic model class has no JSON form, so leaving it to the default
+    serializer raises `PydanticSerializationError`. Its own JSON schema is the
+    shape a dict schema already has, and a dict is returned unchanged.
+    """
+    if isinstance(args_schema, dict):
+        return args_schema
+    return _model_json_schema(
+        args_schema, getattr(args_schema, "__pydantic_validator__", None)
+    )
+
+
+# Attached to the fields via `Annotated` rather than declared with
+# `@field_serializer`, which would take the field's one serializer slot and make
+# any subclass that declares its own serializer for the same field fail at class
+# creation with `PydanticUserError: Multiple field serializer functions ...`.
+_JsonSchemaFallback = PlainSerializer(
+    _serialize_args_schema, when_used="json-unless-none"
+)
+_JsonCallableFallback = PlainSerializer(_serialize_as_str, when_used="json-unless-none")
+
+
+class StructuredTool(BaseTool):
+    """Tool that can operate on any number of inputs."""
+
+    description: str = ""
+
+    args_schema: Annotated[ArgsSchema, SkipValidation(), _JsonSchemaFallback] = Field(
+        ..., description="The tool schema."
+    )
+    """The input arguments' schema."""
+
+    func: Annotated[Callable[..., Any] | None, _JsonCallableFallback] = None
+    """The function to run when the tool is called."""
+
+    coroutine: Annotated[
+        Callable[..., Awaitable[Any]] | None, _JsonCallableFallback
+    ] = None
+    """The asynchronous version of the function."""
+
+    # --- Runnable ---
+
+    # TODO: Is this needed?
+    @override
+    async def ainvoke(
+        self,
+        input: str | dict[str, Any] | ToolCall,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        if not self.coroutine:
+            # If the tool does not implement async, fall back to default implementation
+            return await run_in_executor(config, self.invoke, input, config, **kwargs)
+
+        return await super().ainvoke(input, config, **kwargs)
+
+    # --- Tool ---
+
+    def _run(
+        self,
+        *args: Any,
+        config: RunnableConfig,
+        run_manager: CallbackManagerForToolRun | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Use the tool.
+
+        Args:
+            *args: Positional arguments to pass to the tool
+            config: Configuration for the run
+            run_manager: Optional callback manager to use for the run
+            **kwargs: Keyword arguments to pass to the tool
+
+        Returns:
+            The result of the tool execution
+        """
+        if self.func:
+            if run_manager and signature(self.func).parameters.get("callbacks"):
+                kwargs["callbacks"] = run_manager.get_child()
+            if config_param := _get_runnable_config_param(self.func):
+                kwargs[config_param] = config
+            return self.func(*args, **kwargs)
+        msg = "StructuredTool does not support sync invocation."
+        raise NotImplementedError(msg)
+
+    async def _arun(
+        self,
+        *args: Any,
+        config: RunnableConfig,
+        run_manager: AsyncCallbackManagerForToolRun | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Use the tool asynchronously.
+
+        Args:
+            *args: Positional arguments to pass to the tool
+            config: Configuration for the run
+            run_manager: Optional callback manager to use for the run
+            **kwargs: Keyword arguments to pass to the tool
+
+        Returns:
+            The result of the tool execution
+        """
+        if self.coroutine:
+            if run_manager and signature(self.coroutine).parameters.get("callbacks"):
+                kwargs["callbacks"] = run_manager.get_child()
+            if config_param := _get_runnable_config_param(self.coroutine):
+                kwargs[config_param] = config
+            return await self.coroutine(*args, **kwargs)
+
+        # If self.coroutine is None, then this will delegate to the default
+        # implementation which is expected to delegate to _run on a separate thread.
+        return await super()._arun(
+            *args, config=config, run_manager=run_manager, **kwargs
+        )
+
+    @classmethod
+    def from_function(
+        cls,
+        func: Callable[..., Any] | None = None,
+        coroutine: Callable[..., Awaitable[Any]] | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        return_direct: bool = False,  # noqa: FBT001,FBT002
+        args_schema: ArgsSchema | None = None,
+        infer_schema: bool = True,  # noqa: FBT001,FBT002
+        *,
+        response_format: Literal["content", "content_and_artifact"] = "content",
+        parse_docstring: bool = False,
+        error_on_invalid_docstring: bool = False,
+        **kwargs: Any,
+    ) -> StructuredTool:
+        """Create tool from a given function.
+
+        A classmethod that helps to create a tool from a function.
+
+        Args:
+            func: The function from which to create a tool.
+            coroutine: The async function from which to create a tool.
+            name: The name of the tool.
+
+                Defaults to the function name.
+            description: The description of the tool.
+
+                Defaults to the function docstring.
+            return_direct: Whether to return the result directly or as a callback.
+            args_schema: The schema of the tool's input arguments.
+            infer_schema: Whether to infer the schema from the function's signature.
+            response_format: The tool response format.
+
+                If `'content'` then the output of the tool is interpreted as the
+                contents of a `ToolMessage`. If `'content_and_artifact'` then the output
+                is expected to be a two-tuple corresponding to the `(content, artifact)`
+                of a `ToolMessage`.
+            parse_docstring: If `infer_schema` and `parse_docstring`, will attempt
+                to parse parameter descriptions from Google Style function docstrings.
+            error_on_invalid_docstring: if `parse_docstring` is provided, configure
+                whether to raise `ValueError` on invalid Google Style docstrings.
+            **kwargs: Additional arguments to pass to the tool
+
+        Returns:
+            The tool.
+
+        Raises:
+            ValueError: If the function is not provided.
+            ValueError: If the function does not have a docstring and description
+                is not provided.
+            TypeError: If the `args_schema` is not a `BaseModel` or dict.
+
+        Examples:
+            ```python
+            def add(a: int, b: int) -> int:
+                \"\"\"Add two numbers\"\"\"
+                return a + b
+            tool = StructuredTool.from_function(add)
+            tool.run(1, 2) # 3
+
+            ```
+        """
+        if func is not None:
+            source_function = func
+        elif coroutine is not None:
+            source_function = coroutine
+        else:
+            msg = "Function and/or coroutine must be provided"
+            raise ValueError(msg)
+        name = name or source_function.__name__
+        if args_schema is None and infer_schema:
+            # schema name is appended within function
+            args_schema = create_schema_from_function(
+                name,
+                source_function,
+                parse_docstring=parse_docstring,
+                error_on_invalid_docstring=error_on_invalid_docstring,
+                filter_args=_filter_schema_args(source_function),
+            )
+        description_ = description
+        if description is None and not parse_docstring:
+            description_ = source_function.__doc__ or None
+        if description_ is None and args_schema:
+            if isinstance(args_schema, type) and is_basemodel_subclass(args_schema):
+                description_ = args_schema.__doc__
+                if (
+                    description_
+                    and "A base class for creating Pydantic models" in description_
+                ):
+                    description_ = ""
+                elif not description_:
+                    description_ = None
+            elif isinstance(args_schema, dict):
+                description_ = args_schema.get("description")
+            else:
+                msg = (
+                    "Invalid args_schema: expected BaseModel or dict, "
+                    f"got {args_schema}"
+                )
+                raise TypeError(msg)
+        if description_ is None:
+            msg = "Function must have a docstring if description not provided."
+            raise ValueError(msg)
+        if description is None:
+            # Only apply if using the function's docstring
+            description_ = textwrap.dedent(description_).strip()
+
+        # Description example:
+        # search_api(query: str) - Searches the API for the query.
+        description_ = f"{description_.strip()}"
+        return cls(
+            name=name,
+            func=func,
+            coroutine=coroutine,
+            args_schema=args_schema,
+            description=description_,
+            return_direct=return_direct,
+            response_format=response_format,
+            **kwargs,
+        )
+
+    @functools.cached_property
+    def _injected_args_keys(self) -> frozenset[str]:
+        fn = self.func or self.coroutine
+        if fn is None:
+            return _EMPTY_SET
+        return _get_injected_args_keys_from_signature(fn)
+
+
+def _filter_schema_args(func: Callable[..., Any]) -> list[str]:
+    filter_args = list(FILTERED_ARGS)
+    if config_param := _get_runnable_config_param(func):
+        filter_args.append(config_param)
+    # filter_args.extend(_get_non_model_params(type_hints))
+    return filter_args

coco_core/tracers/__init__.py

@@ -0,0 +1,47 @@
+"""Tracers are classes for tracing runs."""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.tracers.base import BaseTracer
+    from coco_core.tracers.coco import CocoTracer
+    from coco_core.tracers.log_stream import (
+        LogStreamCallbackHandler,
+        RunLog,
+        RunLogPatch,
+    )
+    from coco_core.tracers.schemas import Run
+    from coco_core.tracers.stdout import ConsoleCallbackHandler
+
+__all__ = (
+    "BaseTracer",
+    "ConsoleCallbackHandler",
+    "CocoTracer",
+    "LogStreamCallbackHandler",
+    "Run",
+    "RunLog",
+    "RunLogPatch",
+)
+
+_dynamic_imports = {
+    "BaseTracer": "base",
+    "CocoTracer": "coco",
+    "LogStreamCallbackHandler": "log_stream",
+    "RunLog": "log_stream",
+    "RunLogPatch": "log_stream",
+    "Run": "schemas",
+    "ConsoleCallbackHandler": "stdout",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)

coco_core/tracers/_compat.py

@@ -0,0 +1,95 @@
+"""Compatibility helpers for Pydantic v1/v2 with langsmith `Run` objects.
+
+!!! note
+
+    The generic helpers (`pydantic_to_dict`, `pydantic_copy`) detect Pydanti version
+    based on the langsmith `Run` model. They're intended for langsmith objects (`Run`,
+    `Example`) which migrate together.
+
+For general Pydantic v1/v2 handling, see `coco_core.utils.pydantic`.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+from coco_core.tracers.schemas import Run
+
+# Detect Pydantic version once at import time based on Run model
+_RUN_IS_PYDANTIC_V2 = hasattr(Run, "model_dump")
+
+T = TypeVar("T")
+
+
+def run_to_dict(run: Run, **kwargs: Any) -> dict[str, Any]:
+    """Convert run to dict, compatible with both Pydantic v1 and v2.
+
+    Args:
+        run: The run to convert.
+        **kwargs: Additional arguments passed to `model_dump`/`dict`.
+
+    Returns:
+        Dictionary representation of the run.
+    """
+    if _RUN_IS_PYDANTIC_V2:
+        return run.model_dump(**kwargs)
+    return run.dict(**kwargs)  # type: ignore[deprecated]
+
+
+def run_copy(run: Run, **kwargs: Any) -> Run:
+    """Copy run, compatible with both Pydantic v1 and v2.
+
+    Args:
+        run: The run to copy.
+        **kwargs: Additional arguments passed to `model_copy`/`copy`.
+
+    Returns:
+        A copy of the run.
+    """
+    if _RUN_IS_PYDANTIC_V2:
+        return run.model_copy(**kwargs)
+    return run.copy(**kwargs)  # type: ignore[deprecated]
+
+
+def run_construct(**kwargs: Any) -> Run:
+    """Construct run without validation, compatible with both Pydantic v1 and v2.
+
+    Args:
+        **kwargs: Fields to set on the run.
+
+    Returns:
+        A new `Run` instance constructed without validation.
+    """
+    if _RUN_IS_PYDANTIC_V2:
+        return Run.model_construct(**kwargs)
+    return Run.construct(**kwargs)  # type: ignore[deprecated]
+
+
+def pydantic_to_dict(obj: Any, **kwargs: Any) -> dict[str, Any]:
+    """Convert any Pydantic model to dict, compatible with both v1 and v2.
+
+    Args:
+        obj: The Pydantic model to convert.
+        **kwargs: Additional arguments passed to `model_dump`/`dict`.
+
+    Returns:
+        Dictionary representation of the model.
+    """
+    if _RUN_IS_PYDANTIC_V2:
+        return obj.model_dump(**kwargs)  # type: ignore[no-any-return]
+    return obj.dict(**kwargs)  # type: ignore[no-any-return]
+
+
+def pydantic_copy(obj: T, **kwargs: Any) -> T:
+    """Copy any Pydantic model, compatible with both v1 and v2.
+
+    Args:
+        obj: The Pydantic model to copy.
+        **kwargs: Additional arguments passed to `model_copy`/`copy`.
+
+    Returns:
+        A copy of the model.
+    """
+    if _RUN_IS_PYDANTIC_V2:
+        return obj.model_copy(**kwargs)  # type: ignore[attr-defined,no-any-return]
+    return obj.copy(**kwargs)  # type: ignore[attr-defined,no-any-return]