Skip to content

Commit 8c16f3f

Browse files
TwixesTwixes
committed
Simplify things
Avoiding wonky casts.
1 parent 3356813 commit 8c16f3f

File tree

8 files changed

+27
-90
lines changed

8 files changed

+27
-90
lines changed

‎posthog/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -676,7 +676,7 @@ def shutdown():
676676
_proxy("join")
677677

678678

679-
def setup():
679+
def setup() -> Client:
680680
global default_client
681681
if not default_client:
682682
if not api_key:
@@ -706,6 +706,8 @@ def setup():
706706
default_client.disabled = disabled
707707
default_client.debug = debug
708708

709+
return default_client
710+
709711

710712
def _proxy(method, *args, **kwargs):
711713
"""Create an analytics client if one doesn't exist and send to it."""

‎posthog/ai/anthropic/anthropic.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
with_privacy_mode,
1818
)
1919
from posthog.client import Client as PostHogClient
20+
from posthog import setup
2021

2122

2223
class Anthropic(anthropic.Anthropic):
@@ -33,13 +34,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
3334
**kwargs: Additional arguments passed to the Anthropic client
3435
"""
3536
super().__init__(**kwargs)
36-
if posthog_client is None:
37-
import posthog
38-
39-
posthog.setup()
40-
self._ph_client = cast(PostHogClient, posthog.default_client)
41-
else:
42-
self._ph_client = posthog_client
37+
self._ph_client = posthog_client or setup()
4338
self.messages = WrappedMessages(self)
4439

4540

‎posthog/ai/anthropic/anthropic_async.py‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,9 @@
88

99
import time
1010
import uuid
11-
from typing import Any, Dict, Optional, cast
11+
from typing import Any, Dict, Optional
1212

13+
from posthog import setup
1314
from posthog.ai.utils import (
1415
call_llm_and_track_usage_async,
1516
get_model_params,
@@ -33,13 +34,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
3334
**kwargs: Additional arguments passed to the Anthropic client
3435
"""
3536
super().__init__(**kwargs)
36-
if posthog_client is None:
37-
import posthog
38-
39-
posthog.setup()
40-
self._ph_client = cast(PostHogClient, posthog.default_client)
41-
else:
42-
self._ph_client = posthog_client
37+
self._ph_client = posthog_client or setup()
4338
self.messages = AsyncWrappedMessages(self)
4439

4540

‎posthog/ai/anthropic/anthropic_providers.py‎

Lines changed: 6 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@
55
"Please install the Anthropic SDK to use this feature: 'pip install anthropic'"
66
)
77

8-
from typing import Optional, cast
8+
from typing import Optional
99

1010
from posthog.ai.anthropic.anthropic import WrappedMessages
1111
from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages
1212
from posthog.client import Client as PostHogClient
13+
from posthog import setup
1314

1415

1516
class AnthropicBedrock(anthropic.AnthropicBedrock):
@@ -21,13 +22,7 @@ class AnthropicBedrock(anthropic.AnthropicBedrock):
2122

2223
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
2324
super().__init__(**kwargs)
24-
if posthog_client is None:
25-
import posthog
26-
27-
posthog.setup()
28-
self._ph_client = cast(PostHogClient, posthog.default_client)
29-
else:
30-
self._ph_client = posthog_client
25+
self._ph_client = posthog_client or setup()
3126
self.messages = WrappedMessages(self)
3227

3328

@@ -40,13 +35,7 @@ class AsyncAnthropicBedrock(anthropic.AsyncAnthropicBedrock):
4035

4136
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
4237
super().__init__(**kwargs)
43-
if posthog_client is None:
44-
import posthog
45-
46-
posthog.setup()
47-
self._ph_client = cast(PostHogClient, posthog.default_client)
48-
else:
49-
self._ph_client = posthog_client
38+
self._ph_client = posthog_client or setup()
5039
self.messages = AsyncWrappedMessages(self)
5140

5241

@@ -59,13 +48,7 @@ class AnthropicVertex(anthropic.AnthropicVertex):
5948

6049
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
6150
super().__init__(**kwargs)
62-
if posthog_client is None:
63-
import posthog
64-
65-
posthog.setup()
66-
self._ph_client = cast(PostHogClient, posthog.default_client)
67-
else:
68-
self._ph_client = posthog_client
51+
self._ph_client = posthog_client or setup()
6952
self.messages = WrappedMessages(self)
7053

7154

@@ -78,11 +61,5 @@ class AsyncAnthropicVertex(anthropic.AsyncAnthropicVertex):
7861

7962
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
8063
super().__init__(**kwargs)
81-
if posthog_client is None:
82-
import posthog
83-
84-
posthog.setup()
85-
self._ph_client = cast(PostHogClient, posthog.default_client)
86-
else:
87-
self._ph_client = posthog_client
64+
self._ph_client = posthog_client or setup()
8865
self.messages = AsyncWrappedMessages(self)

‎posthog/ai/gemini/gemini.py‎

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import os
22
import time
33
import uuid
4-
from typing import Any, Dict, Optional, cast
4+
from typing import Any, Dict, Optional
55

66
try:
77
from google import genai
@@ -10,6 +10,7 @@
1010
"Please install the Google Gemini SDK to use this feature: 'pip install google-genai'"
1111
)
1212

13+
from posthog import setup
1314
from posthog.ai.utils import (
1415
call_llm_and_track_usage,
1516
get_model_params,
@@ -58,13 +59,7 @@ def __init__(
5859
posthog_groups: Default groups for all calls (can be overridden per call)
5960
**kwargs: Additional arguments (for future compatibility)
6061
"""
61-
if posthog_client is None:
62-
import posthog
63-
64-
posthog.setup()
65-
self._ph_client = cast(PostHogClient, posthog.default_client)
66-
else:
67-
self._ph_client = posthog_client
62+
self._ph_client = posthog_client or setup()
6863

6964
if self._ph_client is None:
7065
raise ValueError("posthog_client is required for PostHog tracking")
@@ -107,13 +102,7 @@ def __init__(
107102
posthog_groups: Default groups for all calls
108103
**kwargs: Additional arguments (for future compatibility)
109104
"""
110-
if posthog_client is None:
111-
import posthog
112-
113-
posthog.setup()
114-
self._ph_client = cast(PostHogClient, posthog.default_client)
115-
else:
116-
self._ph_client = posthog_client
105+
self._ph_client = posthog_client or setup()
117106

118107
if self._ph_client is None:
119108
raise ValueError("posthog_client is required for PostHog tracking")

‎posthog/ai/openai/openai.py‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import time
22
import uuid
3-
from typing import Any, Dict, List, Optional, cast
3+
from typing import Any, Dict, List, Optional
44

55
try:
66
import openai
@@ -15,6 +15,7 @@
1515
with_privacy_mode,
1616
)
1717
from posthog.client import Client as PostHogClient
18+
from posthog import setup
1819

1920

2021
class OpenAI(openai.OpenAI):
@@ -32,13 +33,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
3233
**openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
3334
"""
3435
super().__init__(**kwargs)
35-
if posthog_client is None:
36-
import posthog
37-
38-
posthog.setup()
39-
self._ph_client = cast(PostHogClient, posthog.default_client)
40-
else:
41-
self._ph_client = posthog_client
36+
self._ph_client = posthog_client or setup()
4237

4338
# Store original objects after parent initialization (only if they exist)
4439
self._original_chat = getattr(self, "chat", None)

‎posthog/ai/openai/openai_async.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
"Please install the OpenAI SDK to use this feature: 'pip install openai'"
1010
)
1111

12+
from posthog import setup
1213
from posthog.ai.utils import (
1314
call_llm_and_track_usage_async,
1415
get_model_params,
@@ -33,13 +34,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
3334
**openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
3435
"""
3536
super().__init__(**kwargs)
36-
if posthog_client is None:
37-
import posthog
38-
39-
posthog.setup()
40-
self._ph_client = cast(PostHogClient, posthog.default_client)
41-
else:
42-
self._ph_client = posthog_client
37+
self._ph_client = posthog_client or setup()
4338

4439
# Store original objects after parent initialization (only if they exist)
4540
self._original_chat = getattr(self, "chat", None)

‎posthog/ai/openai/openai_providers.py‎

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,10 @@
1515
from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
1616
from posthog.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings
1717
from posthog.ai.openai.openai_async import WrappedResponses as AsyncWrappedResponses
18-
from typing import Optional, cast
18+
from typing import Optional
1919

2020
from posthog.client import Client as PostHogClient
21+
from posthog import setup
2122

2223

2324
class AzureOpenAI(openai.AzureOpenAI):
@@ -36,13 +37,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
3637
**openai_config: Any additional keyword args to set on Azure OpenAI (e.g. azure_endpoint="xxx").
3738
"""
3839
super().__init__(**kwargs)
39-
if posthog_client is None:
40-
import posthog
41-
42-
posthog.setup()
43-
self._ph_client = cast(PostHogClient, posthog.default_client)
44-
else:
45-
self._ph_client = posthog_client
40+
self._ph_client = posthog_client or setup()
4641

4742
# Store original objects after parent initialization (only if they exist)
4843
self._original_chat = getattr(self, "chat", None)
@@ -80,13 +75,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
8075
**openai_config: Any additional keyword args to set on Azure OpenAI (e.g. azure_endpoint="xxx").
8176
"""
8277
super().__init__(**kwargs)
83-
if posthog_client is None:
84-
import posthog
85-
86-
posthog.setup()
87-
self._ph_client = cast(PostHogClient, posthog.default_client)
88-
else:
89-
self._ph_client = posthog_client
78+
self._ph_client = posthog_client or setup()
9079

9180
# Store original objects after parent initialization (only if they exist)
9281
self._original_chat = getattr(self, "chat", None)

0 commit comments

Comments
 (0)