diff --git a/.sampo/changesets/cranky-runesinger-tuonetar.md b/.sampo/changesets/cranky-runesinger-tuonetar.md new file mode 100644 index 0000000..84f8d48 --- /dev/null +++ b/.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 diff --git a/lib/posthog/llm_analytics.ex b/lib/posthog/llm_analytics.ex index 9361dbb..32e53a1 100644 --- a/lib/posthog/llm_analytics.ex +++ b/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) diff --git a/public_api.snapshot b/public_api.snapshot index b20b4c7..9a7067f 100644 --- a/public_api.snapshot +++ b/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 diff --git a/test/posthog/integrations/plug_test.exs b/test/posthog/integrations/plug_test.exs index 9d516a4..0db804b 100644 --- a/test/posthog/integrations/plug_test.exs +++ b/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) diff --git a/test/posthog/llm_analytics_test.exs b/test/posthog/llm_analytics_test.exs index c525b45..95f575a 100644 --- a/test/posthog/llm_analytics_test.exs +++ b/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