simonw--llm
Simon Willison
96 行
3.0 KiB
Python
96 行
3.0 KiB
Python
import click
|
|
import httpx
|
|
from httpx._transports.default import ResponseStream
|
|
import json
|
|
import textwrap
|
|
from typing import List, Dict, Iterator
|
|
|
|
|
|
def dicts_to_table_string(
|
|
headings: List[str], dicts: List[Dict[str, str]]
|
|
) -> List[str]:
|
|
max_lengths = [len(h) for h in headings]
|
|
|
|
# Compute maximum length for each column
|
|
for d in dicts:
|
|
for i, h in enumerate(headings):
|
|
if h in d and len(str(d[h])) > max_lengths[i]:
|
|
max_lengths[i] = len(str(d[h]))
|
|
|
|
# Generate formatted table strings
|
|
res = []
|
|
res.append(" ".join(h.ljust(max_lengths[i]) for i, h in enumerate(headings)))
|
|
|
|
for d in dicts:
|
|
row = []
|
|
for i, h in enumerate(headings):
|
|
row.append(str(d.get(h, "")).ljust(max_lengths[i]))
|
|
res.append(" ".join(row))
|
|
|
|
return res
|
|
|
|
|
|
def remove_dict_none_values(d):
|
|
"""
|
|
Recursively remove keys with value of None or value of a dict that is all values of None
|
|
"""
|
|
if not isinstance(d, dict):
|
|
return d
|
|
new_dict = {}
|
|
for key, value in d.items():
|
|
if value is not None:
|
|
if isinstance(value, dict):
|
|
nested = remove_dict_none_values(value)
|
|
if nested:
|
|
new_dict[key] = nested
|
|
elif isinstance(value, list):
|
|
new_dict[key] = [remove_dict_none_values(v) for v in value]
|
|
else:
|
|
new_dict[key] = value
|
|
return new_dict
|
|
|
|
|
|
class _LoggingStream(ResponseStream):
|
|
def __iter__(self) -> Iterator[bytes]:
|
|
for chunk in super().__iter__():
|
|
click.echo(f" {chunk.decode()}", err=True)
|
|
yield chunk
|
|
|
|
|
|
def _no_accept_encoding(request: httpx.Request):
|
|
request.headers.pop("accept-encoding", None)
|
|
|
|
|
|
def _log_response(response: httpx.Response):
|
|
request = response.request
|
|
click.echo(f"Request: {request.method} {request.url}", err=True)
|
|
click.echo(" Headers:", err=True)
|
|
for key, value in request.headers.items():
|
|
if key.lower() == "authorization":
|
|
value = "[...]"
|
|
if key.lower() == "cookie":
|
|
value = value.split("=")[0] + "=..."
|
|
click.echo(f" {key}: {value}", err=True)
|
|
click.echo(" Body:", err=True)
|
|
try:
|
|
request_body = json.loads(request.content)
|
|
click.echo(
|
|
textwrap.indent(json.dumps(request_body, indent=2), " "), err=True
|
|
)
|
|
except json.JSONDecodeError:
|
|
click.echo(textwrap.indent(request.content.decode(), " "), err=True)
|
|
click.echo(f"Response: status_code={response.status_code}", err=True)
|
|
click.echo(" Headers:", err=True)
|
|
for key, value in response.headers.items():
|
|
if key.lower() == "set-cookie":
|
|
value = value.split("=")[0] + "=..."
|
|
click.echo(f" {key}: {value}", err=True)
|
|
click.echo(" Body:", err=True)
|
|
response.stream._stream = _LoggingStream(response.stream._stream) # type: ignore
|
|
|
|
|
|
def logging_client() -> httpx.Client:
|
|
return httpx.Client(
|
|
event_hooks={"request": [_no_accept_encoding], "response": [_log_response]}
|
|
)
|