feat: Added support for stream helper function

Author: 7HR4IZ3

examples/anthrophic.py

@@ -69,6 +69,41 @@ async def test_async_streaming():
         print(event)
 
 
+def test_sync_stream_helper():
+    client = Anthropic()
+    monitor(client)
+
+    with client.messages.stream(
+        max_tokens=1024,
+        messages=[{
+            "role": "user",
+            "content": "Hello, Claude",
+        }],
+        model="claude-3-opus-20240229",
+    ) as stream:
+        for event in stream:
+            print(event)
+
+async def test_async_stream_helper():
+    client = monitor(AsyncAnthropic())
+
+    async with client.messages.stream(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Say hello there!",
+            }
+        ],
+        model="claude-3-opus-20240229",
+    ) as stream:
+        async for event in stream:
+            print(event)
+
+    message = await stream.get_final_message()
+    print(message.to_json())
+
+
 def test_extra_arguments():
     client = Anthropic()
     monitor(client)
@@ -100,4 +135,7 @@ def test_extra_arguments():
 # test_sync_streaming()
 # test_asyncio.run(async_streaming())
 
-test_extra_arguments()
+# test_extra_arguments()
+
+# test_sync_stream_helper()
+asyncio.run(test_async_stream_helper())

lunary/anthrophic.py

@@ -1,14 +1,42 @@
 import typing as t
 from functools import partial
+from inspect import iscoroutine
+from contextlib import AsyncContextDecorator, ContextDecorator
+
 from . import track_event, run_context, run_manager, logging, logger, user_props_ctx, user_ctx, traceback, tags_ctx, filter_params
 
 try:
     from anthropic import Anthropic, AsyncAnthropic
     from anthropic.types import Message
+    from anthropic.lib.streaming import MessageStreamManager, AsyncMessageStreamManager
 except ImportError:
     raise ImportError("Anthrophic SDK not installed!") from None
 
 
+class sync_context_wrapper(ContextDecorator):
+
+    def __init__(self, stream):
+        self.__stream = stream
+
+    def __enter__(self):
+        return self.__stream
+
+    def __exit__(self, *_):
+        return
+
+
+class async_context_wrapper(AsyncContextDecorator):
+
+    def __init__(self, stream):
+        self.__stream = stream
+
+    async def __aenter__(self):
+        return self.__stream
+
+    async def __aexit__(self, *_):
+        return
+
+
 def __input_parser(kwargs: t.Dict):
     return {"input": kwargs.get("messages"), "name": kwargs.get("model")}
 
@@ -35,42 +63,44 @@ def __output_parser(output: t.Union[Message], stream: bool = False):
 
 def __stream_handler(method, run_id, name, type, *args, **kwargs):
     messages = []
+    original_stream = None
     stream = method(*args, **kwargs)
 
+    if isinstance(stream, MessageStreamManager):
+        original_stream = stream
+        stream = original_stream.__enter__()
+
     for event in stream:
         if event.type == "message_start":
-            # print(event.message.model)
             messages.append({
                 "role": event.message.role,
                 "model": event.message.model
             })
         if event.type == "message_delta":
-            # print("*", event.usage.output_tokens)
             if len(messages) >= 1:
                 message = messages[-1]
                 message["usage"] = {"tokens": event.usage.output_tokens}
 
         if event.type == "message_stop": pass
         if event.type == "content_block_start":
-            # print("* START")
-            # print(event.content_block.text)
             if len(messages) >= 1:
                 message = messages[-1]
                 message["output"] = event.content_block.text
 
         if event.type == "content_block_delta":
-            # print(event.delta.text, end="")
             if len(messages) >= 1:
                 message = messages[-1]
                 message["output"] = message.get("output",
                                                 "") + event.delta.text
 
         if event.type == "content_block_stop":
-            # print("* END")
             pass
 
         yield event
 
+    if original_stream:
+        original_stream.__exit__(None, None, None)
+
     track_event(
         type,
         "end",
@@ -86,42 +116,47 @@ def __stream_handler(method, run_id, name, type, *args, **kwargs):
 
 async def __async_stream_handler(method, run_id, name, type, *args, **kwargs):
     messages = []
-    stream = await method(*args, **kwargs)
+    original_stream = None
+    stream = method(*args, **kwargs)
+
+    if iscoroutine(stream):
+        stream = await stream
+
+    if isinstance(stream, AsyncMessageStreamManager):
+        original_stream = stream
+        stream = await original_stream.__aenter__()
 
     async for event in stream:
         if event.type == "message_start":
-            # print(event.message.model)
             messages.append({
                 "role": event.message.role,
                 "model": event.message.model
             })
         if event.type == "message_delta":
-            # print("*", event.usage.output_tokens)
             if len(messages) >= 1:
                 message = messages[-1]
                 message["usage"] = {"tokens": event.usage.output_tokens}
 
         if event.type == "message_stop": pass
         if event.type == "content_block_start":
-            # print("* START")
-            # print(event.content_block.text)
             if len(messages) >= 1:
                 message = messages[-1]
                 message["output"] = event.content_block.text
 
         if event.type == "content_block_delta":
-            # print(event.delta.text, end="")
             if len(messages) >= 1:
                 message = messages[-1]
                 message["output"] = message.get("output",
                                                 "") + event.delta.text
 
         if event.type == "content_block_stop":
-            # print("* END")
             pass
 
         yield event
 
+    if original_stream:
+        await original_stream.__aexit__(None, None, None)
+
     track_event(
         type,
         "end",
@@ -136,9 +171,7 @@ async def __async_stream_handler(method, run_id, name, type, *args, **kwargs):
 
 
 def __metadata_parser(metadata):
-    return {
-        x: metadata[x] for x in metadata if x in ["user_id"]
-    }
+    return {x: metadata[x] for x in metadata if x in ["user_id"]}
 
 
 def __wrap_sync(method: t.Callable,
@@ -152,6 +185,7 @@ def __wrap_sync(method: t.Callable,
                 output_parser=__output_parser,
                 stream_handler=__stream_handler,
                 metadata_parser=__metadata_parser,
+                contextify_stream: t.Optional[t.Callable] = None,
                 *args,
                 **kwargs):
     output = None
@@ -189,10 +223,13 @@ def __wrap_sync(method: t.Callable,
             except Exception as e:
                 logging.exception(e)
 
-            if kwargs.get("stream") == True:
-                return stream_handler(method, run.id, name
-                                      or parsed_input["name"], type, *args,
-                                      **kwargs)
+            if contextify_stream or kwargs.get("stream") == True:
+                generator = stream_handler(method, run.id, name
+                                           or parsed_input["name"], type,
+                                           *args, **kwargs)
+                if contextify_stream:
+                    return contextify_stream(generator)
+                else: return generator
 
             try:
                 output = method(*args, **kwargs)
@@ -241,6 +278,7 @@ async def __wrap_async(method: t.Callable,
                        output_parser=__output_parser,
                        stream_handler=__async_stream_handler,
                        metadata_parser=__metadata_parser,
+                       contextify_stream: t.Optional[bool] = False,
                        *args,
                        **kwargs):
     output = None
@@ -274,14 +312,17 @@ async def __wrap_async(method: t.Callable,
                                   or tags_ctx.get()),
                             template_id=(kwargs.get("extra_headers", {}).get(
                                 "Template-Id", None)),
-                            is_openai=True)
+                            is_openai=False)
             except Exception as e:
                 logging.exception(e)
 
-            if kwargs.get("stream") == True:
-                return stream_handler(method, run.id, name
-                                            or parsed_input["name"], type,
-                                            *args, **kwargs)
+            if contextify_stream or kwargs.get("stream") == True:
+                generator = stream_handler(method, run.id, name
+                                           or parsed_input["name"], type,
+                                           *args, **kwargs)
+                if contextify_stream:
+                    return contextify_stream(generator)
+                else: return generator
 
             try:
                 output = await method(*args, **kwargs)
@@ -325,11 +366,22 @@ async def __wrap_async(method: t.Callable,
 
 def monitor(client: "ClientType") -> "ClientType":
     if isinstance(client, Anthropic):
-        client.messages.create = partial(__wrap_sync, client.messages.create,
-                                         "llm")
+        client.messages.create = partial(__wrap_sync,
+                                         client.messages.create,
+                                         type="llm")
+        client.messages.stream = partial(__wrap_sync,
+                                         client.messages.stream,
+                                         type="llm",
+                                         contextify_stream=sync_context_wrapper)
     elif isinstance(client, AsyncAnthropic):
-        client.messages.create = partial(__wrap_async, client.messages.create,
-                                         "llm")
+        client.messages.create = partial(__wrap_async,
+                                         client.messages.create,
+                                         type="llm")
+        client.messages.stream = partial(__wrap_sync,
+                                         client.messages.stream,
+                                         type="llm",
+                                         stream_handler=__async_stream_handler,
+                                         contextify_stream=async_context_wrapper)
     else:
         raise Exception(
             "Invalid argument. Expected instance of Anthropic Client")