feat: parking custom websocket client #634
Conversation
WalkthroughThis pull request introduces a pluggable WebSocket client abstraction layer for the Deepgram SDK. A new module defines protocol interfaces (WebSocketProtocol, SyncWebSocketProtocol, WebSocketFactory) and a default factory implementation using the websockets library. This abstraction is integrated into client wrappers and the main DeepgramClient constructors, allowing custom WebSocket clients to be injected. Documentation for SageMaker BiDi API integration is included. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant DeepgramClient
participant ClientWrapper
participant WebSocketFactory
participant WebSocket as Custom WebSocket<br/>(User-provided)
User->>DeepgramClient: __init__(..., websocket_client=factory)
DeepgramClient->>DeepgramClient: Extract websocket_client from kwargs
alt websocket_client provided
DeepgramClient->>ClientWrapper: Create wrapper with websocket_client
ClientWrapper->>WebSocketFactory: Store factory reference
else websocket_client not provided
DeepgramClient->>ClientWrapper: Create wrapper (uses default)
ClientWrapper->>WebSocketFactory: Store get_default_factory()
end
Note over DeepgramClient,ClientWrapper: Preserve existing state<br/>(api_key, headers, etc.)
DeepgramClient-->>User: Client instance ready
User->>DeepgramClient: Use client (WebSocket needed)
DeepgramClient->>ClientWrapper: Access websocket_client
ClientWrapper->>WebSocketFactory: connect_async/connect_sync()
WebSocketFactory-->>WebSocket: Create connection
WebSocket-->>ClientWrapper: WebSocket connection
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
sagemaker_bidi_api.md (1)
47-47: Standardize hyphenation of "unfragmented".Line 47 uses "un-fragmented" (with hyphen) while line 99 uses "unfragmented" (without hyphen). Choose one form consistently throughout the document.
Recommend using "unfragmented" (without hyphen) as the standard form:
- An un-fragmented binary chunk always has **`CompletionStatus = COMPLETE`**. + An unfragmented binary chunk always has **`CompletionStatus = COMPLETE`**.Also applies to: 99-99
src/deepgram/core/client_wrapper.py (1)
8-9: WebSocket factory plumbing through client wrappers is consistentImporting
WebSocketFactory/get_default_factoryand threadingwebsocket_clientthroughBaseClientWrapper,SyncClientWrapper, andAsyncClientWrapperlooks correct, and defaulting withwebsocket_client or get_default_factory()preserves existing behavior for callers that don’t specify a factory. The version header bump aligns with a new release.You might optionally tighten the internal attribute type to
WebSocketFactoryinstead ofOptional[WebSocketFactory], since it’s always populated after__init__, which can simplify downstream typing.Also applies to: 19-26, 32-33, 58-66, 81-89
src/deepgram/client.py (1)
191-288: Customwebsocket_clientsupport in sync/async clients is wired correctlyAllowing a
websocket_client: Optional[WebSocketFactory]kwarg, then reconstructing_client_wrapperasSyncClientWrapper/AsyncClientWrapperwith preservedapi_key, headers, environment, timeout, and the underlyinghttpx_clientkeeps backward compatibility while exposing the new extension point. The access‑token override and telemetry setup both operate on the final wrapper, so HTTP and WebSocket transports see consistent headers and instrumentation.One small robustness improvement: if the underlying wrapper implementation ever changes such that
original_wrapper.httpx_clientis not the thinHttpClient/AsyncHttpClientwrapper,original_wrapper.httpx_client.httpx_clientcould break. Using a safe fallback like:- self._client_wrapper = SyncClientWrapper( + httpx_client = getattr(original_wrapper.httpx_client, "httpx_client", original_wrapper.httpx_client) + self._client_wrapper = SyncClientWrapper( @@ - httpx_client=original_wrapper.httpx_client.httpx_client, + httpx_client=httpx_client,(and similarly for the async variant) would future‑proof that path with minimal cost.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
.fernignore(1 hunks)sagemaker_bidi_api.md(1 hunks)src/deepgram/client.py(13 hunks)src/deepgram/core/client_wrapper.py(4 hunks)src/deepgram/core/websocket_client.py(1 hunks)
🧰 Additional context used
🧠 Learnings (14)
📓 Common learnings
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/v1/websocket/__init__.py:8-8
Timestamp: 2024-10-09T02:19:46.086Z
Learning: Unused imports in `deepgram/clients/listen/v1/websocket/__init__.py` are retained to maintain backward compatibility and should not be flagged for removal in reviews.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/v1/websocket/__init__.py:8-8
Timestamp: 2024-07-01T19:21:39.778Z
Learning: Unused imports in `deepgram/clients/listen/v1/websocket/__init__.py` are retained to maintain backward compatibility and should not be flagged for removal in reviews.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/client.py:30-30
Timestamp: 2024-10-09T02:19:46.087Z
Learning: The `LiveOptions` import in `deepgram/clients/listen/v1/websocket/client.py` is intentionally present for future use and should not be flagged for removal in reviews.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/client.py:30-30
Timestamp: 2024-07-11T14:10:17.231Z
Learning: The `LiveOptions` import in `deepgram/clients/listen/v1/websocket/client.py` is intentionally present for future use and should not be flagged for removal in reviews.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/async_client.py:30-30
Timestamp: 2024-10-09T02:19:46.087Z
Learning: The unused import `LiveOptions` in `deepgram/clients/listen/v1/websocket/async_client.py` is intentional and will be used in PR #432.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/async_client.py:30-30
Timestamp: 2024-07-11T14:10:24.647Z
Learning: The unused import `LiveOptions` in `deepgram/clients/listen/v1/websocket/async_client.py` is intentional and will be used in PR #432.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 424
File: deepgram/client.py:81-81
Timestamp: 2024-06-27T00:06:01.811Z
Learning: Imports for SpeakStreamClient and AsyncSpeakStreamClient in `deepgram/client.py` are necessary for export purposes and should not be flagged as unused in reviews.
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 424
File: deepgram/client.py:81-81
Timestamp: 2024-10-09T02:19:46.087Z
Learning: Imports for SpeakStreamClient and AsyncSpeakStreamClient in `deepgram/client.py` are necessary for export purposes and should not be flagged as unused in reviews.
📚 Learning: 2024-10-09T02:19:46.086Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/v1/websocket/__init__.py:8-8
Timestamp: 2024-10-09T02:19:46.086Z
Learning: Unused imports in `deepgram/clients/listen/v1/websocket/__init__.py` are retained to maintain backward compatibility and should not be flagged for removal in reviews.
Applied to files:
src/deepgram/core/client_wrapper.py.fernignoresrc/deepgram/core/websocket_client.pysrc/deepgram/client.py
📚 Learning: 2024-07-11T14:10:17.231Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/client.py:30-30
Timestamp: 2024-07-11T14:10:17.231Z
Learning: The `LiveOptions` import in `deepgram/clients/listen/v1/websocket/client.py` is intentionally present for future use and should not be flagged for removal in reviews.
Applied to files:
src/deepgram/core/client_wrapper.pysrc/deepgram/core/websocket_client.pysrc/deepgram/client.py
📚 Learning: 2024-10-09T02:19:46.087Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 431
File: deepgram/clients/listen/v1/websocket/async_client.py:30-30
Timestamp: 2024-10-09T02:19:46.087Z
Learning: The unused import `LiveOptions` in `deepgram/clients/listen/v1/websocket/async_client.py` is intentional and will be used in PR #432.
Applied to files:
src/deepgram/core/client_wrapper.pysrc/deepgram/core/websocket_client.pysrc/deepgram/client.py
📚 Learning: 2024-06-27T00:06:01.811Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 424
File: deepgram/client.py:81-81
Timestamp: 2024-06-27T00:06:01.811Z
Learning: Imports for SpeakStreamClient and AsyncSpeakStreamClient in `deepgram/client.py` are necessary for export purposes and should not be flagged as unused in reviews.
Applied to files:
src/deepgram/core/client_wrapper.pysrc/deepgram/core/websocket_client.pysrc/deepgram/client.py
📚 Learning: 2024-07-11T14:25:40.772Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 432
File: deepgram/client.py:65-66
Timestamp: 2024-07-11T14:25:40.772Z
Learning: The `ListenRESTClient` and `AsyncListenRESTClient` imports in `deepgram/client.py` are intentional to support direct invocation, as demonstrated in the `examples/advanced/rest/direct_invocation` example.
Applied to files:
src/deepgram/core/client_wrapper.py.fernignoresrc/deepgram/client.py
📚 Learning: 2024-10-18T16:19:48.637Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 472
File: examples/speech-to-text/websocket/replay/main.py:23-27
Timestamp: 2024-10-18T16:19:48.637Z
Learning: When initializing `DeepgramClient` without providing an API key, it uses the `DEEPGRAM_API_KEY` environment variable internally.
Applied to files:
src/deepgram/core/client_wrapper.pysrc/deepgram/client.py
📚 Learning: 2024-07-01T19:13:29.909Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/__init__.py:37-38
Timestamp: 2024-07-01T19:13:29.909Z
Learning: Unused imports in `deepgram/clients/listen/__init__.py` are retained to maintain backward compatibility.
Applied to files:
src/deepgram/core/client_wrapper.py.fernignoresrc/deepgram/client.py
📚 Learning: 2024-07-01T19:17:04.194Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/v1/__init__.py:36-43
Timestamp: 2024-07-01T19:17:04.194Z
Learning: Unused imports in `deepgram/clients/listen/v1/__init__.py` are retained to maintain backward compatibility and should not be flagged for removal in reviews.
Applied to files:
.fernignore
📚 Learning: 2024-07-01T19:14:20.539Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/__init__.py:19-24
Timestamp: 2024-07-01T19:14:20.539Z
Learning: Unused imports in `deepgram/clients/listen/__init__.py` are retained to maintain backward compatibility and should not be flagged for removal in reviews.
Applied to files:
.fernignore
📚 Learning: 2024-10-09T02:19:46.086Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/listen/__init__.py:10-13
Timestamp: 2024-10-09T02:19:46.086Z
Learning: Unused imports in `deepgram/clients/listen/__init__.py` are retained for backward compatibility and should not be flagged for removal in reviews.
Applied to files:
.fernignore
📚 Learning: 2024-07-01T19:12:57.715Z
Learnt from: dvonthenen
Repo: deepgram/deepgram-python-sdk PR: 426
File: deepgram/clients/speak/__init__.py:16-19
Timestamp: 2024-07-01T19:12:57.715Z
Learning: Unused imports in `deepgram/clients/speak/__init__.py` are retained for backward compatibility and should not be flagged for removal in reviews.
Applied to files:
src/deepgram/client.py
📚 Learning: 2025-06-22T17:02:32.416Z
Learnt from: lukeocodes
Repo: deepgram/deepgram-python-sdk PR: 543
File: tests/unit_test/test_unit_authentication.py:67-87
Timestamp: 2025-06-22T17:02:32.416Z
Learning: In the Deepgram Python SDK, DeepgramClient("arbitrary-string") should be treated as an API key for backward compatibility with existing code patterns.
Applied to files:
src/deepgram/client.py
📚 Learning: 2025-06-22T17:02:32.416Z
Learnt from: lukeocodes
Repo: deepgram/deepgram-python-sdk PR: 543
File: tests/unit_test/test_unit_authentication.py:67-87
Timestamp: 2025-06-22T17:02:32.416Z
Learning: When both api_key and access_token are explicitly provided via named parameters in the Deepgram Python SDK, access tokens should take precedence according to the documented priority order: explicit access_token parameter (highest), explicit api_key parameter, DEEPGRAM_ACCESS_TOKEN environment variable, and DEEPGRAM_API_KEY environment variable (lowest).
Applied to files:
src/deepgram/client.py
🧬 Code graph analysis (2)
src/deepgram/core/client_wrapper.py (1)
src/deepgram/core/websocket_client.py (2)
WebSocketFactory(68-110)get_default_factory(149-151)
src/deepgram/client.py (2)
src/deepgram/core/client_wrapper.py (6)
AsyncClientWrapper(72-92)BaseClientWrapper(11-46)SyncClientWrapper(49-69)get_custom_headers(39-40)get_environment(42-43)get_timeout(45-46)src/deepgram/core/websocket_client.py (1)
WebSocketFactory(68-110)
🪛 GitHub Actions: CI
src/deepgram/core/websocket_client.py
[error] 76-76: Argument 1 to 'asynccontextmanager' has incompatible type 'Callable[[WebSocketFactory, str, Optional[Dict[str, str]]], Coroutine[Any, Any, AsyncIterator[WebSocketProtocol]]]'; expected 'Callable[[WebSocketFactory, str, Optional[Dict[str, str]]], AsyncIterator[Never]]' [arg-type]
🪛 LanguageTool
sagemaker_bidi_api.md
[grammar] ~21-~21: Ensure spelling is correct
Context: ...you are already familiar with SageMaker publi APIs):** - [Optional] TargetVariant: ...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
[grammar] ~81-~81: Ensure spelling is correct
Context: ...you are already familiar with SageMaker publi APIs):** - InvokedProductionVariant: ...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
🔇 Additional comments (3)
.fernignore (1)
22-27: .fernignore updates for websocket client/custom wrappers look goodThe added comments and the new ignore entry for
src/deepgram/core/websocket_client.pycorrectly document and protect the new WebSocket abstraction and wrapper customizations from regeneration. No functional risk here; just keep these comments in sync with behavior as the WebSocket layer evolves.Also applies to: 30-31, 37-39
src/deepgram/client.py (1)
29-56: Telemetry context version detection is robustThe nested try/except around importing
.versionand falling back to"unknown"provides a safe package_version for telemetry without risking import-time failures, and the outer fallback still guarantees a minimal context. This is a solid improvement with no behavior regressions.src/deepgram/core/websocket_client.py (1)
27-66: WebSocket abstraction design is clear and extension‑friendlyThe split into
WebSocketProtocol/SyncWebSocketProtocolplus aWebSocketFactoryandDefaultWebSocketFactorygives a clean, minimal surface for custom implementations (e.g., SageMaker) while keeping the defaultwebsocketspath well encapsulated. Usingruntime_checkableProtocols and a single shared_default_factoryinstance should play nicely with both type‑checking and runtime checks.Also applies to: 113-151
| } | ||
| ``` | ||
|
|
||
| More detailed explanation **(you can skip this part if you are already familiar with SageMaker publi APIs):** |
There was a problem hiding this comment.
Fix spelling errors in user-facing text.
The word "publi" should be "public" in both instances.
Apply this diff to correct the spelling:
- More detailed explanation **(you can skip this part if you are already familiar with SageMaker publi APIs):**
+ More detailed explanation **(you can skip this part if you are already familiar with SageMaker public APIs):**Apply the same fix to line 81.
Also applies to: 81-81
🧰 Tools
🪛 LanguageTool
[grammar] ~21-~21: Ensure spelling is correct
Context: ...you are already familiar with SageMaker publi APIs):** - [Optional] TargetVariant: ...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
🤖 Prompt for AI Agents
In sagemaker_bidi_api.md around lines 21 and 81 the user-facing text contains
the misspelling "publi" instead of "public"; update both occurrences to read
"public" so the phrases become "SageMaker public APIs" (or equivalent) ensuring
consistent, correct spelling in both locations.
| - An un-fragmented binary chunk always has **`CompletionStatus = COMPLETE`**. | ||
| - Note that SageMaker does not aggregate PayloadParts on customers’ behalf, the fragments are passed to model containers as-is. | ||
| - [Optional] P: | ||
| - Padding string defending against token length side channel attack . |
There was a problem hiding this comment.
Remove encoding artifacts from padding field descriptions.
Lines 50 and 102 contain zero-width or special characters that should be removed for clarity.
Apply this diff to clean up the encoding artifacts:
- - [Optional] P:
- - Padding string defending against token length side channel attack .
+ - [Optional] P:
+ - Padding string defending against token length side channel attack.Apply the same fix to line 102.
Also applies to: 102-102
🤖 Prompt for AI Agents
In sagemaker_bidi_api.md around lines 50 and 102, remove hidden
zero-width/special Unicode characters present in the padding field descriptions
(e.g., after words like "defending against" or elsewhere) so the text reads
normally; open those lines, delete any zero-width spaces or non-printing
characters, and save the file ensuring the descriptions use only regular
ASCII/UTF-8 printable characters and spacing.
| try: | ||
| from websockets.client import connect as websockets_client_connect # type: ignore | ||
| except ImportError: | ||
| from websockets.legacy.client import connect as websockets_client_connect # type: ignore |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
🧩 Analysis chain
Avoid aliasing connect so WebSocket telemetry instrumentation can hook it
Right now the async path imports connect as a function alias:
try:
from websockets.client import connect as websockets_client_connect
except ImportError:
from websockets.legacy.client import connect as websockets_client_connectand then uses it in DefaultWebSocketFactory.connect_async. If apply_websocket_instrumentation monkey‑patches websockets.client.connect / websockets.legacy.client.connect, this alias will bypass the patched function and telemetry won’t see those connections.
You can preserve compatibility with instrumentation by importing the module instead and resolving .connect at call time:
-try:
- from websockets.client import connect as websockets_client_connect # type: ignore
-except ImportError:
- from websockets.legacy.client import connect as websockets_client_connect # type: ignore
+try:
+ import websockets.client as websockets_client # type: ignore
+except ImportError:
+ import websockets.legacy.client as websockets_client # type: ignore
@@
- async with websockets_client_connect(url, extra_headers=headers) as protocol:
+ async with websockets_client.connect(url, extra_headers=headers) as protocol:
yield protocolThe sync path already imports the module (websockets.sync.client) and accesses connect as an attribute, so it should continue to respect any monkey‑patches.
Also applies to: 120-143
🏁 Script executed:
#!/bin/bash
# Sanity‑check that websocket instrumentation still wraps the functions you expect.
set -e
# Show where apply_websocket_instrumentation hooks into websockets:
rg -n "apply_websocket_instrumentation" -C3 src
# Confirm that only module‑level `websockets.client.connect` is imported elsewhere:
rg -n "websockets.client import connect" src || echo "No direct function aliases found (good)."Length of output: 2551
🏁 Script executed:
#!/bin/bash
# Show the implementation of apply_websocket_instrumentation to understand what it patches
rg -n "def apply_websocket_instrumentation" -A30 src/deepgram/extensions/core/instrumented_socket.py
# Show how websockets_client_connect is used in connect_async
rg -n "connect_async" -A15 src/deepgram/core/websocket_client.py
# Show the sync path implementation to compare
rg -n "websockets.sync.client" -B3 -A3 src/deepgram/core/websocket_client.pyLength of output: 247
🏁 Script executed:
#!/bin/bash
# Show the implementation of apply_websocket_instrumentation to understand what it patches
rg -n "def apply_websocket_instrumentation" -A30 src/deepgram/extensions/core/instrumented_socket.py
# Show how websockets_client_connect is used in connect_async
rg -n "connect_async" -A15 src/deepgram/core/websocket_client.py
# Show the sync path implementation to compare
rg -n "websockets.sync.client" -B3 -A3 src/deepgram/core/websocket_client.pyLength of output: 3393
🏁 Script executed:
#!/bin/bash
# Find all usages of websockets_client_connect to confirm scope
rg -n "websockets_client_connect" srcLength of output: 3718
Refactor to support WebSocket telemetry instrumentation
The analysis is correct: importing connect as a function alias prevents telemetry instrumentation from capturing WebSocket connections. Since apply_websocket_instrumentation monkey-patches the module-level connect function at runtime, the aliased import references the pre-patched version and bypasses instrumentation.
Verified details:
apply_websocket_instrumentationpatcheswebsockets.legacy.client.connectandwebsockets.connectat the module level (lines 388 and ~397 ininstrumented_socket.py)- The alias
websockets_client_connectis resolved at import time (before instrumentation runs), so it holds a reference to the original, unpatched function - The sync path already uses module-level import with attribute access (
websockets_sync_client.connect()), which correctly resolves the patched function at call time - This pattern is also present in
instrumented_socket.pyitself (lines 18–20), preventing self-instrumentation
Import the module instead and resolve .connect at call time:
try:
- from websockets.client import connect as websockets_client_connect # type: ignore
+ import websockets.client as websockets_client # type: ignore
except ImportError:
- from websockets.legacy.client import connect as websockets_client_connect # type: ignore
+ import websockets.legacy.client as websockets_client # type: ignoreThen update line 129 to use attribute access:
-async with websockets_client_connect(url, extra_headers=headers) as protocol:
+async with websockets_client.connect(url, extra_headers=headers) as protocol:Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In src/deepgram/core/websocket_client.py around lines 19–22 and update at line
129: the current code aliases the connect function at import time which prevents
runtime monkey-patching for telemetry; change the import to import the module
(e.g., import websockets.client or websockets.legacy.client as a module
variable) so the connect attribute is resolved at call time, and then replace
the direct function call at line 129 to use module.connect (i.e., access the
connect attribute on the imported module when calling) so instrumentation can
patch it at runtime.
| class WebSocketFactory(Protocol): | ||
| """ | ||
| Protocol for WebSocket connection factories. | ||
|
|
||
| Users can implement this protocol to provide custom WebSocket clients | ||
| (e.g., for AWS SageMaker bidirectional streaming). | ||
| """ | ||
|
|
||
| @asynccontextmanager | ||
| async def connect_async( | ||
| self, | ||
| url: str, | ||
| headers: Optional[Dict[str, str]] = None, | ||
| ) -> AsyncIterator[WebSocketProtocol]: | ||
| """ | ||
| Create an async WebSocket connection. | ||
|
|
||
| Args: | ||
| url: The WebSocket URL to connect to | ||
| headers: Optional headers to include in the connection request | ||
|
|
||
| Yields: | ||
| A WebSocket connection implementing WebSocketProtocol | ||
| """ | ||
| ... | ||
|
|
||
| @contextmanager | ||
| def connect_sync( | ||
| self, | ||
| url: str, | ||
| headers: Optional[Dict[str, str]] = None, | ||
| ) -> Iterator[SyncWebSocketProtocol]: | ||
| """ | ||
| Create a synchronous WebSocket connection. | ||
|
|
||
| Args: | ||
| url: The WebSocket URL to connect to | ||
| headers: Optional headers to include in the connection request | ||
|
|
||
| Yields: | ||
| A WebSocket connection implementing SyncWebSocketProtocol | ||
| """ | ||
| ... | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
Fix Protocol connect_* methods to describe context managers (resolves mypy error)
The mypy failure in CI comes from decorating Protocol methods with @asynccontextmanager / @contextmanager. Protocols should only declare the interface; concrete implementations (like DefaultWebSocketFactory) should own the decorators. You can resolve the error and make the expected usage (async with factory.connect_async(...)) explicit by:
- Removing the decorators from
WebSocketFactory.connect_async/connect_sync. - Typing their return values as
AsyncContextManager[WebSocketProtocol]andContextManager[SyncWebSocketProtocol].
A concrete diff:
-from contextlib import asynccontextmanager, contextmanager
-from typing import Any, AsyncIterator, Dict, Iterator, Optional, Protocol, runtime_checkable
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+ Any,
+ AsyncContextManager,
+ AsyncIterator,
+ ContextManager,
+ Dict,
+ Iterator,
+ Optional,
+ Protocol,
+ runtime_checkable,
+)
@@
class WebSocketFactory(Protocol):
@@
- @asynccontextmanager
- async def connect_async(
+ def connect_async(
self,
url: str,
headers: Optional[Dict[str, str]] = None,
- ) -> AsyncIterator[WebSocketProtocol]:
+ ) -> AsyncContextManager[WebSocketProtocol]:
@@
- @contextmanager
- def connect_sync(
+ def connect_sync(
self,
url: str,
headers: Optional[Dict[str, str]] = None,
- ) -> Iterator[SyncWebSocketProtocol]:
+ ) -> ContextManager[SyncWebSocketProtocol]:DefaultWebSocketFactory can continue to use @asynccontextmanager / @contextmanager on its concrete methods; their resulting types (AsyncContextManager[...] / ContextManager[...]) will then structurally satisfy the Protocol.
🏁 Script executed:
#!/bin/bash
# Verify the mypy error is resolved after adjusting WebSocketFactory.
set -e
# Run mypy just for the websocket_client module
mypy src/deepgram/core/websocket_client.pyLength of output: 42
I need to verify the actual code state and the proposed typing approach. Let me inspect the codebase and check the documentation.
Now let me inspect the actual code to verify the current state:
---
🏁 Script executed:
```shell
#!/bin/bash
# Inspect the current websocket_client.py file
cat src/deepgram/core/websocket_client.py | head -120
Length of output: 3356
Let me inspect the DefaultWebSocketFactory implementation and check for other implementations:
---
🏁 Script executed:
```shell
#!/bin/bash
# Get full DefaultWebSocketFactory implementation
sed -n '113,200p' src/deepgram/core/websocket_client.py
Length of output: 1208
🏁 Script executed:
#!/bin/bash
# Search for other implementations of WebSocketFactory
rg -n "class.*WebSocketFactory" src/Length of output: 147
🏁 Script executed:
#!/bin/bash
# Check imports in the file to see current state
rg -n "from typing import|from contextlib import" src/deepgram/core/websocket_client.pyLength of output: 156
Now let me check for usages of WebSocketFactory and the connect_async/connect_sync methods to ensure the proposed changes are compatible:
---
🏁 Script executed:
```shell
#!/bin/bash
# Search for usages of connect_async and connect_sync
rg -n -C2 "\.connect_async\(|\.connect_sync\(" src/
Length of output: 0
🏁 Script executed:
#!/bin/bash
# Search for WebSocketFactory type references
rg -n "WebSocketFactory" src/ --type pyLength of output: 1081
Let me check how the factory methods are actually invoked:
---
🏁 Script executed:
```shell
#!/bin/bash
# Search for async with and with statements that might use factory methods
rg -n "async with|with.*connect" src/deepgram/core/client_wrapper.py -B2 -A2
Length of output: 0
🏁 Script executed:
#!/bin/bash
# More specific search for factory usage patterns
rg -n "factory\." src/deepgram/core/client_wrapper.py -B2 -A2Length of output: 0
---
🏁 Script executed:
```shell
#!/bin/bash
# View client_wrapper.py to see how WebSocketFactory is used
cat src/deepgram/core/client_wrapper.py
Length of output: 3017
---
🏁 Script executed:
```shell
#!/bin/bash
# Search for where websocket_client is used in client.py
rg -n "websocket_client\." src/deepgram/client.py -B3 -A3
Length of output: 0
🏁 Script executed:
#!/bin/bash
# Broader search for connect_async/connect_sync across entire src
rg -rn "connect_async|connect_sync" src/ --type pyLength of output: 207
Fix Protocol connect_* methods to resolve mypy errors
The mypy failure comes from decorating Protocol methods with @asynccontextmanager / @contextmanager. Protocols should only declare the interface; concrete implementations (like DefaultWebSocketFactory) should own the decorators. Resolve this by:
- Removing the decorators from
WebSocketFactory.connect_async/connect_sync(lines 71 and 85) - Changing their return types to
AsyncContextManager[WebSocketProtocol]andContextManager[SyncWebSocketProtocol] - Adding imports for
AsyncContextManagerandContextManagerfrom typing
The concrete DefaultWebSocketFactory can continue using @asynccontextmanager / @contextmanager; its resulting return types will then satisfy the Protocol's interface.
🧰 Tools
🪛 GitHub Actions: CI
[error] 76-76: Argument 1 to 'asynccontextmanager' has incompatible type 'Callable[[WebSocketFactory, str, Optional[Dict[str, str]]], Coroutine[Any, Any, AsyncIterator[WebSocketProtocol]]]'; expected 'Callable[[WebSocketFactory, str, Optional[Dict[str, str]]], AsyncIterator[Never]]' [arg-type]
🤖 Prompt for AI Agents
In src/deepgram/core/websocket_client.py around lines 68 to 111, the Protocol
methods connect_async and connect_sync are incorrectly decorated with
@asynccontextmanager / @contextmanager which causes mypy errors; remove those
decorators, change their return annotations to
AsyncContextManager[WebSocketProtocol] and ContextManager[SyncWebSocketProtocol]
respectively, and add imports for AsyncContextManager and ContextManager from
typing; leave concrete DefaultWebSocketFactory implementations free to use the
decorators so their return types satisfy the Protocol.
Summary by CodeRabbit
New Features
Documentation