refactor 3

tests/trace/test_op_coroutines.py

@@ -0,0 +1,89 @@
+import asyncio
+from typing import  Coroutine
+
+import pytest
+import weave
+from weave.trace.weave_client import Call
+
+
+def test_non_async_non_coro(client):
+    @weave.op()
+    def non_async_non_coro():
+        return 1
+
+    res = non_async_non_coro()
+    assert res == 1
+    res, call = non_async_non_coro.call()
+    assert isinstance(call, Call)
+    assert call.ended_at is not None
+    assert res == 1
+    assert call.ended_at is not None
+
+@pytest.mark.asyncio
+async def test_non_async_coro(client):
+    @weave.op()
+    def non_async_coro():
+        return asyncio.to_thread(lambda: 1)
+
+    res = non_async_coro()
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    res, call = non_async_coro.call()
+    assert isinstance(call, Call)
+    assert call.ended_at is not None
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    assert call.ended_at is not None
+
+@pytest.mark.asyncio
+async def test_async_coro(client):
+    @weave.op()
+    async def async_coro():
+        return asyncio.to_thread(lambda: 1)
+
+    res = async_coro()
+    assert isinstance(res, Coroutine)
+    res2 = await res
+    assert isinstance(res2, Coroutine)
+    assert await res2 == 1
+    res, call = async_coro.call()
+    assert isinstance(call, Call)
+    assert call.ended_at is None # BIG DIFFERENCE! We haven't ended the outer op yet!
+    assert isinstance(res, Coroutine)
+    res2 = await res
+    assert isinstance(res2, Coroutine) # Since we did not await the inner coroutine, we still have a coroutine here.
+    assert await res2 == 1 
+    assert call.ended_at is not None
+
+@pytest.mark.asyncio
+async def test_async_awaited_coro(client):
+    @weave.op()
+    async def async_awaited_coro():
+        return await asyncio.to_thread(lambda: 1)
+
+    res = async_awaited_coro()
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    res, call = async_awaited_coro.call()
+    assert isinstance(call, Call)
+    assert call.ended_at is None # BIG DIFFERENCE! We haven't ended the outer op yet!
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    assert call.ended_at is not None
+
+
+@pytest.mark.asyncio
+async def test_async_non_coro(client):
+    @weave.op()
+    async def async_non_coro():
+        return 1
+
+    res = async_non_coro()
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    res, call = async_non_coro.call()
+    assert isinstance(call, Call)
+    assert call.ended_at is None # BIG DIFFERENCE! We haven't ended the outer op yet!
+    assert isinstance(res, Coroutine)
+    assert await res == 1
+    assert call.ended_at is not None

weave/trace/op.py

@@ -341,7 +341,7 @@ def add(a: int, b: int) -> int:
 
     call = Call(_op_name="", trace_id="", parent_id=None, project_id="", inputs={})
     func = op.resolve_fn
-
+    is_async = inspect.iscoroutinefunction(func)
     if settings.should_disable_weave():
         res = func(*args, **kwargs)
     elif weave_client_context.get_weave_client() is None:
@@ -350,19 +350,24 @@ def add(a: int, b: int) -> int:
         res = func(*args, **kwargs)
     else:
         try:
-            # This try/except allows us to fail gracefully and
-            # still let the user code continue to execute
+            # Weird consideration: if the user does `op.call` and does not
+            # await the result, then the `started_at` timestamp will be
+            # set before the op itself executes!
             call = _create_call(op, *args, __weave=__weave, **kwargs)
             execute_res = _execute_call(
                 op, call, *args, __should_raise=__should_raise, **kwargs
             )
             return execute_res
         except Exception as e:
+            # This try/except allows us to fail gracefully and
+            # still let the user code continue to execute
             if get_raise_on_captured_errors():
                 raise
             log_once(
                 logger.error,
-                ASYNC_CALL_CREATE_MSG.format(traceback.format_exc()),
+                (ASYNC_CALL_CREATE_MSG if is_async else CALL_CREATE_MSG).format(
+                    traceback.format_exc()
+                ),
             )
             res = func(*args, **kwargs)
 

weave/trace/weave_client.py

@@ -731,6 +731,8 @@ def finish_call(
         *,
         op: Optional[Op] = None,
     ) -> None:
+        ended_at = datetime.datetime.now(tz=datetime.timezone.utc)
+        call.ended_at = ended_at
         original_output = output
 
         if op is not None and op.postprocess_output:
@@ -787,7 +789,6 @@ def finish_call(
             call.exception = exception_str
 
         project_id = self._project_id()
-        ended_at = datetime.datetime.now(tz=datetime.timezone.utc)
 
         # The finish handler serves as a last chance for integrations
         # to customize what gets logged for a call.