feat: add groups and privacy mode support to langchain callback

Author: k11kirky

posthog/ai/langchain/callbacks.py

@@ -23,7 +23,7 @@
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
-from posthog.ai.utils import get_model_params
+from posthog.ai.utils import get_model_params, with_privacy_mode
 from posthog.client import Client
 
 log = logging.getLogger("posthog")
@@ -69,18 +69,24 @@ def __init__(
         distinct_id: Optional[Union[str, int, float, UUID]] = None,
         trace_id: Optional[Union[str, int, float, UUID]] = None,
         properties: Optional[Dict[str, Any]] = None,
+        privacy_mode: bool = False,
+        groups: Optional[Dict[str, Any]] = None,
     ):
         """
         Args:
             client: PostHog client instance.
             distinct_id: Optional distinct ID of the user to associate the trace with.
             trace_id: Optional trace ID to use for the event.
             properties: Optional additional metadata to use for the trace.
+            privacy_mode: Whether to redact the input and output of the trace.
+            groups: Optional additional PostHog groups to use for the trace.
         """
         self._client = client
         self._distinct_id = distinct_id
         self._trace_id = trace_id
         self._properties = properties or {}
+        self._privacy_mode = privacy_mode
+        self._groups = groups or {}
         self._runs = {}
         self._parent_tree = {}
 
@@ -164,8 +170,8 @@ def on_llm_end(
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
-            "$ai_output": {"choices": output},
+            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
+            "$ai_output": with_privacy_mode(self._client, self._privacy_mode, {"choices": output}),
             "$ai_http_status": 200,
             "$ai_input_tokens": input_tokens,
             "$ai_output_tokens": output_tokens,
@@ -180,6 +186,7 @@ def on_llm_end(
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
             properties=event_properties,
+            groups=self._groups,
         )
 
     def on_chain_error(
@@ -212,7 +219,7 @@ def on_llm_error(
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
+            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
             "$ai_http_status": _get_http_status(error),
             "$ai_latency": latency,
             "$ai_trace_id": trace_id,
@@ -225,6 +232,7 @@ def on_llm_error(
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
             properties=event_properties,
+            groups=self._groups,
         )
 
     def _set_parent_of_run(self, run_id: UUID, parent_run_id: Optional[UUID] = None):

posthog/test/ai/langchain/test_callbacks.py

@@ -595,3 +595,54 @@ def test_base_url_retrieval(mock_client):
     assert mock_client.capture.call_count == 1
     call = mock_client.capture.call_args[1]
     assert call["properties"]["$ai_base_url"] == "https://test.posthog.com"
+
+
+def test_groups(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client, groups={"company": "test_company"})
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 1
+    call = mock_client.capture.call_args[1]
+    assert call["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client, privacy_mode=True)
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 1
+    call = mock_client.capture.call_args[1]
+    assert call["properties"]["$ai_input"] is None
+    assert call["properties"]["$ai_output"] is None
+
+
+def test_privacy_mode_global(mock_client):
+    mock_client.privacy_mode = True
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client)
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 1
+    call = mock_client.capture.call_args[1]
+    assert call["properties"]["$ai_input"] is None
+    assert call["properties"]["$ai_output"] is None