caveman-middleware 0.1.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. caveman_middleware-0.1.0a1/LICENSE +21 -0
  2. caveman_middleware-0.1.0a1/PKG-INFO +143 -0
  3. caveman_middleware-0.1.0a1/README.md +95 -0
  4. caveman_middleware-0.1.0a1/caveman_middleware/__init__.py +2 -0
  5. caveman_middleware-0.1.0a1/caveman_middleware/_google_wire.py +84 -0
  6. caveman_middleware-0.1.0a1/caveman_middleware/_httpx2.py +190 -0
  7. caveman_middleware-0.1.0a1/caveman_middleware/_native.py +292 -0
  8. caveman_middleware-0.1.0a1/caveman_middleware/_streams.py +60 -0
  9. caveman_middleware-0.1.0a1/caveman_middleware/_usage.py +91 -0
  10. caveman_middleware-0.1.0a1/caveman_middleware/_versions.py +63 -0
  11. caveman_middleware-0.1.0a1/caveman_middleware/agno.py +557 -0
  12. caveman_middleware-0.1.0a1/caveman_middleware/anthropic.py +158 -0
  13. caveman_middleware-0.1.0a1/caveman_middleware/asgi.py +210 -0
  14. caveman_middleware-0.1.0a1/caveman_middleware/autogen.py +507 -0
  15. caveman_middleware-0.1.0a1/caveman_middleware/crewai.py +338 -0
  16. caveman_middleware-0.1.0a1/caveman_middleware/google.py +493 -0
  17. caveman_middleware-0.1.0a1/caveman_middleware/langchain.py +398 -0
  18. caveman_middleware-0.1.0a1/caveman_middleware/litellm.py +375 -0
  19. caveman_middleware-0.1.0a1/caveman_middleware/llama_index.py +807 -0
  20. caveman_middleware-0.1.0a1/caveman_middleware/mcp.py +126 -0
  21. caveman_middleware-0.1.0a1/caveman_middleware/openai.py +215 -0
  22. caveman_middleware-0.1.0a1/caveman_middleware/py.typed +0 -0
  23. caveman_middleware-0.1.0a1/caveman_middleware/pydantic_ai.py +310 -0
  24. caveman_middleware-0.1.0a1/caveman_middleware/strands.py +210 -0
  25. caveman_middleware-0.1.0a1/caveman_middleware.egg-info/PKG-INFO +143 -0
  26. caveman_middleware-0.1.0a1/caveman_middleware.egg-info/SOURCES.txt +32 -0
  27. caveman_middleware-0.1.0a1/caveman_middleware.egg-info/dependency_links.txt +1 -0
  28. caveman_middleware-0.1.0a1/caveman_middleware.egg-info/requires.txt +45 -0
  29. caveman_middleware-0.1.0a1/caveman_middleware.egg-info/top_level.txt +1 -0
  30. caveman_middleware-0.1.0a1/pyproject.toml +51 -0
  31. caveman_middleware-0.1.0a1/setup.cfg +4 -0
  32. caveman_middleware-0.1.0a1/tests/test_adapters.py +36 -0
  33. caveman_middleware-0.1.0a1/tests/test_fail_open.py +125 -0
  34. caveman_middleware-0.1.0a1/tests/test_versions.py +35 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Julius Brussee
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,143 @@
1
+ Metadata-Version: 2.4
2
+ Name: caveman-middleware
3
+ Version: 0.1.0a1
4
+ Summary: Native framework adapters for the Caveman compression runtime.
5
+ License-Expression: MIT
6
+ Project-URL: Repository, https://github.com/JuliusBrussee/caveman
7
+ Project-URL: Issues, https://github.com/JuliusBrussee/caveman/issues
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Programming Language :: Python :: 3.13
11
+ Classifier: Typing :: Typed
12
+ Requires-Python: >=3.13
13
+ Description-Content-Type: text/markdown
14
+ License-File: LICENSE
15
+ Requires-Dist: caveman-sdk<2,>=1.1
16
+ Provides-Extra: openai
17
+ Requires-Dist: openai<4,>=2.20; extra == "openai"
18
+ Provides-Extra: anthropic
19
+ Requires-Dist: anthropic<2,>=1.4; extra == "anthropic"
20
+ Provides-Extra: google
21
+ Requires-Dist: google-genai<3,>=2.22; extra == "google"
22
+ Provides-Extra: langchain
23
+ Requires-Dist: langchain<2,>=1.4; extra == "langchain"
24
+ Requires-Dist: langchain-core<2,>=1.6; extra == "langchain"
25
+ Requires-Dist: langgraph<2,>=1.2; extra == "langchain"
26
+ Provides-Extra: litellm
27
+ Requires-Dist: litellm<2,>=1.100; extra == "litellm"
28
+ Provides-Extra: strands
29
+ Requires-Dist: strands-agents<2,>=1.55; extra == "strands"
30
+ Provides-Extra: agno
31
+ Requires-Dist: agno<4,>=3.0; extra == "agno"
32
+ Provides-Extra: asgi
33
+ Requires-Dist: fastapi<1,>=0.141; extra == "asgi"
34
+ Requires-Dist: starlette<2,>=1.6; extra == "asgi"
35
+ Provides-Extra: mcp
36
+ Requires-Dist: mcp<3,>=2.2; extra == "mcp"
37
+ Provides-Extra: crewai
38
+ Requires-Dist: crewai<2,>=1.15; extra == "crewai"
39
+ Provides-Extra: pydantic-ai
40
+ Requires-Dist: pydantic-ai<3,>=2.42; extra == "pydantic-ai"
41
+ Provides-Extra: autogen
42
+ Requires-Dist: autogen-agentchat<0.8,>=0.7; extra == "autogen"
43
+ Requires-Dist: autogen-core<0.8,>=0.7; extra == "autogen"
44
+ Requires-Dist: autogen-ext<0.8,>=0.7; extra == "autogen"
45
+ Provides-Extra: llama-index
46
+ Requires-Dist: llama-index-core<0.15,>=0.14; extra == "llama-index"
47
+ Dynamic: license-file
48
+
49
+ # Caveman Python middleware
50
+
51
+ Native framework adapters over `caveman_cloud.middleware`. Your framework keeps
52
+ its models, tool loop, retries, streams and stored conversation. Caveman only
53
+ changes the text of selected tool results in the outbound model view, and only
54
+ when the local runtime answers in time.
55
+
56
+ Alpha. Install the extra for the framework you use:
57
+
58
+ ```sh
59
+ pip install "caveman-middleware[langchain]"
60
+ ```
61
+
62
+ ```python
63
+ from caveman_cloud.middleware import MiddlewareRuntime, Scope
64
+ from caveman_middleware.langchain import with_caveman_agent
65
+ from langchain.agents import create_agent
66
+
67
+ runtime = MiddlewareRuntime(endpoint="http://127.0.0.1:8787")
68
+ runtime.ready()
69
+ scope = Scope("my-app", session_id="conversation-1", branch_id="main", cache_epoch="0")
70
+
71
+ agent = create_agent(**with_caveman_agent(existing_agent_options, runtime=runtime, scope=scope))
72
+ ```
73
+
74
+ Call `runtime.close()` at shutdown. Nothing in your stored message history is
75
+ rewritten; only the request handed to the provider is.
76
+
77
+ ## Supported frameworks
78
+
79
+ One extra per adapter. The range is the band the adapter's version gate accepts;
80
+ the left-hand value is the version its wire handling was written against.
81
+
82
+ | Extra | Module | Tested range |
83
+ |---|---|---|
84
+ | `langchain` | `caveman_middleware.langchain` | `langchain>=1.4,<2`, `langchain-core>=1.6,<2`, `langgraph>=1.2,<2` |
85
+ | `litellm` | `caveman_middleware.litellm` | `litellm>=1.100,<2` |
86
+ | `openai` | `caveman_middleware.openai` | `openai>=3.10,<4` |
87
+ | `anthropic` | `caveman_middleware.anthropic` | `anthropic>=1.4,<2` |
88
+ | `google` | `caveman_middleware.google` | `google-genai>=2.22,<3` |
89
+ | `strands` | `caveman_middleware.strands` | `strands-agents>=1.55,<2` |
90
+ | `agno` | `caveman_middleware.agno` | `agno>=3.0,<4` |
91
+ | `crewai` | `caveman_middleware.crewai` | `crewai>=1.15,<2` |
92
+ | `autogen` | `caveman_middleware.autogen` | `autogen-core`, `autogen-agentchat`, `autogen-ext` all `>=0.7,<0.8` |
93
+ | `pydantic-ai` | `caveman_middleware.pydantic_ai` | `pydantic-ai-slim>=2.42,<3` |
94
+ | `llama-index` | `caveman_middleware.llama_index` | `llama-index-core>=0.14,<0.15` |
95
+ | `asgi` | `caveman_middleware.asgi` | `fastapi>=0.141,<1`, `starlette>=1.6,<2` |
96
+ | `mcp` | `caveman_middleware.mcp` | `mcp>=2.2,<3` |
97
+
98
+ Outside its range an adapter keeps the caller's native input unchanged and
99
+ reports `unsupported_version`. It never raises.
100
+
101
+ Extras carry no provider pins, so most combine: `pip install
102
+ "caveman-middleware[langchain,agno,litellm]"` resolves. Three pairs cannot, and
103
+ the cause is upstream, not this package: `strands-agents` and `crewai` both
104
+ require `mcp<2.2` while the MCP adapter needs `>=2.2`, and `agno` and `crewai`
105
+ disagree over `litellm`. Install those families in separate environments.
106
+
107
+ ## What you lose
108
+
109
+ Read this before turning on `compress`.
110
+
111
+ - **Markers are scoped.** A replacement marker is bound to
112
+ `namespace` + `session_id` + `branch_id` + `cache_epoch` and to the
113
+ authenticated principal that created it. Replaying it under any other scope
114
+ returns nothing.
115
+ - **Markers expire.** A scope lives 24 hours, renewed each time it is used.
116
+ After that the stored original is released.
117
+ - **Persisted compressed history is a trap.** If you save the transformed
118
+ messages and replay them later — under a different scope, or after the scope
119
+ expired — that text is gone for good. Caveman does not rewrite your stored
120
+ history for this reason: keep the originals, let Caveman transform the
121
+ outbound copy on every call.
122
+ - **Recovery needs a registered executor.** In `compress` mode a lossy
123
+ transform is only used when your framework really holds the
124
+ `caveman_retrieve` tool. Without it, only recovery-free transforms apply.
125
+ - **Eligible content is narrow.** Successful tool-result text and explicitly
126
+ passed document bodies. Never system prompts, user messages, assistant
127
+ reasoning, errored tool results, images or other non-text parts, anything
128
+ marked protected, and never a payload the adapter cannot parse exactly.
129
+ - **Failure is silent and safe.** Runtime down, over deadline, out of capacity,
130
+ expired scope: the original request goes to the provider and the decision is
131
+ reported as skipped. Set `strict=True` only if you want a pre-inference error
132
+ instead.
133
+
134
+ ## Reports
135
+
136
+ `on_report=` receives one immutable record per native call: status, reason,
137
+ transform ids, replacement and reuse counts. No original content, no provider
138
+ credentials. `runtime.last_report` holds the most recent one. A callback that
139
+ raises cannot affect inference.
140
+
141
+ Token counts in reports are estimates from the runtime's tokenizer over the
142
+ segments it saw. They are not measured billing savings, and this package does
143
+ not claim a reduction figure.
@@ -0,0 +1,95 @@
1
+ # Caveman Python middleware
2
+
3
+ Native framework adapters over `caveman_cloud.middleware`. Your framework keeps
4
+ its models, tool loop, retries, streams and stored conversation. Caveman only
5
+ changes the text of selected tool results in the outbound model view, and only
6
+ when the local runtime answers in time.
7
+
8
+ Alpha. Install the extra for the framework you use:
9
+
10
+ ```sh
11
+ pip install "caveman-middleware[langchain]"
12
+ ```
13
+
14
+ ```python
15
+ from caveman_cloud.middleware import MiddlewareRuntime, Scope
16
+ from caveman_middleware.langchain import with_caveman_agent
17
+ from langchain.agents import create_agent
18
+
19
+ runtime = MiddlewareRuntime(endpoint="http://127.0.0.1:8787")
20
+ runtime.ready()
21
+ scope = Scope("my-app", session_id="conversation-1", branch_id="main", cache_epoch="0")
22
+
23
+ agent = create_agent(**with_caveman_agent(existing_agent_options, runtime=runtime, scope=scope))
24
+ ```
25
+
26
+ Call `runtime.close()` at shutdown. Nothing in your stored message history is
27
+ rewritten; only the request handed to the provider is.
28
+
29
+ ## Supported frameworks
30
+
31
+ One extra per adapter. The range is the band the adapter's version gate accepts;
32
+ the left-hand value is the version its wire handling was written against.
33
+
34
+ | Extra | Module | Tested range |
35
+ |---|---|---|
36
+ | `langchain` | `caveman_middleware.langchain` | `langchain>=1.4,<2`, `langchain-core>=1.6,<2`, `langgraph>=1.2,<2` |
37
+ | `litellm` | `caveman_middleware.litellm` | `litellm>=1.100,<2` |
38
+ | `openai` | `caveman_middleware.openai` | `openai>=3.10,<4` |
39
+ | `anthropic` | `caveman_middleware.anthropic` | `anthropic>=1.4,<2` |
40
+ | `google` | `caveman_middleware.google` | `google-genai>=2.22,<3` |
41
+ | `strands` | `caveman_middleware.strands` | `strands-agents>=1.55,<2` |
42
+ | `agno` | `caveman_middleware.agno` | `agno>=3.0,<4` |
43
+ | `crewai` | `caveman_middleware.crewai` | `crewai>=1.15,<2` |
44
+ | `autogen` | `caveman_middleware.autogen` | `autogen-core`, `autogen-agentchat`, `autogen-ext` all `>=0.7,<0.8` |
45
+ | `pydantic-ai` | `caveman_middleware.pydantic_ai` | `pydantic-ai-slim>=2.42,<3` |
46
+ | `llama-index` | `caveman_middleware.llama_index` | `llama-index-core>=0.14,<0.15` |
47
+ | `asgi` | `caveman_middleware.asgi` | `fastapi>=0.141,<1`, `starlette>=1.6,<2` |
48
+ | `mcp` | `caveman_middleware.mcp` | `mcp>=2.2,<3` |
49
+
50
+ Outside its range an adapter keeps the caller's native input unchanged and
51
+ reports `unsupported_version`. It never raises.
52
+
53
+ Extras carry no provider pins, so most combine: `pip install
54
+ "caveman-middleware[langchain,agno,litellm]"` resolves. Three pairs cannot, and
55
+ the cause is upstream, not this package: `strands-agents` and `crewai` both
56
+ require `mcp<2.2` while the MCP adapter needs `>=2.2`, and `agno` and `crewai`
57
+ disagree over `litellm`. Install those families in separate environments.
58
+
59
+ ## What you lose
60
+
61
+ Read this before turning on `compress`.
62
+
63
+ - **Markers are scoped.** A replacement marker is bound to
64
+ `namespace` + `session_id` + `branch_id` + `cache_epoch` and to the
65
+ authenticated principal that created it. Replaying it under any other scope
66
+ returns nothing.
67
+ - **Markers expire.** A scope lives 24 hours, renewed each time it is used.
68
+ After that the stored original is released.
69
+ - **Persisted compressed history is a trap.** If you save the transformed
70
+ messages and replay them later — under a different scope, or after the scope
71
+ expired — that text is gone for good. Caveman does not rewrite your stored
72
+ history for this reason: keep the originals, let Caveman transform the
73
+ outbound copy on every call.
74
+ - **Recovery needs a registered executor.** In `compress` mode a lossy
75
+ transform is only used when your framework really holds the
76
+ `caveman_retrieve` tool. Without it, only recovery-free transforms apply.
77
+ - **Eligible content is narrow.** Successful tool-result text and explicitly
78
+ passed document bodies. Never system prompts, user messages, assistant
79
+ reasoning, errored tool results, images or other non-text parts, anything
80
+ marked protected, and never a payload the adapter cannot parse exactly.
81
+ - **Failure is silent and safe.** Runtime down, over deadline, out of capacity,
82
+ expired scope: the original request goes to the provider and the decision is
83
+ reported as skipped. Set `strict=True` only if you want a pre-inference error
84
+ instead.
85
+
86
+ ## Reports
87
+
88
+ `on_report=` receives one immutable record per native call: status, reason,
89
+ transform ids, replacement and reuse counts. No original content, no provider
90
+ credentials. `runtime.last_report` holds the most recent one. A callback that
91
+ raises cannot affect inference.
92
+
93
+ Token counts in reports are estimates from the runtime's tokenizer over the
94
+ segments it saw. They are not measured billing savings, and this package does
95
+ not claim a reduction figure.
@@ -0,0 +1,2 @@
1
+ """Optional native adapters. Import only the framework module your app uses."""
2
+ __version__ = "0.1.0a1"
@@ -0,0 +1,84 @@
1
+ """Bounded Google JSON string offsets; untouched serialized text stays exact."""
2
+ import json
3
+
4
+
5
+ def parse(text):
6
+ if len(text) > 2 << 20:
7
+ return None
8
+ decoder, strings, index, nodes = json.JSONDecoder(), {}, 0, 0
9
+
10
+ def white():
11
+ nonlocal index
12
+ while index < len(text) and text[index] in " \t\r\n":
13
+ index += 1
14
+
15
+ def walk(path, depth=0):
16
+ nonlocal index, nodes
17
+ nodes += 1
18
+ if depth > 64 or nodes > 65536:
19
+ raise ValueError("bounded")
20
+ white()
21
+ start = index
22
+ if text[index] == '"':
23
+ value, index = decoder.raw_decode(text, index)
24
+ strings[path] = (start, index, value)
25
+ return value
26
+ if text[index] == "{":
27
+ index += 1
28
+ result = {}
29
+ white()
30
+ if text[index] == "}":
31
+ index += 1
32
+ return result
33
+ while True:
34
+ white()
35
+ key, index = decoder.raw_decode(text, index)
36
+ if type(key) is not str or key in result:
37
+ raise ValueError("duplicate")
38
+ white()
39
+ if text[index] != ":":
40
+ raise ValueError("colon")
41
+ index += 1
42
+ result[key] = walk((*path, key), depth + 1)
43
+ white()
44
+ end = text[index]
45
+ index += 1
46
+ if end == "}":
47
+ return result
48
+ if end != ",":
49
+ raise ValueError("comma")
50
+ if text[index] == "[":
51
+ index += 1
52
+ result = []
53
+ white()
54
+ if text[index] == "]":
55
+ index += 1
56
+ return result
57
+ while True:
58
+ result.append(walk((*path, len(result)), depth + 1))
59
+ white()
60
+ end = text[index]
61
+ index += 1
62
+ if end == "]":
63
+ return result
64
+ if end != ",":
65
+ raise ValueError("comma")
66
+ value, index = decoder.raw_decode(text, index)
67
+ return value
68
+
69
+ try:
70
+ value = walk(())
71
+ white()
72
+ return (value, strings) if index == len(text) else None
73
+ except (ValueError, IndexError, RecursionError):
74
+ return None
75
+
76
+
77
+ def patch(text, replacements):
78
+ end = len(text)
79
+ for (start, stop, original), replacement in sorted(replacements, reverse=True):
80
+ if start < 0 or stop > end or start >= stop or json.loads(text[start:stop]) != original:
81
+ return None
82
+ text = text[:start] + json.dumps(replacement, ensure_ascii=False) + text[stop:]
83
+ end = start
84
+ return text
@@ -0,0 +1,190 @@
1
+ """Pull-driven response accounting for the pinned native provider SDK transport."""
2
+ from __future__ import annotations
3
+
4
+ import httpx2
5
+ import asyncio
6
+ import hashlib
7
+ import uuid
8
+
9
+ from ._native import Attempt, owner
10
+
11
+
12
+ from ._usage import UsageReader
13
+
14
+
15
+ class OpenAICall:
16
+ """One prepared native operation and its observed physical HTTP attempts."""
17
+ def __init__(self, attempt, transport):
18
+ self.attempt, self.transport, self.attempt_count = attempt, transport, 0
19
+
20
+ def next_attempt(self, request):
21
+ original = self.attempt
22
+ attempt = original if self.attempt_count == 0 else Attempt(original.runtime, original.scope,
23
+ original.logical_call_id, str(uuid.uuid4()), optimization=original.optimization, plan_id=original.plan_id)
24
+ self.attempt_count += 1
25
+ try:
26
+ attempt.wire_sha256 = hashlib.sha256(request.content).hexdigest()
27
+ except httpx2.RequestNotRead:
28
+ pass
29
+ return attempt
30
+
31
+
32
+ class CavemanOpenAITransport(httpx2.BaseTransport):
33
+ """Wrap an application-owned HTTPX2 transport before client construction.
34
+
35
+ Pass this same object as ``transport`` to ``with_caveman_openai``. Requests
36
+ outside that native wrapper pass through. The supplied transport retains its
37
+ connection pool, TLS/proxy settings, and close behavior; this adds no retries.
38
+ """
39
+ def __init__(self, transport):
40
+ if not isinstance(transport, httpx2.BaseTransport):
41
+ raise TypeError("Expected an HTTPX2 BaseTransport")
42
+ self.transport = transport
43
+
44
+ def handle_request(self, request):
45
+ call = owner.get()
46
+ if not isinstance(call, OpenAICall) or call.transport is not self:
47
+ return self.transport.handle_request(request)
48
+ attempt = call.next_attempt(request)
49
+ attempt.observe("dispatch_intent")
50
+ token = owner.set(attempt)
51
+ try:
52
+ response = self.transport.handle_request(request)
53
+ observe_response(response, attempt)
54
+ return response
55
+ except BaseException:
56
+ attempt.observe("failed")
57
+ raise
58
+ finally:
59
+ owner.reset(token)
60
+
61
+ def close(self):
62
+ self.transport.close()
63
+
64
+
65
+ class CavemanAsyncOpenAITransport(httpx2.AsyncBaseTransport):
66
+ """Async counterpart of CavemanOpenAITransport, using the native pool."""
67
+ def __init__(self, transport):
68
+ if not isinstance(transport, httpx2.AsyncBaseTransport):
69
+ raise TypeError("Expected an HTTPX2 AsyncBaseTransport")
70
+ self.transport = transport
71
+
72
+ async def handle_async_request(self, request):
73
+ call = owner.get()
74
+ if not isinstance(call, OpenAICall) or call.transport is not self:
75
+ return await self.transport.handle_async_request(request)
76
+ attempt = call.next_attempt(request)
77
+ attempt.observe("dispatch_intent")
78
+ token = owner.set(attempt)
79
+ try:
80
+ response = await self.transport.handle_async_request(request)
81
+ observe_response(response, attempt)
82
+ return response
83
+ except (asyncio.CancelledError, GeneratorExit):
84
+ attempt.observe("cancelled")
85
+ raise
86
+ except BaseException:
87
+ attempt.observe("failed")
88
+ raise
89
+ finally:
90
+ owner.reset(token)
91
+
92
+ async def aclose(self):
93
+ await self.transport.aclose()
94
+
95
+
96
+ class SyncObservedStream(httpx2.SyncByteStream):
97
+ def __init__(self, original, attempt, parser):
98
+ self.original, self.attempt, self.parser, self.ended = original, attempt, parser, False
99
+
100
+ def finish(self, event):
101
+ if not self.ended:
102
+ self.ended = True
103
+ self.attempt.observe(event, self.parser.finish() if event == "completed" else None)
104
+
105
+ def __iter__(self):
106
+ iterator = iter(self.original)
107
+ try:
108
+ while True:
109
+ token = owner.set(self.attempt)
110
+ try:
111
+ chunk = next(iterator)
112
+ except StopIteration:
113
+ self.finish("completed")
114
+ return
115
+ finally:
116
+ owner.reset(token)
117
+ self.parser.feed(chunk)
118
+ yield chunk
119
+ except GeneratorExit:
120
+ self.finish("completed" if self.parser.terminal else "cancelled")
121
+ raise
122
+ except BaseException:
123
+ self.finish("failed")
124
+ raise
125
+
126
+ def close(self):
127
+ try:
128
+ self.original.close()
129
+ finally:
130
+ self.finish("completed" if self.parser.terminal else "cancelled")
131
+
132
+
133
+ class AsyncObservedStream(httpx2.AsyncByteStream):
134
+ def __init__(self, original, attempt, parser):
135
+ self.original, self.attempt, self.parser, self.ended = original, attempt, parser, False
136
+
137
+ def finish(self, event):
138
+ if not self.ended:
139
+ self.ended = True
140
+ self.attempt.observe(event, self.parser.finish() if event == "completed" else None)
141
+
142
+ async def __aiter__(self):
143
+ iterator = self.original.__aiter__()
144
+ try:
145
+ while True:
146
+ token = owner.set(self.attempt)
147
+ try:
148
+ chunk = await anext(iterator)
149
+ except StopAsyncIteration:
150
+ self.finish("completed")
151
+ return
152
+ finally:
153
+ owner.reset(token)
154
+ self.parser.feed(chunk)
155
+ yield chunk
156
+ except GeneratorExit:
157
+ self.finish("completed" if self.parser.terminal else "cancelled")
158
+ raise
159
+ except asyncio.CancelledError:
160
+ self.finish("cancelled")
161
+ raise
162
+ except BaseException:
163
+ self.finish("failed")
164
+ raise
165
+
166
+ async def aclose(self):
167
+ try:
168
+ await self.original.aclose()
169
+ finally:
170
+ self.finish("completed" if self.parser.terminal else "cancelled")
171
+
172
+
173
+ def observe_response(response, attempt):
174
+ if not isinstance(response, httpx2.Response):
175
+ return
176
+ if response.status_code >= 400:
177
+ attempt.observe("failed")
178
+ return
179
+ content_type = response.headers.get("content-type", "")
180
+ if "json" not in content_type and "text/event-stream" not in content_type:
181
+ attempt.observe("completed")
182
+ return
183
+ parser = UsageReader("text/event-stream" in content_type)
184
+ if response.is_stream_consumed:
185
+ parser.feed(response.content)
186
+ attempt.observe("completed", parser.finish())
187
+ elif isinstance(response.stream, httpx2.SyncByteStream):
188
+ response.stream = SyncObservedStream(response.stream, attempt, parser)
189
+ elif isinstance(response.stream, httpx2.AsyncByteStream):
190
+ response.stream = AsyncObservedStream(response.stream, attempt, parser)