Fix logfire fstring weirdness. Add some additional fields to Message.extra from litellm.

Author: monoxgas

rigging/generator/litellm_.py

@@ -216,14 +216,24 @@ def _parse_model_response(
         ):
             tool_calls = [call.model_dump() for call in choice.message.tool_calls]
 
-        extra = {"response_id": response.id}
+        extra: dict[str, t.Any] = {"response_id": response.id}
         if hasattr(response, "provider"):
             extra["provider"] = response.provider
         if (
             hasattr(choice.message, "provider_specific_fields")
             and choice.message.provider_specific_fields is not None
         ):
             extra.update(choice.message.provider_specific_fields)
+        if (
+            hasattr(choice.message, "reasoning_content")
+            and choice.message.reasoning_content is not None
+        ):
+            extra["reasoning_content"] = choice.message.reasoning_content
+        if (
+            hasattr(choice.message, "thinking_blocks")
+            and choice.message.thinking_blocks is not None
+        ):
+            extra["thinking_blocks"] = choice.message.thinking_blocks
 
         message = Message(
             role="assistant",

rigging/tracing.py

@@ -1,5 +1,32 @@
+import typing as t
+
 import logfire_api
 
 Span = logfire_api.LogfireSpan
 
-tracer = logfire_api.Logfire(otel_scope="rigging")
+
+class Tracer(logfire_api.Logfire):
+    def span(
+        self,
+        msg_template: str,
+        /,
+        *,
+        _tags: t.Sequence[str] | None = None,
+        _span_name: str | None = None,
+        _level: t.Any | None = None,
+        _links: t.Any = (),
+        **attributes: t.Any,
+    ) -> logfire_api.LogfireSpan:
+        # Pass msg_template as the span name
+        # to avoid weird fstring behaviors
+        return super().span(
+            msg_template,
+            _tags=_tags,
+            _span_name=msg_template,
+            _level=_level,
+            _links=_links,
+            **attributes,
+        )
+
+
+tracer = Tracer(otel_scope="rigging")