feat(llmo): Use default PH client for LangChain (#293)

* feat(llmo): Use default PH client for langchain

* chore: Run formatter

* feat: Test the CallbackHandler without any PH client

* chore: Run formatter

Author: Radu-Raicea

posthog/ai/langchain/callbacks.py

@@ -33,7 +33,7 @@
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
-from posthog import default_client
+from posthog import setup
 from posthog.ai.utils import get_model_params, with_privacy_mode
 from posthog.client import Client
 
@@ -81,7 +81,7 @@ class CallbackHandler(BaseCallbackHandler):
     The PostHog LLM observability callback handler for LangChain.
     """
 
-    _client: Client
+    _ph_client: Client
     """PostHog client instance."""
 
     _distinct_id: Optional[Union[str, int, UUID]]
@@ -127,10 +127,7 @@ def __init__(
             privacy_mode: Whether to redact the input and output of the trace.
             groups: Optional additional PostHog groups to use for the trace.
         """
-        posthog_client = client or default_client
-        if posthog_client is None:
-            raise ValueError("PostHog client is required")
-        self._client = posthog_client
+        self._ph_client = client or setup()
         self._distinct_id = distinct_id
         self._trace_id = trace_id
         self._properties = properties or {}
@@ -481,7 +478,7 @@ def _capture_trace_or_span(
         event_properties = {
             "$ai_trace_id": trace_id,
             "$ai_input_state": with_privacy_mode(
-                self._client, self._privacy_mode, run.input
+                self._ph_client, self._privacy_mode, run.input
             ),
             "$ai_latency": run.latency,
             "$ai_span_name": run.name,
@@ -497,13 +494,13 @@ def _capture_trace_or_span(
             event_properties["$ai_is_error"] = True
         elif outputs is not None:
             event_properties["$ai_output_state"] = with_privacy_mode(
-                self._client, self._privacy_mode, outputs
+                self._ph_client, self._privacy_mode, outputs
             )
 
         if self._distinct_id is None:
             event_properties["$process_person_profile"] = False
 
-        self._client.capture(
+        self._ph_client.capture(
             distinct_id=self._distinct_id or run_id,
             event=event_name,
             properties=event_properties,
@@ -550,14 +547,16 @@ def _capture_generation(
             "$ai_provider": run.provider,
             "$ai_model": run.model,
             "$ai_model_parameters": run.model_params,
-            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.input),
+            "$ai_input": with_privacy_mode(
+                self._ph_client, self._privacy_mode, run.input
+            ),
             "$ai_http_status": 200,
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
         }
         if run.tools:
             event_properties["$ai_tools"] = with_privacy_mode(
-                self._client,
+                self._ph_client,
                 self._privacy_mode,
                 run.tools,
             )
@@ -589,7 +588,7 @@ def _capture_generation(
                     _extract_raw_esponse(generation) for generation in generation_result
                 ]
             event_properties["$ai_output_choices"] = with_privacy_mode(
-                self._client, self._privacy_mode, completions
+                self._ph_client, self._privacy_mode, completions
             )
 
         if self._properties:
@@ -598,7 +597,7 @@ def _capture_generation(
         if self._distinct_id is None:
             event_properties["$process_person_profile"] = False
 
-        self._client.capture(
+        self._ph_client.capture(
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
             properties=event_properties,

posthog/test/ai/langchain/test_callbacks.py

@@ -1727,3 +1727,29 @@ def test_openai_reasoning_tokens(mock_client):
     assert call["properties"]["$ai_reasoning_tokens"] is not None
     assert call["properties"]["$ai_input_tokens"] is not None
     assert call["properties"]["$ai_output_tokens"] is not None
+
+
+def test_callback_handler_without_client():
+    """Test that CallbackHandler works properly when no PostHog client is passed."""
+    with patch("posthog.ai.langchain.callbacks.setup") as mock_setup:
+        mock_client = mock_setup.return_value
+
+        callbacks = CallbackHandler()
+
+        # Verify that setup() was called
+        mock_setup.assert_called_once()
+
+        # Verify that the client was set to the result of setup()
+        assert callbacks._ph_client == mock_client
+
+        # Test that the callback handler works with a simple chain
+        prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+        model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+        chain = prompt | model
+
+        # This should work and call the mock client
+        result = chain.invoke({}, config={"callbacks": [callbacks]})
+        assert result.content == "Bar"
+
+        # Verify that the mock client was used for capturing events
+        assert mock_client.capture.call_count == 3