Repository navigation
Capture AI usage from HTTP responses - #780
iacobdaniel wants to merge 13 commits into
Conversation
![aikido-pr-checks[bot]](https://avatars.githubusercontent.com/in/898896?s=60&v=4){kind=small}
| def patch_client(m): | ||
| """Patches the sync and async clients of httpx, or of httpx2 (same API)""" | ||
| patch_function(m, "Client.send", _send) | ||
| patch_function(m, "AsyncClient.send", _async_send) |
There was a problem hiding this comment.
worth adding a test to ensure HTTP client instrumentation doesn't break anything? just like some basic HTTP client usage
There was a problem hiding this comment.
Added tests that plain GET/POST calls through requests, httpx, httpx2 (sync and async) and urllib3 return the same status, headers and body with the hooks installed, and that connection errors still raise the client's own exception.
| @on_import("requests") | ||
| def patch(m): | ||
| if is_ai_network_capture_enabled(): | ||
| patch_function(m, "Session.send", _send) |
There was a problem hiding this comment.
worth adding a test to ensure HTTP client instrumentation doesn't break anything? just like some basic HTTP client usage
There was a problem hiding this comment.
Covered by the same new tests.
| @on_import("urllib3.connectionpool") | ||
| def patch(m): | ||
| if is_ai_network_capture_enabled(): | ||
| patch_function(m, "HTTPConnectionPool.urlopen", _urlopen) |
There was a problem hiding this comment.
worth adding a test to ensure HTTP client instrumentation doesn't break anything? just like some basic HTTP client usage
There was a problem hiding this comment.
Covered by the same tests.
| if cache_read: | ||
| tokens["cacheRead"] = tokens.get("cacheRead", 0) + cache_read | ||
| if cache_write: | ||
| tokens["cacheWrite"] = tokens.get("cacheWrite", 0) + cache_write |
There was a problem hiding this comment.
I think it's safe to report these as 0? * is still 0?
There was a problem hiding this comment.
Yes, it's safe. After the fix every entry reports cacheRead and cacheWrite, starting at 0, so the code no longer needs a special case to leave them out.
| DEFAULT_PORTS = {"http": 80, "https": 443} | ||
|
|
||
| # A redirect makes requests and urllib3 pass the final response through twice. | ||
| SEEN_ATTR = "_aikido_ai_seen" |
There was a problem hiding this comment.
Yes, both are needed.
DEFAULT_PORTS exists because the HTTP clients describe the same address differently. urllib3 always gives the port (443), httpx never does. We leave out the scheme's standard port (80 for http, 443 for https), as the URL standard.
And without SEEN_ATTR, a redirected call through requests or an urllib3 pool is counted twice, because the same final response comes back. Added a test for it also.
There are alternative solutions, we can write them inline in the code if we do not want the constants, but would make the code a bit more complicated and I think a constant with a name keeps each one clear.
| "import aikido_zen.sinks.urllib3\n" | ||
| "targets = [httpx.Client.send, httpx.AsyncClient.send, requests.Session.send,\n" | ||
| " urllib3.connectionpool.HTTPConnectionPool.urlopen]\n" | ||
| "print([hasattr(t, '__wrapped__') for t in targets])\n" |
There was a problem hiding this comment.
not sure if we need to test whether these methods are wrapped
if the flag is enabled, call the HTTP client like you normally would and see if nothing breaks, that's fine! (and then indeed call the get_ai_stats method to see what we've captured)
Also best to move the client specific bits into separate files per sink?
There was a problem hiding this comment.
Done. The tests call each client normally with the flag on and check get_ai_stats, and each client has its own file. I kept the wrapped check for the flag-off case, since it's the only test that proves the clients stay untouched without the flag.
| @@ -0,0 +1,49 @@ | |||
| """ | |||
| patching module httpx, when AIKIDO_FEATURE_AI_NETWORK_CAPTURE is on | |||
There was a problem hiding this comment.
Let's revisit whether these comments are actually needed, if we remove the feature flag, it might not be updated? only add info that the code doesn't tell?
There was a problem hiding this comment.
Removed them, the remaining comments only explain why.
Count AI API calls from the responses httpx, httpx2, requests and urllib3 receive, when AIKIDO_FEATURE_AI_NETWORK_CAPTURE is on (off by default). A POST whose path ends like an AI API counts on any host and is reported with its endpoint URL as provider, with its tokens, and with cached tokens as cacheRead/cacheWrite when the provider reports them. The SDK sinks are unchanged, so with the flag on an SDK call shows up under both its URL and the SDK name. Without the flag nothing is patched; with it, an HTTP call costs about 1 µs more (httpx 23.1 → 24.3 µs, requests no measurable change).