39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from typing import Any, Dict, List, Optional
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
|
|
|
|
class _WireBodyCapture(CustomLogger):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.records: List[Dict[str, Any]] = []
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
self.records.append(
|
|
{
|
|
"model": model,
|
|
"body": kwargs.get("additional_args", {}).get("complete_input_dict"),
|
|
"api_base": kwargs.get("additional_args", {}).get("api_base"),
|
|
}
|
|
)
|
|
|
|
async def async_log_pre_api_call(self, model, messages, kwargs):
|
|
self.log_pre_api_call(model, messages, kwargs)
|
|
|
|
def latest(self) -> Optional[Dict[str, Any]]:
|
|
return self.records[-1] if self.records else None
|
|
|
|
|
|
@pytest.fixture()
|
|
def wire_capture():
|
|
capture = _WireBodyCapture()
|
|
previous = list(litellm.callbacks)
|
|
litellm.callbacks = previous + [capture]
|
|
try:
|
|
yield capture
|
|
finally:
|
|
litellm.callbacks = previous
|