Parse '|' annotations in tool decorator

Author: Andrei Bratu

.github/workflows/ci.yml

@@ -10,7 +10,7 @@ jobs:
       - name: Set up python
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: 3.12
       - name: Bootstrap poetry
         run: |
           curl -sSL https://install.python-poetry.org | python - -y --version 1.5.1
@@ -41,6 +41,32 @@ jobs:
           REPLICATE_API_KEY: ${{ secrets.REPLICATE_API_KEY }}
           GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+  test_3_12:
+    # Test with Python 3.12
+    # Some tool decorator tests assert the ability to parse the signature
+    # if functions that use typing features introduced in Python 3.10 e.g. '|'
+    runs-on: ubuntu-20.04
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v3
+      - name: Set up python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.12
+      - name: Bootstrap poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python - -y --version 1.5.1
+      - name: Install dependencies
+        run: poetry install
+
+      - name: Test
+        run: poetry run pytest -rP .
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          REPLICATE_API_KEY: ${{ secrets.REPLICATE_API_KEY }}
+          GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
 
   publish:
     needs: [compile, test]

src/humanloop/decorators/tool.py

@@ -1,6 +1,7 @@
 import builtins
 import inspect
 import logging
+import sys
 import textwrap
 import typing
 from dataclasses import dataclass
@@ -24,6 +25,9 @@
 from humanloop.requests.tool_function import ToolFunctionParams
 from humanloop.requests.tool_kernel_request import ToolKernelRequestParams
 
+if sys.version_info >= (3, 10):
+    import types
+
 logger = logging.getLogger("humanloop.sdk")
 
 
@@ -335,7 +339,7 @@ def _parse_annotation(annotation: typing.Type) -> _ParsedAnnotation:
             annotation=[_parse_annotation(arg) for arg in typing.get_args(annotation)],
         )
 
-    if origin is typing.Union:
+    if origin is typing.Union or (sys.version_info >= (3, 10) and origin is types.UnionType):
         sub_types = typing.get_args(annotation)
         if sub_types[-1] is type(None):
             # type(None) in sub_types indicates Optional type
@@ -495,4 +499,8 @@ def _parameter_is_optional(
     origin = typing.get_origin(parameter.annotation)
     # sub_types refers to T inside the annotation
     sub_types = typing.get_args(parameter.annotation)
-    return origin is typing.Union and len(sub_types) > 0 and sub_types[-1] is type(None)
+    return (
+        (origin is typing.Union or (sys.version_info >= (3, 10) and origin is types.UnionType))
+        and len(sub_types) > 0
+        and sub_types[-1] is type(None)
+    )

tests/decorators/test_tool_decorator.py

@@ -1,3 +1,4 @@
+import sys
 from typing import Any, Optional, TypedDict, Union
 
 import pytest
@@ -460,3 +461,72 @@ def calculator(operation: str, num1: float, num2: float) -> float:
         key=HUMANLOOP_FILE_KEY,
     )
     assert hl_file_higher_order_fn["tool"]["source_code"] == hl_file_decorated_fn["tool"]["source_code"]  # type: ignore
+
+
+def test_python310_syntax(
+    opentelemetry_test_configuration: tuple[Tracer, InMemorySpanExporter],
+):
+    if sys.version_info < (3, 10):
+        pytest.skip("Requires Python 3.10")
+    # GIVEN an OTel configuration
+    tracer, _ = opentelemetry_test_configuration
+
+    # GIVEN a function annotated with @tool where a parameter uses `|` for Optional
+    @tool(opentelemetry_tracer=tracer)
+    def calculator(a: float, b: float | None = None) -> float:
+        if a is None:
+            a = 0
+        return a + b
+
+    # WHEN building the Tool kernel
+    # THEN the JSON schema is correct
+    assert calculator.json_schema == {
+        "description": "",
+        "name": "calculator",
+        "parameters": {
+            "properties": {
+                "a": {"type": "number"},
+                "b": {"type": ["number", "null"]},
+            },
+            "required": ("a",),
+            "type": "object",
+            "additionalProperties": False,
+        },
+        "strict": True,
+    }
+
+    Validator.check_schema(calculator.json_schema)
+
+
+def test_python310_union_syntax(
+    opentelemetry_test_configuration: tuple[Tracer, InMemorySpanExporter],
+):
+    if sys.version_info < (3, 10):
+        pytest.skip("Requires Python 3.10")
+
+    # GIVEN an OTel configuration
+    tracer, _ = opentelemetry_test_configuration
+
+    # GIVEN a function annotated with @tool where a parameter uses `|` for Union
+    @tool(opentelemetry_tracer=tracer)
+    def calculator(a: float, b: float | int | str) -> float:
+        return a + b
+
+    # WHEN building the Tool kernel
+    # THEN the JSON schema is correct
+    assert calculator.json_schema == {
+        "description": "",
+        "name": "calculator",
+        "parameters": {
+            "properties": {
+                "a": {"type": "number"},
+                "b": {"anyOf": [{"type": "number"}, {"type": "integer"}, {"type": "string"}]},
+            },
+            "required": ("a", "b"),
+            "type": "object",
+            "additionalProperties": False,
+        },
+        "strict": True,
+    }
+
+    Validator.check_schema(calculator.json_schema)