Skip to content

Commit

Permalink
Re-implemented LLM_OPENAI_SHOW_RESPONSES against httpx
Browse files Browse the repository at this point in the history
Closes #404
Refs #364
  • Loading branch information
simonw committed Jan 26, 2024
1 parent 214fcaa commit 1f67f14
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 2 deletions.
5 changes: 4 additions & 1 deletion llm/default_plugins/openai_models.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from llm import EmbeddingModel, Model, hookimpl
import llm
from llm.utils import dicts_to_table_string, remove_dict_none_values
from llm.utils import dicts_to_table_string, remove_dict_none_values, logging_client
import click
import datetime
import httpx
import openai
import os

try:
# Pydantic 2
Expand Down Expand Up @@ -342,6 +343,8 @@ def get_client(self):
kwargs["api_key"] = "DUMMY_KEY"
if self.headers:
kwargs["headers"] = self.headers
if os.environ.get("LLM_OPENAI_SHOW_RESPONSES"):
kwargs["http_client"] = logging_client()
return openai.OpenAI(**kwargs)

def build_kwargs(self, prompt):
Expand Down
52 changes: 51 additions & 1 deletion llm/utils.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,9 @@
from typing import List, Dict
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(
Expand Down Expand Up @@ -43,3 +48,48 @@ def remove_dict_none_values(d):
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_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_encoding], "response": [_log_response]}
)

0 comments on commit 1f67f14

Please sign in to comment.