From 9dd89b7ef12e6e6abef3a304952271f2aaa0178f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bar=C4=B1=C5=9F=20=C3=96zmen?= <hbaristr@gmail.com>
Date: Sun, 3 Nov 2024 01:07:58 +0300
Subject: [PATCH 1/2] logging now takes into account streaming functions +
 logfire API now uses span and structured outputs

---
 simplemind/logging.py | 51 +++++++++++++++++++++++++++++++++----------
 1 file changed, 39 insertions(+), 12 deletions(-)

diff --git a/simplemind/logging.py b/simplemind/logging.py
index 06defc3..22600ad 100644
--- a/simplemind/logging.py
+++ b/simplemind/logging.py
@@ -15,19 +15,46 @@ def logger(func: Callable[..., Any]) -> Callable[..., Any]:
         if not settings.logging.is_enabled:
             return func(*args, **kwargs)
 
-        logfire.info(f"Calling {func.__name__} with args: {args}, kwargs: {kwargs}")
-        t1 = time.perf_counter()
+        # See logfire manual tracing docs: https://logfire.pydantic.dev/docs/guides/onboarding-checklist/add-manual-tracing/#exceptions
+        with logfire.span("{event}", event="function called", function=func.__name__, args=args, kwargs=kwargs):
+            
+            t1 = time.perf_counter()
 
-        try:
-            result = func(*args, **kwargs)
-            t2 = time.perf_counter()
-            logfire.info(f"{func.__name__} returned: {result} in {t2-t1} seconds")
+            try:
+                is_streaming = "generate_stream_text" in func.__name__ or kwargs.get("stream")
+                if is_streaming:
+                    chunks = []
+                    for chunk in func(*args, **kwargs):
+                        chunks.append(chunk)
+                        yield chunk
+                    result = "".join(chunks)
+                    # note: no need to log the function name here, as it's already in the span
+                    logfire.info(
+                        "{event}",
+                        event="function completed",
+                        result=result[:2000],
+                        chunk_count=len(chunks),
+                        duration=time.perf_counter() - t1
+                    )
+                else:
+                    result = func(*args, **kwargs)
+                    logfire.info(
+                        "{event}",
+                        event="function completed",
+                        result=str(result)[:2000],
+                        duration=time.perf_counter() - t1
+                    )
+                    return result
 
-            return result
-
-        except Exception as e:
-            t2 = time.perf_counter()
-            logfire.error(f"Error in {func.__name__}: {e} in {t2-t1} seconds")
-            raise e
+            except Exception as e:
+                logfire.error(
+                    "{event}",
+                    event="function failed",
+                    error=str(e),
+                    duration=time.perf_counter() - t1,
+                    streamed_chunks=len(chunks) if is_streaming else None,
+                    partial_streaming_result="".join(chunks)[:2000] if is_streaming else None,
+                )
+                raise e
 
     return wrapper

From 9968f162d6ba64d6115d9c9723dcecfdf65e789f Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Sun, 3 Nov 2024 07:07:04 -0500
Subject: [PATCH 2/2] Update simplemind/logging.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
---
 simplemind/logging.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/simplemind/logging.py b/simplemind/logging.py
index 22600ad..bd2c1cc 100644
--- a/simplemind/logging.py
+++ b/simplemind/logging.py
@@ -46,14 +46,19 @@ def logger(func: Callable[..., Any]) -> Callable[..., Any]:
                     )
                     return result
 
-            except Exception as e:
+            if not settings.logging.is_enabled:
+                return func(*args, **kwargs)
+
+            chunks = []
+            try:
+                is_streaming = "generate_stream_text" in func.__name__ or kwargs.get("stream")
                 logfire.error(
                     "{event}",
                     event="function failed",
                     error=str(e),
                     duration=time.perf_counter() - t1,
-                    streamed_chunks=len(chunks) if is_streaming else None,
-                    partial_streaming_result="".join(chunks)[:2000] if is_streaming else None,
+                    streamed_chunks=len(chunks) if chunks else None,
+                    partial_streaming_result="".join(chunks)[:2000] if chunks else None,
                 )
                 raise e