diff --git a/llm/default_plugins/openai_models.py b/llm/default_plugins/openai_models.py index d067109b..b1c66e26 100644 --- a/llm/default_plugins/openai_models.py +++ b/llm/default_plugins/openai_models.py @@ -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 @@ -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): diff --git a/llm/utils.py b/llm/utils.py index 480aa0f8..00ee95b2 100644 --- a/llm/utils.py +++ b/llm/utils.py @@ -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( @@ -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]} + )