feat: add LLMAnalytics.pop_span/1 function

.sampo/changesets/cranky-runesinger-tuonetar.md

@@ -0,0 +1,5 @@
+---
+hex/posthog: minor
+---
+
+Add PostHog.LLMAnalytics.pop_span/1 function for convenient span context handoff between processes

lib/posthog/llm_analytics.ex

@@ -127,7 +127,8 @@ defmodule PostHog.LLMAnalytics do
 
   Just as with Context, LLMAnalytics tracks the current trace and span in the
   process dictionary. Any time you spawn a new process, you'll need to propagate
-  this information. Use `set_session/2`, `set_trace/2` and `set_root_span/2`:
+  this information. Use `set_session/2`, `set_trace/2` and `set_root_span/2` or
+  `pop_span/1`:
 
   ```
   def generate_response(user_message) do
@@ -401,7 +402,28 @@ defmodule PostHog.LLMAnalytics do
     capture_current_span(name, type, properties)
   end
 
-  defp pop_span(name) do
+  @doc """
+  Pop current span from process dictionary.
+
+  This is useful for when you started a span in the current process, but want to
+  hand it off to another process and capture there.
+
+  ## Examples
+
+      iex> PostHog.LLMAnalytics.start_span(%{"$ai_span_name": "LLM Call"})
+      "019ecdb3-ba1d-7504-978a-60242768140f"
+      iex> current_span = PostHog.LLMAnalytics.pop_span()
+      %{
+        "$ai_span_id": "019ecdb3-ba1d-7504-978a-60242768140f",
+        "$ai_span_name": "LLM Call"
+      }
+      iex> Task.async(fn ->
+        PostHog.LLMAnalytics.start_span(current_span)
+        PostHog.LLMAnalytics.capture_current_span("$ai_span")
+      end)
+  """
+  @spec pop_span(PostHog.supervisor_name()) :: PostHog.properties()
+  def pop_span(name \\ PostHog) do
     case Process.get({name, @span_backlog_key}) do
       [span | rest] ->
         Process.put({name, @span_backlog_key}, rest)

public_api.snapshot

@@ -123,6 +123,8 @@ PostHog.LLMAnalytics
     get_session/1 doc: documented
     get_trace/0 doc: none
     get_trace/1 doc: documented
+    pop_span/0 doc: none
+    pop_span/1 doc: documented
     set_root_span/1 doc: none
     set_root_span/2 doc: documented
     set_session/0 doc: none

test/posthog/integrations/plug_test.exs

@@ -24,7 +24,6 @@ defmodule PostHog.Integrations.PlugTest do
 
   defmodule MyRouter do
     use Plug.Router
-    require Logger
 
     plug(PostHog.Integrations.Plug)
     plug(:match)

test/posthog/llm_analytics_test.exs

@@ -512,4 +512,35 @@ defmodule Posthog.LLMAnalyticsTest do
              ] = all_captured(MyPostHog)
     end
   end
+
+  describe "pop_span/1" do
+    test "pops current span from the backlog" do
+      span_id = LLMAnalytics.start_span(%{foo: "bar"})
+
+      assert %{foo: "bar", "$ai_span_id": span_id} == LLMAnalytics.pop_span()
+      assert [] == Process.get({PostHog, :__llm_analytics_spans})
+    end
+
+    test "creates a new span if none is in the backlog" do
+      assert %{"$ai_span_id": "" <> _span_id} = LLMAnalytics.pop_span()
+      assert nil == Process.get({PostHog, :__llm_analytics_spans})
+    end
+
+    test "respects root span when popping a new one" do
+      LLMAnalytics.set_root_span("root_span_id")
+
+      assert %{"$ai_span_id": "" <> _span_id, "$ai_parent_id": "root_span_id"} =
+               LLMAnalytics.pop_span()
+
+      assert nil == Process.get({PostHog, :__llm_analytics_spans})
+    end
+
+    @tag config: [supervisor_name: MyPostHog]
+    test "custom PostHog instance" do
+      span_id = LLMAnalytics.start_span(MyPostHog, %{foo: "bar"})
+
+      assert %{foo: "bar", "$ai_span_id": span_id} == LLMAnalytics.pop_span(MyPostHog)
+      assert [] == Process.get({MyPostHog, :__llm_analytics_spans})
+    end
+  end
 end