mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
test(ocr): integrate parity harness with Rust bridge stack
This commit is contained in:
commit
b36325d585
405 changed files with 27620 additions and 4474 deletions
|
|
@ -112,10 +112,10 @@ commands:
|
|||
node --version
|
||||
npm --version
|
||||
install_rust:
|
||||
description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself."
|
||||
description: "Install pinned rustup (1.28.2) and Rust toolchain (1.98.0) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself."
|
||||
steps:
|
||||
- run:
|
||||
name: Install Rust (rustup 1.28.2, toolchain 1.97.1)
|
||||
name: Install Rust (rustup 1.28.2, toolchain 1.98.0)
|
||||
command: |
|
||||
case "$(uname -m)" in
|
||||
x86_64)
|
||||
|
|
@ -135,7 +135,7 @@ commands:
|
|||
"https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init"
|
||||
echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c -
|
||||
chmod +x /tmp/rustup-init
|
||||
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1
|
||||
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0
|
||||
rm -f /tmp/rustup-init
|
||||
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV"
|
||||
export PATH="$HOME/.cargo/bin:$PATH"
|
||||
|
|
@ -300,7 +300,7 @@ jobs:
|
|||
if ($rustupActual -ne $rustupExpected) {
|
||||
throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual"
|
||||
}
|
||||
& $rustupInit -y --profile minimal --default-toolchain stable
|
||||
& $rustupInit -y --profile minimal --default-toolchain 1.98.0
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
exit $LASTEXITCODE
|
||||
}
|
||||
|
|
@ -1483,7 +1483,7 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
|
||||
installing_litellm_on_python_3_13:
|
||||
docker:
|
||||
|
|
@ -1507,9 +1507,9 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
|
||||
installing_litellm_on_python_legacy_migration_resolver:
|
||||
installing_litellm_on_python_v2_migration_resolver:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
|
|
@ -1536,10 +1536,10 @@ jobs:
|
|||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Run legacy migration resolver proxy smoke test
|
||||
name: Run v2 migration resolver proxy smoke test
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv \
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
|
||||
|
||||
helm_chart_testing:
|
||||
machine:
|
||||
|
|
@ -2879,8 +2879,7 @@ jobs:
|
|||
command: |
|
||||
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
|
||||
(grep -q "Database setup failed after multiple retries" docker_output.log || \
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log || \
|
||||
grep -q "Database migration cannot proceed" docker_output.log); then
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
|
||||
echo "Expected error found. Test passed."
|
||||
else
|
||||
echo "Expected error not found. Test failed."
|
||||
|
|
@ -3012,7 +3011,7 @@ workflows:
|
|||
filters: *main_branches
|
||||
- installing_litellm_on_python_3_13:
|
||||
filters: *main_branches
|
||||
- installing_litellm_on_python_legacy_migration_resolver:
|
||||
- installing_litellm_on_python_v2_migration_resolver:
|
||||
filters: *main_branches
|
||||
- helm_chart_testing:
|
||||
requires:
|
||||
|
|
|
|||
68
.github/scripts/smoke_test_native_wheel.py
vendored
Normal file
68
.github/scripts/smoke_test_native_wheel.py
vendored
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
CHILD_SCRIPT: Final = """
|
||||
from importlib.util import module_from_spec, spec_from_file_location
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
native_path = Path(sys.argv[1])
|
||||
spec = spec_from_file_location("litellm.rust_bridge._native", native_path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise RuntimeError("cannot create native extension import specification")
|
||||
module = module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
before = module.gil_stats()
|
||||
if not isinstance(before.get("releases"), int):
|
||||
raise AssertionError(f"unexpected gil_stats result: {before!r}")
|
||||
|
||||
try:
|
||||
module._panic_for_test()
|
||||
except BaseException as error:
|
||||
if type(error).__name__ != "PanicException":
|
||||
raise AssertionError(f"expected PanicException, got {type(error).__name__}") from error
|
||||
else:
|
||||
raise AssertionError("Rust panic returned without raising")
|
||||
|
||||
after = module.gil_stats()
|
||||
if not isinstance(after.get("releases"), int):
|
||||
raise AssertionError(f"native module unusable after panic: {after!r}")
|
||||
"""
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) != 2:
|
||||
sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
|
||||
return 2
|
||||
|
||||
wheel: Final = Path(sys.argv[1])
|
||||
with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive:
|
||||
native_members: Final = tuple(
|
||||
member
|
||||
for member in archive.infolist()
|
||||
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
|
||||
)
|
||||
if len(native_members) != 1:
|
||||
sys.stderr.write(f"expected one native extension, found {len(native_members)}\n")
|
||||
return 1
|
||||
|
||||
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
|
||||
native_path.write_bytes(archive.read(native_members[0]))
|
||||
result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
|
||||
|
||||
if result.returncode != 0:
|
||||
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
282
.github/scripts/verify_linux_native_wheel.py
vendored
Normal file
282
.github/scripts/verify_linux_native_wheel.py
vendored
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import zipfile
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from itertools import product
|
||||
from pathlib import Path, PurePosixPath
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import Final, Protocol
|
||||
|
||||
EXPECTED_PYTHON_TAG: Final = "cp310"
|
||||
EXPECTED_ABI_TAG: Final = "abi3"
|
||||
EXPECTED_PLATFORM_TAG: Final = "linux_x86_64"
|
||||
|
||||
|
||||
class CommandRunner(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
command: tuple[str, ...],
|
||||
*,
|
||||
check: bool,
|
||||
capture_output: bool,
|
||||
text: bool,
|
||||
) -> subprocess.CompletedProcess[str]: ...
|
||||
|
||||
|
||||
def _run_command(
|
||||
command: tuple[str, ...],
|
||||
*,
|
||||
check: bool,
|
||||
capture_output: bool,
|
||||
text: bool,
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(command, check=check, capture_output=capture_output, text=text)
|
||||
|
||||
|
||||
def _dist_info_directory(member: zipfile.ZipInfo) -> str | None:
|
||||
parts: Final = PurePosixPath(member.filename).parts
|
||||
if not parts or not parts[0].endswith(".dist-info"):
|
||||
return None
|
||||
return parts[0]
|
||||
|
||||
|
||||
def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]:
|
||||
if len(members) != 1:
|
||||
return ()
|
||||
lines: Final = archive.read(members[0]).splitlines()
|
||||
return tuple(line.removeprefix(b"Tag:").strip().decode("ascii") for line in lines if line.startswith(b"Tag:"))
|
||||
|
||||
|
||||
def _load_native_module(native_path: Path) -> ModuleType | None:
|
||||
module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path)
|
||||
if module_spec is None or module_spec.loader is None:
|
||||
return None
|
||||
try:
|
||||
native_module: Final = importlib.util.module_from_spec(module_spec)
|
||||
module_spec.loader.exec_module(native_module)
|
||||
except Exception as error: # noqa: BLE001 # native module initialization can raise arbitrary exceptions
|
||||
sys.stderr.write(f"native module load failed: {error}\n")
|
||||
return None
|
||||
return native_module
|
||||
|
||||
|
||||
def main(
|
||||
argv: Sequence[str] | None = None,
|
||||
environment: Mapping[str, str] | None = None,
|
||||
load_native_module: Callable[[Path], ModuleType | None] = _load_native_module,
|
||||
run_command: CommandRunner = _run_command,
|
||||
) -> int:
|
||||
arguments: Final = tuple(sys.argv if argv is None else argv)
|
||||
resolved_environment: Final = os.environ if environment is None else environment
|
||||
if len(arguments) != 2:
|
||||
sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL\n")
|
||||
return 2
|
||||
|
||||
wheel: Final = Path(arguments[1])
|
||||
wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3)
|
||||
if len(wheel_tags) != 4:
|
||||
sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n")
|
||||
return 1
|
||||
|
||||
wheel_identity: Final = wheel_tags[0].split("-")
|
||||
if len(wheel_identity) != 2 or wheel_identity[0] != "litellm" or not wheel_identity[1]:
|
||||
sys.stderr.write(f"unexpected wheel identity: {wheel_tags[0]}\n")
|
||||
return 1
|
||||
|
||||
expected_dist_info_directory: Final = f"{wheel_tags[0]}.dist-info"
|
||||
expected_dist_info_directories: Final = frozenset((expected_dist_info_directory,))
|
||||
python_tag: Final = wheel_tags[1]
|
||||
abi_tag: Final = wheel_tags[2]
|
||||
platform_tag: Final = wheel_tags[3]
|
||||
expanded_filename_tags: Final = frozenset(
|
||||
"-".join(tag) for tag in product(python_tag.split("."), abi_tag.split("."), platform_tag.split("."))
|
||||
)
|
||||
|
||||
with zipfile.ZipFile(wheel) as archive:
|
||||
wheel_members: Final = archive.infolist()
|
||||
dist_info_directories: Final = frozenset(
|
||||
directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None
|
||||
)
|
||||
required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL")
|
||||
dist_info_file_counts: Final = MappingProxyType(
|
||||
{
|
||||
filename: sum(
|
||||
member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members
|
||||
)
|
||||
for filename in required_dist_info_files
|
||||
}
|
||||
)
|
||||
wheel_metadata_members: Final = tuple(
|
||||
member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL"
|
||||
)
|
||||
wheel_metadata_tags: Final = _wheel_metadata_tags(archive, wheel_metadata_members)
|
||||
native_members: Final = tuple(
|
||||
member
|
||||
for member in wheel_members
|
||||
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
|
||||
)
|
||||
if len(native_members) != 1:
|
||||
sys.stderr.write(f"expected one native extension, found {len(native_members)}\n")
|
||||
return 1
|
||||
|
||||
unexpected_members: Final = tuple(
|
||||
member.filename
|
||||
for member in wheel_members
|
||||
if member.filename.endswith((".pdb", ".dwp", ".rlib", ".rmeta", "Cargo.toml", "Cargo.lock"))
|
||||
or any(part.endswith(".dSYM") for part in PurePosixPath(member.filename).parts)
|
||||
)
|
||||
native_member: Final = native_members[0]
|
||||
uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members)
|
||||
native_path: Final = wheel.parent / "native" / Path(native_member.filename).name
|
||||
native_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
native_path.write_bytes(archive.read(native_member))
|
||||
|
||||
wheel_metadata_tags_match: Final = (
|
||||
len(wheel_metadata_tags) == len(expanded_filename_tags)
|
||||
and frozenset(wheel_metadata_tags) == expanded_filename_tags
|
||||
)
|
||||
commit_sha: Final = resolved_environment.get(
|
||||
"RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown")
|
||||
)
|
||||
rustc_version: Final = run_command(
|
||||
("rustc", "--version"),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stdout.strip()
|
||||
pyproject: Final = (Path(__file__).parents[2] / "pyproject.toml").read_text()
|
||||
maturin_match: Final = re.search(r'"maturin==([^";]+)', pyproject)
|
||||
if maturin_match is None:
|
||||
sys.stderr.write("build-system does not pin an exact Maturin version\n")
|
||||
return 1
|
||||
|
||||
maturin_version: Final = maturin_match.group(1)
|
||||
native_percentage: Final = native_member.file_size / uncompressed_wheel_size * 100
|
||||
size_report: Final = "\n".join(
|
||||
(
|
||||
"## Native wheel build report",
|
||||
"",
|
||||
"| Build | Value |",
|
||||
"| --- | --- |",
|
||||
f"| Commit | `{commit_sha}` |",
|
||||
f"| Platform | `{platform_tag}` |",
|
||||
f"| Python ABI | `{python_tag}-{abi_tag}` |",
|
||||
f"| Rust compiler | `{rustc_version}` |",
|
||||
f"| Maturin | `{maturin_version}` |",
|
||||
"| Cargo profile | `release` |",
|
||||
"",
|
||||
"| Artifact | Size |",
|
||||
"| --- | ---: |",
|
||||
f"| Compressed wheel | {wheel.stat().st_size / 1_000_000:.2f} MB |",
|
||||
f"| Uncompressed wheel | {uncompressed_wheel_size / 1_000_000:.2f} MB |",
|
||||
f"| Native extension | {native_member.file_size / 1_000_000:.2f} MB |",
|
||||
f"| Native share | {native_percentage:.2f}% |",
|
||||
"",
|
||||
)
|
||||
)
|
||||
summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY")
|
||||
if summary_path is None:
|
||||
sys.stdout.write(size_report)
|
||||
else:
|
||||
Path(summary_path).write_text(size_report)
|
||||
|
||||
sections: Final = run_command(
|
||||
("readelf", "--sections", "--wide", str(native_path)),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stdout
|
||||
debug_sections: Final = tuple(section for section in (".debug_", ".zdebug_") if section in sections)
|
||||
debug_sections_absent: Final = not debug_sections
|
||||
static_symbol_table_absent: Final = ".symtab" not in sections
|
||||
|
||||
dynamic_symbols: Final = run_command(
|
||||
("readelf", "--dyn-syms", "--wide", str(native_path)),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stdout
|
||||
extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols
|
||||
native_module: Final = load_native_module(native_path)
|
||||
native_module_loads: Final = native_module is not None
|
||||
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
|
||||
native_size_limit: Final = 20_000_000
|
||||
native_size_within_limit: Final = native_member.file_size <= native_size_limit
|
||||
validations: Final = (
|
||||
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
|
||||
(f"ABI tag is {EXPECTED_ABI_TAG}", abi_tag == EXPECTED_ABI_TAG),
|
||||
(f"Platform tag is {EXPECTED_PLATFORM_TAG}", platform_tag == EXPECTED_PLATFORM_TAG),
|
||||
("Wheel dist-info directory matches the filename", dist_info_directories == expected_dist_info_directories),
|
||||
(
|
||||
"Required dist-info files are present exactly once",
|
||||
all(count == 1 for count in dist_info_file_counts.values()),
|
||||
),
|
||||
("Wheel metadata tags match the filename", wheel_metadata_tags_match),
|
||||
("Debug sections are absent", debug_sections_absent),
|
||||
("Static symbol table is absent", static_symbol_table_absent),
|
||||
("Python extension entry point is present", extension_entry_point_present),
|
||||
("Native module loads", native_module_loads),
|
||||
("Production module omits the panic test hook", panic_test_hook_absent),
|
||||
("Native extension does not exceed 20 MB", native_size_within_limit),
|
||||
("Wheel contents are valid", not unexpected_members),
|
||||
)
|
||||
|
||||
verified_report: Final = size_report + "\n".join(
|
||||
("", "| Validation | Expected | Result |", "| --- | --- | :---: |")
|
||||
+ tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations)
|
||||
+ ("",)
|
||||
)
|
||||
if summary_path is not None:
|
||||
Path(summary_path).write_text(verified_report)
|
||||
|
||||
invalid_dist_info_files: Final = any(count != 1 for count in dist_info_file_counts.values())
|
||||
validation_errors: Final = tuple(
|
||||
message
|
||||
for failed, message in (
|
||||
(bool(debug_sections), f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}"),
|
||||
(not static_symbol_table_absent, f"{native_member.filename} contains a static symbol table"),
|
||||
(not extension_entry_point_present, "native extension does not export PyInit__native"),
|
||||
(
|
||||
python_tag != EXPECTED_PYTHON_TAG,
|
||||
f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}",
|
||||
),
|
||||
(abi_tag != EXPECTED_ABI_TAG, f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}"),
|
||||
(
|
||||
platform_tag != EXPECTED_PLATFORM_TAG,
|
||||
f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}",
|
||||
),
|
||||
(
|
||||
dist_info_directories != expected_dist_info_directories,
|
||||
f"unexpected dist-info directories: expected {expected_dist_info_directory}, "
|
||||
f"found {', '.join(sorted(dist_info_directories))}",
|
||||
),
|
||||
(invalid_dist_info_files, f"required dist-info file counts are invalid: {dist_info_file_counts}"),
|
||||
(
|
||||
not invalid_dist_info_files and not wheel_metadata_tags_match,
|
||||
f"WHEEL tags do not match filename: expected {', '.join(sorted(expanded_filename_tags))}, "
|
||||
f"found {', '.join(sorted(wheel_metadata_tags))}",
|
||||
),
|
||||
(
|
||||
native_module is not None and not panic_test_hook_absent,
|
||||
"production native module exposes _panic_for_test",
|
||||
),
|
||||
(
|
||||
not native_size_within_limit,
|
||||
f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB",
|
||||
),
|
||||
(bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
|
||||
)
|
||||
if failed
|
||||
)
|
||||
sys.stderr.write("".join(f"{message}\n" for message in validation_errors))
|
||||
|
||||
return 0 if all(passed for _, passed in validations) else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
6
.github/workflows/image-scan.yml
vendored
6
.github/workflows/image-scan.yml
vendored
|
|
@ -80,7 +80,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
# Scans the whole shipped artifact: OS/apk plus every language package
|
||||
# baked into the image, including ones no lockfile declares (e.g. prisma's
|
||||
|
|
@ -124,7 +124,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
|
|
@ -185,7 +185,7 @@ jobs:
|
|||
LITELLM_COMPONENT_PORT: "4000"
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
ui-image:
|
||||
name: ui-image
|
||||
|
|
|
|||
130
.github/workflows/report-rust-release-wheel.yml
vendored
Normal file
130
.github/workflows/report-rust-release-wheel.yml
vendored
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
name: Report LiteLLM Rust release wheel
|
||||
|
||||
on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs
|
||||
workflow_run:
|
||||
workflows:
|
||||
- LiteLLM Rust
|
||||
types:
|
||||
- completed
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
report-release-wheel:
|
||||
name: report release wheel
|
||||
if: >-
|
||||
github.event.workflow_run.event == 'pull_request' &&
|
||||
github.event.workflow_run.path == '.github/workflows/test-rust.yml' &&
|
||||
github.event.workflow_run.head_repository.full_name == github.repository &&
|
||||
github.event.workflow_run.pull_requests[0].number != null
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
issues: write # PR comments use the issues API
|
||||
pull-requests: read # Current-head validation rejects stale workflow runs
|
||||
|
||||
steps:
|
||||
- name: Link release wheel report on PR
|
||||
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
|
||||
env:
|
||||
COMMENT_MARKER: "<!-- litellm-release-wheel-size -->"
|
||||
with:
|
||||
script: |
|
||||
const marker = process.env.COMMENT_MARKER;
|
||||
const workflowRun = context.payload.workflow_run;
|
||||
const allowedConclusions = new Set([
|
||||
"action_required",
|
||||
"cancelled",
|
||||
"failure",
|
||||
"neutral",
|
||||
"skipped",
|
||||
"stale",
|
||||
"startup_failure",
|
||||
"success",
|
||||
"timed_out",
|
||||
]);
|
||||
if (
|
||||
!allowedConclusions.has(workflowRun.conclusion) ||
|
||||
workflowRun.event !== "pull_request" ||
|
||||
workflowRun.path !== ".github/workflows/test-rust.yml" ||
|
||||
workflowRun.head_repository?.full_name !==
|
||||
`${context.repo.owner}/${context.repo.repo}` ||
|
||||
workflowRun.pull_requests?.length !== 1
|
||||
) {
|
||||
throw new Error("unexpected source workflow");
|
||||
}
|
||||
const pullRequest = workflowRun.pull_requests[0];
|
||||
const pullRequestNumber = pullRequest.number;
|
||||
const headSha = workflowRun.head_sha;
|
||||
const runId = workflowRun.id;
|
||||
if (
|
||||
!Number.isSafeInteger(pullRequestNumber) ||
|
||||
pullRequestNumber <= 0 ||
|
||||
!Number.isSafeInteger(runId) ||
|
||||
runId <= 0 ||
|
||||
!/^[0-9a-f]{40}$/.test(headSha) ||
|
||||
pullRequest.head?.sha !== headSha
|
||||
) {
|
||||
throw new Error("invalid source workflow metadata");
|
||||
}
|
||||
const runUrl =
|
||||
`${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` +
|
||||
`/actions/runs/${runId}`;
|
||||
const result =
|
||||
workflowRun.conclusion === "success"
|
||||
? "successfully"
|
||||
: `with \`${workflowRun.conclusion}\``;
|
||||
const body = [
|
||||
marker,
|
||||
"## LiteLLM Rust workflow",
|
||||
"",
|
||||
`Workflow completed ${result} for \`${headSha}\``,
|
||||
"",
|
||||
`[View workflow run](${runUrl})`,
|
||||
].join("\n");
|
||||
const comments = await github.paginate(github.rest.issues.listComments, {
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
issue_number: pullRequestNumber,
|
||||
per_page: 100,
|
||||
});
|
||||
const existing = comments.find(
|
||||
(comment) =>
|
||||
comment.user?.login === "github-actions[bot]" &&
|
||||
comment.body?.startsWith(marker),
|
||||
);
|
||||
const currentPullRequest = (
|
||||
await github.rest.pulls.get({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
pull_number: pullRequestNumber,
|
||||
})
|
||||
).data;
|
||||
if (
|
||||
currentPullRequest.state !== "open" ||
|
||||
currentPullRequest.head.repo?.full_name !==
|
||||
`${context.repo.owner}/${context.repo.repo}` ||
|
||||
currentPullRequest.head.sha !== headSha
|
||||
) {
|
||||
core.info("source workflow no longer matches the current pull request head");
|
||||
return;
|
||||
}
|
||||
if (existing) {
|
||||
await github.rest.issues.updateComment({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
comment_id: existing.id,
|
||||
body,
|
||||
});
|
||||
} else {
|
||||
await github.rest.issues.createComment({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
issue_number: pullRequestNumber,
|
||||
body,
|
||||
});
|
||||
}
|
||||
65
.github/workflows/test-rust.yml
vendored
65
.github/workflows/test-rust.yml
vendored
|
|
@ -4,6 +4,12 @@ on:
|
|||
push:
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".cargo/**"
|
||||
- "pyproject.toml"
|
||||
- "rust-toolchain.toml"
|
||||
- ".github/scripts/smoke_test_native_wheel.py"
|
||||
- ".github/scripts/verify_linux_native_wheel.py"
|
||||
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
pull_request:
|
||||
branches:
|
||||
|
|
@ -13,6 +19,12 @@ on:
|
|||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".cargo/**"
|
||||
- "pyproject.toml"
|
||||
- "rust-toolchain.toml"
|
||||
- ".github/scripts/smoke_test_native_wheel.py"
|
||||
- ".github/scripts/verify_linux_native_wheel.py"
|
||||
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
|
||||
permissions:
|
||||
|
|
@ -40,9 +52,7 @@ jobs:
|
|||
persist-credentials: false
|
||||
|
||||
- name: Set up Rust
|
||||
run: |
|
||||
rustup toolchain install stable --profile minimal --component clippy,rustfmt
|
||||
rustup default stable
|
||||
run: rustup toolchain install
|
||||
|
||||
- name: Cache Cargo registry and target
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
|
|
@ -51,7 +61,7 @@ jobs:
|
|||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
|
|
@ -69,3 +79,50 @@ jobs:
|
|||
|
||||
- name: Run core tests with Bedrock auth
|
||||
run: cargo test -p litellm-core --features bedrock-auth --locked
|
||||
|
||||
release-wheel:
|
||||
name: release wheel
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
permissions:
|
||||
contents: read
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up Rust
|
||||
run: rustup toolchain install
|
||||
|
||||
- name: Build release wheel
|
||||
run: uv build --wheel --out-dir dist
|
||||
|
||||
- name: Build panic contract wheel
|
||||
run: >-
|
||||
uv build --wheel --out-dir panic-dist
|
||||
--config-setting "maturin.build-args=--features panic-test,extension-module"
|
||||
|
||||
- name: Smoke-test native panic unwinding
|
||||
run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
|
||||
|
||||
- name: Verify stripped native extension
|
||||
env:
|
||||
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
|
||||
|
||||
- name: Test native route wheel
|
||||
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
@ -101,6 +103,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
# The base image only configures Chainguard's authenticated apk repo, which
|
||||
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
|
||||
# also works for anyone installing extra packages into a running container.
|
||||
# https://github.com/BerriAI/litellm/issues/33518
|
||||
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
# Stage 2 — copy source and install the project + workspace members.
|
||||
|
|
@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -117,7 +117,7 @@
|
|||
"limit": 111
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 695
|
||||
"limit": 692
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13 \
|
||||
--no-sources-package litellm-proxy-extras; \
|
||||
else \
|
||||
|
|
@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13; \
|
||||
fi
|
||||
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
|
||||
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
|
||||
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A |
|
||||
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
|
||||
|
|
@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
|
|||
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
|
||||
was not provided to the helm command line, the `masterkey` is a randomly
|
||||
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
|
||||
The key is generated once on the first install; later `helm upgrade` runs reuse the
|
||||
value already in that Secret, so upgrading never rotates the master key.
|
||||
|
||||
```bash
|
||||
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
{{- if not .Values.masterkeySecretName }}
|
||||
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
|
||||
{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }}
|
||||
{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }}
|
||||
{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: {{ include "litellm.fullname" . }}-masterkey
|
||||
name: {{ $secretName }}
|
||||
data:
|
||||
masterkey: {{ $masterkey | b64enc }}
|
||||
type: Opaque
|
||||
|
|
|
|||
|
|
@ -15,6 +15,53 @@ tests:
|
|||
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
|
||||
# but stored as base64 encoded in Kubernetes secret (requirement).
|
||||
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
|
||||
- it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhpc3Rpbmcta2V5
|
||||
- it: should let an explicit masterkey value override the one already stored in the cluster
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
masterkey: sk-explicit
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhwbGljaXQ=
|
||||
- it: should not create a secret if masterkeySecretName is set
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -7,8 +7,7 @@ import subprocess
|
|||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Optional
|
||||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
|
|
@ -51,17 +50,6 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
|
|||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_PRISMA_ATTEMPTS: Final = 4
|
||||
|
||||
_TRANSIENT_PRISMA_FAILURES: Final = MappingProxyType(
|
||||
{
|
||||
"deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)",
|
||||
"P1001": "an unreachable database server",
|
||||
"P1002": "a database server that timed out",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
|
||||
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
|
||||
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
|
||||
|
|
@ -286,23 +274,6 @@ class ProxyExtrasDBManager:
|
|||
env=prisma_env,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _transient_prisma_failure(stderr: str) -> str | None:
|
||||
"""Why a failed prisma command is worth retrying, or None.
|
||||
|
||||
v1 retried every failure, so it absorbed a database that was not up yet
|
||||
or another instance holding the migration lock. v2 fails fast, which is
|
||||
right for a broken migration and wrong for these.
|
||||
"""
|
||||
return next(
|
||||
(
|
||||
reason
|
||||
for marker, reason in _TRANSIENT_PRISMA_FAILURES.items()
|
||||
if marker in stderr
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_permission_error(error_message: str) -> bool:
|
||||
"""
|
||||
|
|
@ -684,7 +655,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
"""
|
||||
v2 migration resolver (what the proxy CLI selects by default).
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
|
||||
|
|
@ -705,46 +676,20 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(_PRISMA_ATTEMPTS):
|
||||
try:
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.info(
|
||||
"prisma db push attempt %s timed out, retrying",
|
||||
attempt + 1,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
except subprocess.CalledProcessError as e:
|
||||
stderr = e.stderr or ""
|
||||
transient = ProxyExtrasDBManager._transient_prisma_failure(
|
||||
stderr
|
||||
)
|
||||
# Re-raise as RuntimeError so proxy_cli.py's
|
||||
# `except RuntimeError` catches it and exits cleanly.
|
||||
if transient is None or attempt == _PRISMA_ATTEMPTS - 1:
|
||||
raise RuntimeError(
|
||||
f"prisma db push failed.\n\nDetail: {e}"
|
||||
f"\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
logger.info(
|
||||
"prisma db push attempt %s failed on %s, retrying. "
|
||||
"Prisma error:\n%s",
|
||||
attempt + 1,
|
||||
transient,
|
||||
stderr,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
raise RuntimeError(
|
||||
f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts."
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
) as e:
|
||||
# Re-raise as RuntimeError so proxy_cli.py's
|
||||
# `except RuntimeError` catches it and exits cleanly.
|
||||
raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
|
|
@ -754,7 +699,7 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(_PRISMA_ATTEMPTS):
|
||||
for attempt in range(4):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
|
|
@ -869,36 +814,16 @@ class ProxyExtrasDBManager:
|
|||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
transient = ProxyExtrasDBManager._transient_prisma_failure(stderr)
|
||||
if transient is None:
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
if attempt == _PRISMA_ATTEMPTS - 1:
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after "
|
||||
f"{_PRISMA_ATTEMPTS} attempts on {transient}. "
|
||||
"Check database connectivity and load."
|
||||
f"\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s failed on %s, retrying. "
|
||||
"Prisma error:\n%s",
|
||||
attempt + 1,
|
||||
transient,
|
||||
stderr,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after {_PRISMA_ATTEMPTS} "
|
||||
"attempts (retry loop exhausted by timeouts or repeated "
|
||||
"idempotent-recovery continues). Check database connectivity, "
|
||||
"load, and _prisma_migrations ledger state."
|
||||
"Database migration failed after 4 attempts (retry loop "
|
||||
"exhausted by timeouts or repeated idempotent-recovery "
|
||||
"continues). Check database connectivity, load, and "
|
||||
"_prisma_migrations ledger state."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -946,11 +871,10 @@ class ProxyExtrasDBManager:
|
|||
|
||||
Args:
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
use_v2_resolver: Run the v2 migration resolver (safer during
|
||||
use_v2_resolver: Opt into the v2 migration resolver (safer during
|
||||
rolling deploys; does not run the diff-and-force recovery
|
||||
that causes schema thrashing). Defaults to False here so
|
||||
direct callers keep the old behavior; the proxy CLI passes
|
||||
True, so the proxy's runtime default is v2.
|
||||
that causes schema thrashing). Defaults to False for
|
||||
backwards compatibility.
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
|
|
@ -968,7 +892,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
if use_v2_resolver:
|
||||
logger.info("Using v2 migration resolver")
|
||||
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
|
||||
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
|
|
|
|||
0
litellm-proxy-extras/tests/__init__.py
Normal file
0
litellm-proxy-extras/tests/__init__.py
Normal file
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
|
||||
|
||||
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
||||
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
|
||||
(default) behavior is unchanged from pre-fix.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_proxy_extras.utils import (
|
||||
ProxyExtrasDBManager,
|
||||
_max_migration_timestamp,
|
||||
_migration_timestamp,
|
||||
)
|
||||
|
||||
|
||||
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
||||
def _run(*args, **kwargs):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=returncode,
|
||||
cmd=args[0],
|
||||
stderr=stderr,
|
||||
output="",
|
||||
)
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
"Database error code: 42501\npermission denied for schema public"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="permission"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_removes_connection_limit():
|
||||
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
|
||||
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
|
||||
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
assert "connection_limit" not in stripped
|
||||
assert "pool_timeout" not in stripped
|
||||
assert "sslmode=require" in stripped
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_passthrough_no_query():
|
||||
"""URLs without query strings are returned unchanged."""
|
||||
url = "postgresql://u:p@h:5432/db"
|
||||
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
|
||||
|
||||
|
||||
def test_migration_timestamp_extracts_leading_digits():
|
||||
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
|
||||
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
|
||||
|
||||
|
||||
def test_migration_timestamp_returns_zero_on_malformed():
|
||||
assert _migration_timestamp("0_init") == 0
|
||||
assert _migration_timestamp("not_a_migration") == 0
|
||||
|
||||
|
||||
def test_max_migration_timestamp():
|
||||
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
|
||||
assert _max_migration_timestamp(names) == 20260415000000
|
||||
|
||||
|
||||
def test_max_migration_timestamp_empty_set():
|
||||
assert _max_migration_timestamp(set()) == 0
|
||||
|
||||
|
||||
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
|
||||
|
||||
This is the existing buggy behavior — we're not fixing it in v1, only
|
||||
offering v2 as opt-in. This test pins the default so that a future
|
||||
inadvertent default flip is caught.
|
||||
"""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
# Stub `prisma migrate deploy` to claim success with pending migrations
|
||||
# applied, which is the code path that triggers the legacy post-migration
|
||||
# sanity check (a call to _resolve_all_migrations).
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
def fake_run(cmd, *args, **kwargs):
|
||||
return FakeResult()
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
|
||||
def fake_resolve(*args, **kwargs):
|
||||
resolve_called["n"] += 1
|
||||
|
||||
monkeypatch.setattr("subprocess.run", fake_run)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
|
||||
|
||||
|
||||
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
|
||||
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = "db push error"
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
||||
"""_warn_if_db_ahead_of_head must never raise — it's informational.
|
||||
|
||||
Non-connection DB errors (e.g. InsufficientPrivilege from a user
|
||||
without SELECT on _prisma_migrations) must be caught, not propagated.
|
||||
"""
|
||||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class _FakeConn:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def execute(self, *a, **kw):
|
||||
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
|
||||
raise psycopg.errors.InsufficientPrivilege("permission denied")
|
||||
|
||||
def _fake_connect(*a, **kw):
|
||||
return _FakeConn()
|
||||
|
||||
monkeypatch.setattr("psycopg.connect", _fake_connect)
|
||||
|
||||
# Must not raise.
|
||||
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
|
||||
|
||||
|
||||
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""If marking a migration as applied fails inside P3009 idempotent
|
||||
recovery, the subprocess error must be re-raised as RuntimeError so
|
||||
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
|
||||
)
|
||||
|
||||
# First call: migrate deploy -> P3009 idempotent error.
|
||||
# Recovery path tries _resolve_specific_migration; that also raises.
|
||||
def _failing_resolve(*a, **kw):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd="prisma migrate resolve --applied",
|
||||
stderr="resolve failed",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
|
||||
)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Failed to mark migration .* as applied"
|
||||
):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_all_migrations",
|
||||
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
# AGENTS.md
|
||||
|
||||
litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers.
|
||||
litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
|
||||
|
||||
## Crates
|
||||
|
||||
|
|
@ -8,9 +8,10 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (
|
|||
|-------|------|
|
||||
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
|
||||
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
|
||||
|
||||
## Where a route lives
|
||||
|
||||
|
|
@ -28,7 +29,7 @@ core/src/messages/
|
|||
|
||||
Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
|
||||
|
||||
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
|
||||
Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
|
||||
|
||||
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
|
||||
|
||||
|
|
|
|||
|
|
@ -21,12 +21,13 @@ variants of it. The test for a good abstraction is that adding the next provider
|
|||
is a few declarative lines, not a new file of duplicated flow. Only diverge from
|
||||
the base when behavior is genuinely different, and say so explicitly in the PR.
|
||||
|
||||
## Crates (exactly three — see AGENTS.md)
|
||||
## Crates (see AGENTS.md)
|
||||
|
||||
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
|
||||
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
|
||||
`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not
|
||||
a route — add modules, not crates.
|
||||
`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
|
||||
holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate
|
||||
is a layer or shared foundation, not a route; add modules, not crates.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
|
|
@ -175,7 +176,7 @@ cd litellm-rust
|
|||
cargo fmt --check
|
||||
# the ai-gateway binary + server code is behind the `server` feature
|
||||
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
|
|
|
|||
174
litellm-rust/Cargo.lock
generated
174
litellm-rust/Cargo.lock
generated
|
|
@ -919,6 +919,12 @@ version = "0.3.33"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109"
|
||||
|
||||
[[package]]
|
||||
name = "futures-timer"
|
||||
version = "3.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
|
||||
|
||||
[[package]]
|
||||
name = "futures-util"
|
||||
version = "0.3.33"
|
||||
|
|
@ -972,6 +978,12 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "glob"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b"
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.3.27"
|
||||
|
|
@ -1380,6 +1392,12 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.186"
|
||||
|
|
@ -1416,6 +1434,10 @@ dependencies = [
|
|||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-types",
|
||||
"base64",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"mime",
|
||||
"rand 0.8.7",
|
||||
"reqwest",
|
||||
"serde",
|
||||
|
|
@ -1423,6 +1445,8 @@ dependencies = [
|
|||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1430,14 +1454,29 @@ name = "litellm-python-bridge"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"futures-util",
|
||||
"litellm-ai-gateway",
|
||||
"litellm-core",
|
||||
"litellm-python-interop",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"pythonize",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-python-interop"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"pyo3",
|
||||
"pythonize",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1627,6 +1666,15 @@ dependencies = [
|
|||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro-crate"
|
||||
version = "3.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
|
||||
dependencies = [
|
||||
"toml_edit",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.107"
|
||||
|
|
@ -1638,9 +1686,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "pyo3"
|
||||
version = "0.29.0"
|
||||
version = "0.29.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
|
||||
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"once_cell",
|
||||
|
|
@ -1666,18 +1714,18 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.29.0"
|
||||
version = "0.29.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
|
||||
checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
|
||||
dependencies = [
|
||||
"target-lexicon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-ffi"
|
||||
version = "0.29.0"
|
||||
version = "0.29.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
|
||||
checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"pyo3-build-config",
|
||||
|
|
@ -1685,9 +1733,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "pyo3-macros"
|
||||
version = "0.29.0"
|
||||
version = "0.29.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
|
||||
checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"pyo3-macros-backend",
|
||||
|
|
@ -1697,9 +1745,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "pyo3-macros-backend"
|
||||
version = "0.29.0"
|
||||
version = "0.29.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
|
||||
checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
|
|
@ -1899,6 +1947,12 @@ version = "0.8.11"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
|
||||
|
||||
[[package]]
|
||||
name = "relative-path"
|
||||
version = "1.9.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.12.28"
|
||||
|
|
@ -1956,6 +2010,35 @@ dependencies = [
|
|||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rstest"
|
||||
version = "0.26.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49"
|
||||
dependencies = [
|
||||
"futures-timer",
|
||||
"futures-util",
|
||||
"rstest_macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rstest_macros"
|
||||
version = "0.26.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"glob",
|
||||
"proc-macro-crate",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"regex",
|
||||
"relative-path",
|
||||
"rustc_version",
|
||||
"syn 2.0.119",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "2.1.3"
|
||||
|
|
@ -2210,6 +2293,15 @@ dependencies = [
|
|||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sharded-slab"
|
||||
version = "0.1.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "2.0.1"
|
||||
|
|
@ -2348,6 +2440,15 @@ dependencies = [
|
|||
"syn 3.0.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thread_local"
|
||||
version = "1.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.53"
|
||||
|
|
@ -2488,6 +2589,36 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_datetime"
|
||||
version = "1.1.1+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_edit"
|
||||
version = "0.25.13+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"toml_datetime",
|
||||
"toml_parser",
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_parser"
|
||||
version = "1.1.3+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
|
||||
dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
|
|
@ -2566,6 +2697,17 @@ dependencies = [
|
|||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-subscriber"
|
||||
version = "0.3.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
|
||||
dependencies = [
|
||||
"sharded-slab",
|
||||
"thread_local",
|
||||
"tracing-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "try-lock"
|
||||
version = "0.2.5"
|
||||
|
|
@ -2620,6 +2762,7 @@ dependencies = [
|
|||
"idna",
|
||||
"percent-encoding",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2903,6 +3046,15 @@ version = "0.52.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "writeable"
|
||||
version = "0.6.3"
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
members = [
|
||||
"crates/core",
|
||||
"crates/ai-gateway",
|
||||
"crates/python-interop",
|
||||
"crates/python-bridge",
|
||||
]
|
||||
resolver = "2"
|
||||
|
|
@ -13,14 +14,18 @@ license = "MIT"
|
|||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
[workspace.dependencies]
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
|
||||
litellm-python-interop = { path = "crates/python-interop" }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.29.0"
|
||||
pyo3 = "0.29.2"
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
pythonize = "0.29.0"
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
|
||||
rstest = "0.26.1"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = { version = "1.0", features = ["float_roundtrip"] }
|
||||
sha2 = "0.10"
|
||||
|
|
@ -30,6 +35,9 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
|
|||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
base64 = "0.22"
|
||||
bytes = "1"
|
||||
mime = "0.3"
|
||||
url = { version = "2", features = ["serde"] }
|
||||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
|
|
|
|||
|
|
@ -26,9 +26,10 @@ coverage and production evidence.
|
|||
|-------|------|
|
||||
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
|
||||
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
|
||||
|
||||
## Layout
|
||||
|
||||
|
|
@ -38,7 +39,8 @@ crates/
|
|||
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
|
||||
src/providers/anthropic/messages/transformation.rs
|
||||
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
|
||||
python-bridge/ PyO3 bridge for Python LiteLLM.
|
||||
python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
|
||||
python-bridge/ PyO3 API adapter for Python LiteLLM.
|
||||
```
|
||||
|
||||
The folder shape follows the Python provider tree:
|
||||
|
|
|
|||
|
|
@ -54,6 +54,6 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages`
|
|||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
|
|
|
|||
|
|
@ -6,15 +6,16 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
|
|||
|
||||
## Crates
|
||||
|
||||
`litellm-rust` is exactly three crates (a crate is a **layer**, not a route):
|
||||
`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route:
|
||||
|
||||
| Crate | Role |
|
||||
|-------|------|
|
||||
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
|
||||
| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
|
||||
|
||||
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) fn audio_transcription_provider_config(
|
||||
provider: &str,
|
||||
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
|
||||
match provider {
|
||||
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<BTreeMap<String, String>> {
|
||||
headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"audio transcription extra_headers.{key} must be a string"
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
|
||||
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
let truncated: String = body.chars().take(256).collect();
|
||||
if truncated.chars().count() == body.chars().count() {
|
||||
truncated
|
||||
} else {
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,89 +0,0 @@
|
|||
use std::time::SystemTime;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
|
||||
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::common_utils::truncate_error_body;
|
||||
use super::types::ProviderAudioTranscriptionRequest;
|
||||
use crate::client::http_client;
|
||||
|
||||
pub(crate) async fn execute_audio_transcription_provider_call(
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> CoreResult<Value> {
|
||||
let body = serde_json::to_vec(&request.body).map_err(|error| {
|
||||
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
|
||||
})?;
|
||||
let mut request_builder = http_client().post(&request.url).body(body.clone());
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|error| {
|
||||
CoreError::InvalidResponse(format!("invalid audio response JSON: {error}"))
|
||||
})?;
|
||||
Ok(request
|
||||
.config
|
||||
.transform_transcription_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
pub(crate) async fn sign_request(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
optional_params: &serde_json::Map<String, Value>,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
let env_lookup = environment_lookup;
|
||||
let auth = request
|
||||
.config
|
||||
.auth_strategy(&request.model, optional_params, &env_lookup)?;
|
||||
let body = serde_json::to_vec(&request.body).map_err(|error| {
|
||||
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
|
||||
})?;
|
||||
let mut headers = super::common_utils::string_headers(None)?;
|
||||
headers.insert("Content-Type".to_string(), "application/json".to_string());
|
||||
headers.extend(request.upstream_headers.iter().cloned());
|
||||
match auth {
|
||||
AudioTranscriptionAuth::Bearer => {}
|
||||
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
|
||||
let credentials =
|
||||
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
|
||||
.await?;
|
||||
headers.extend(sign_bedrock_post(
|
||||
&request.url,
|
||||
&body,
|
||||
&headers,
|
||||
®ion,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?);
|
||||
}
|
||||
}
|
||||
Ok(ProviderAudioTranscriptionRequest {
|
||||
upstream_headers: headers.into_iter().collect(),
|
||||
..request.clone()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn environment_lookup(key: &str) -> Option<String> {
|
||||
std::env::var(key).ok()
|
||||
}
|
||||
|
|
@ -1,15 +1,14 @@
|
|||
use litellm_core::audio_transcription::{
|
||||
AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
|
||||
prepare_audio_transcription_provider_call,
|
||||
};
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::Error;
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
|
||||
use super::handler::sign_request;
|
||||
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
use super::types::PreparedAudioTranscriptionRequest;
|
||||
use crate::integrations::custom_guardrail::{
|
||||
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
|
||||
};
|
||||
|
|
@ -26,7 +25,7 @@ pub(crate) struct AudioTranscriptionLifecycleHooks {
|
|||
request_metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
|
||||
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
|
||||
|
||||
impl AudioTranscriptionLifecycleHooks {
|
||||
|
|
@ -45,7 +44,7 @@ impl AudioTranscriptionLifecycleHooks {
|
|||
async fn run_pre_call_guardrails(
|
||||
&self,
|
||||
request: PreparedAudioTranscriptionRequest,
|
||||
) -> CoreResult<PreparedAudioTranscriptionRequest> {
|
||||
) -> Result<PreparedAudioTranscriptionRequest, Error> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
|
|
@ -63,17 +62,17 @@ impl AudioTranscriptionLifecycleHooks {
|
|||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
let Value::Object(mut data) = guardrail_request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"audio transcription pre_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let audio = data.remove("audio").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string())
|
||||
Error::InvalidRequest("audio transcription guardrail removed audio".to_string())
|
||||
})?;
|
||||
let optional_params = match data.remove("optional_params") {
|
||||
Some(Value::Object(value)) => value,
|
||||
Some(_) => {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"audio transcription optional_params must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -89,53 +88,36 @@ impl AudioTranscriptionLifecycleHooks {
|
|||
async fn prepare_provider_request(
|
||||
&self,
|
||||
request: PreparedAudioTranscriptionRequest,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
let config = audio_transcription_provider_config(&request.custom_llm_provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
|
||||
let env_lookup = super::handler::environment_lookup;
|
||||
let headers = string_headers(request.extra_headers)?;
|
||||
let url = config.complete_url(
|
||||
request.api_base.as_deref(),
|
||||
&request.model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_transcription_params(&request.optional_params);
|
||||
let body = config.transform_transcription_request(
|
||||
&request.model,
|
||||
request.audio,
|
||||
filtered_params,
|
||||
)?;
|
||||
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
|
||||
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
|
||||
if matches!(auth, AudioTranscriptionAuth::Bearer)
|
||||
&& !has_header(
|
||||
&upstream_headers
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeMap<_, _>>(),
|
||||
"authorization",
|
||||
)
|
||||
&& let Some(api_key) = request.api_key.as_deref()
|
||||
{
|
||||
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
|
||||
}
|
||||
let provider_request = ProviderAudioTranscriptionRequest {
|
||||
model: request.model,
|
||||
config,
|
||||
url,
|
||||
body: body.body,
|
||||
upstream_headers,
|
||||
timeout: request.timeout,
|
||||
};
|
||||
let provider_request = self.run_during_call_guardrails(provider_request).await?;
|
||||
sign_request(&provider_request, &request.optional_params).await
|
||||
) -> Result<ProviderAudioTranscriptionRequest, Error> {
|
||||
let PreparedAudioTranscriptionRequest {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
audio,
|
||||
api_key,
|
||||
api_base,
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout,
|
||||
..
|
||||
} = request;
|
||||
let provider_request =
|
||||
prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
|
||||
model: &model,
|
||||
audio,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: Some(&custom_llm_provider),
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout,
|
||||
})?;
|
||||
self.run_during_call_guardrails(provider_request).await
|
||||
}
|
||||
|
||||
async fn run_during_call_guardrails(
|
||||
&self,
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
) -> Result<ProviderAudioTranscriptionRequest, Error> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
|
|
@ -144,23 +126,23 @@ impl AudioTranscriptionLifecycleHooks {
|
|||
.run_during_call(
|
||||
&guardrail_context(&self.request_metadata),
|
||||
GuardrailRequest::new(json!({
|
||||
"model": request.model,
|
||||
"custom_llm_provider": "bedrock",
|
||||
"url": request.url,
|
||||
"body": request.body,
|
||||
"model": request.model(),
|
||||
"custom_llm_provider": request.custom_llm_provider(),
|
||||
"url": request.url(),
|
||||
"body": request.body(),
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
let Value::Object(mut data) = guardrail_request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"audio transcription during_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let body = data.remove("body").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio transcription guardrail removed body".to_string())
|
||||
Error::InvalidRequest("audio transcription guardrail removed body".to_string())
|
||||
})?;
|
||||
Ok(ProviderAudioTranscriptionRequest { body, ..request })
|
||||
Ok(request.with_body(body))
|
||||
}
|
||||
|
||||
fn logging_payload(
|
||||
|
|
@ -241,7 +223,7 @@ impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscri
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
error: &'a CoreError,
|
||||
error: &'a Error,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
|
|
@ -281,22 +263,21 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
|
|||
}
|
||||
}
|
||||
|
||||
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
|
||||
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
|
||||
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
|
||||
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
|
||||
}
|
||||
|
||||
fn core_error_kind(error: &CoreError) -> &'static str {
|
||||
fn core_error_kind(error: &Error) -> &'static str {
|
||||
match error {
|
||||
CoreError::Auth(_) => "AuthError",
|
||||
CoreError::InvalidProvider(_) => "InvalidProvider",
|
||||
CoreError::InvalidRequest(_) => "InvalidRequest",
|
||||
CoreError::InvalidType { .. } => "InvalidType",
|
||||
CoreError::MissingField(_) => "MissingField",
|
||||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
Error::Auth(_) => "AuthError",
|
||||
Error::InvalidProvider(_) => "InvalidProvider",
|
||||
Error::InvalidRequest(_) => "InvalidRequest",
|
||||
Error::InvalidType { .. } => "InvalidType",
|
||||
Error::MissingField(_) => "MissingField",
|
||||
Error::Http { .. } => "HttpError",
|
||||
Error::InvalidResponse(_) => "InvalidResponse",
|
||||
Error::Network(_) => "NetworkError",
|
||||
Error::Routing(_) => "RoutingError",
|
||||
Error::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,19 +1,17 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::Error;
|
||||
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
|
||||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use serde_json::Value;
|
||||
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod hooks;
|
||||
mod prepare;
|
||||
mod types;
|
||||
|
||||
pub use types::AudioTranscriptionRequest;
|
||||
|
||||
use handler::execute_audio_transcription_provider_call;
|
||||
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
|
||||
|
||||
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult<Value> {
|
||||
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
let PreparedAudioTranscriptionCall { request, hooks } =
|
||||
prepare_audio_transcription_call(request);
|
||||
CallLifecycle::default()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -46,13 +45,3 @@ impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
|
|||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ProviderAudioTranscriptionRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,8 +15,7 @@ use std::time::Duration;
|
|||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use tokio::net::TcpStream;
|
||||
|
|
@ -48,7 +47,7 @@ pub(crate) type UpstreamRx = SplitStream<UpstreamWs>;
|
|||
/// Resolve the OpenAI API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent (guard at resolution time).
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
|
|
@ -58,7 +57,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
|||
.ok()
|
||||
.filter(|key| !key.trim().is_empty())
|
||||
})
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.
|
||||
|
|
@ -70,24 +69,24 @@ pub(crate) async fn dial_upstream(
|
|||
model: &str,
|
||||
api_key: &str,
|
||||
api_base: Option<&str>,
|
||||
) -> CoreResult<UpstreamWs> {
|
||||
) -> Result<UpstreamWs, Error> {
|
||||
let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model);
|
||||
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
// GA realtime: only Authorization. The legacy OpenAI-Beta header triggers
|
||||
// beta_api_shape_disabled, so we do not send it.
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||
.map_err(|err| CoreError::Auth(err.to_string()))?,
|
||||
.map_err(|err| Error::Auth(err.to_string()))?,
|
||||
);
|
||||
|
||||
let (upstream, _response) = connect_async(request)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
Ok(upstream)
|
||||
}
|
||||
|
||||
|
|
@ -96,22 +95,22 @@ pub(crate) async fn dial_upstream(
|
|||
/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an
|
||||
/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can
|
||||
/// discard a misbehaving socket rather than warm it.
|
||||
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<RealtimeEvent> {
|
||||
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> Result<RealtimeEvent, Error> {
|
||||
loop {
|
||||
let message = upstream_rx
|
||||
.next()
|
||||
.await
|
||||
.ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))?
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.ok_or_else(|| Error::Network("upstream closed before first event".to_string()))?
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
match message {
|
||||
Message::Text(text) => {
|
||||
return serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()));
|
||||
.map_err(|err| Error::InvalidResponse(err.to_string()));
|
||||
}
|
||||
// Ignore protocol frames (ping/pong) while waiting for the first event.
|
||||
Message::Ping(_) | Message::Pong(_) => continue,
|
||||
Message::Close(_) => {
|
||||
return Err(CoreError::Network(
|
||||
return Err(Error::Network(
|
||||
"upstream closed before first event".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -139,7 +138,7 @@ pub(crate) async fn splice<In, Out>(
|
|||
mut observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
|
|
@ -154,7 +153,7 @@ where
|
|||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -175,26 +174,26 @@ where
|
|||
// inflate its own spend log. Logging observes upstream events only.
|
||||
for outbound in config.transform_realtime_request(&event, model)?.events {
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
|
||||
upstream_tx
|
||||
.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
// upstream -> client
|
||||
upstream_message = upstream_rx.next() => {
|
||||
let Some(message) = upstream_message else { break }; // upstream closed
|
||||
match message.map_err(|err| CoreError::Network(err.to_string()))? {
|
||||
match message.map_err(|err| Error::Network(err.to_string()))? {
|
||||
Message::Text(text) => {
|
||||
let event: RealtimeEvent = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
|
||||
observe(&event);
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
|
|
@ -225,7 +224,7 @@ pub async fn realtime<In, Out>(
|
|||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
|
|
@ -258,7 +257,7 @@ pub async fn realtime_warm<In, Out>(
|
|||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ use std::sync::{Arc, Mutex};
|
|||
use std::time::{Duration, Instant};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::Error;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
|
||||
use crate::io::realtime::{
|
||||
|
|
@ -438,7 +438,7 @@ impl RealtimePool {
|
|||
///
|
||||
/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends
|
||||
/// unprompted is `session.created`; we buffer exactly that and read nothing more.
|
||||
async fn warm_one(key: &UpstreamKey) -> CoreResult<WarmConnection> {
|
||||
async fn warm_one(key: &UpstreamKey) -> Result<WarmConnection, Error> {
|
||||
let upstream: UpstreamWs =
|
||||
dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?;
|
||||
let (tx, mut rx) = upstream.split();
|
||||
|
|
|
|||
|
|
@ -4,10 +4,10 @@ use std::time::Duration;
|
|||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::Error;
|
||||
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
|
||||
use litellm_core::responses::types::ResponsesWsEvent;
|
||||
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
|
@ -37,51 +37,49 @@ impl ResponsesWebSocketConnection {
|
|||
url: &str,
|
||||
headers: &HashMap<String, String>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Self> {
|
||||
) -> Result<Self, Error> {
|
||||
let mut request = url
|
||||
.into_client_request()
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
for (name, value) in headers {
|
||||
let header_name = name
|
||||
.parse::<HeaderName>()
|
||||
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
|
||||
request.headers_mut().insert(header_name, header_value);
|
||||
}
|
||||
let connect = connect_async(request);
|
||||
let result = match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
|
||||
CoreError::Network("Responses WebSocket connection timed out".to_string())
|
||||
Error::Network("Responses WebSocket connection timed out".to_string())
|
||||
})?,
|
||||
None => connect.await,
|
||||
};
|
||||
let (socket, _) = result.map_err(|error| match error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
},
|
||||
other => CoreError::Network(other.to_string()),
|
||||
other => Error::Network(other.to_string()),
|
||||
})?;
|
||||
Ok(Self {
|
||||
socket: Arc::new(Mutex::new(Some(socket))),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send_text(&self, text: String) -> CoreResult<()> {
|
||||
pub async fn send_text(&self, text: String) -> Result<(), Error> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Err(CoreError::Network(
|
||||
"Responses WebSocket is closed".to_string(),
|
||||
));
|
||||
return Err(Error::Network("Responses WebSocket is closed".to_string()));
|
||||
};
|
||||
socket
|
||||
.send(Message::Text(text))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))
|
||||
.map_err(|error| Error::Network(error.to_string()))
|
||||
}
|
||||
|
||||
pub async fn recv_text(&self) -> CoreResult<Option<String>> {
|
||||
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
|
||||
let mut socket_guard = self.socket.lock().await;
|
||||
let Some(socket) = socket_guard.as_mut() else {
|
||||
return Ok(None);
|
||||
|
|
@ -90,27 +88,27 @@ impl ResponsesWebSocketConnection {
|
|||
Some(Ok(Message::Text(text))) => Ok(Some(text)),
|
||||
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
|
||||
.map(Some)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string())),
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string())),
|
||||
Some(Ok(Message::Close(_))) | None => Ok(None),
|
||||
Some(Ok(_)) => Ok(None),
|
||||
Some(Err(error)) => Err(CoreError::Network(error.to_string())),
|
||||
Some(Err(error)) => Err(Error::Network(error.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> CoreResult<()> {
|
||||
pub async fn close(&self) -> Result<(), Error> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
if let Some(socket) = socket.as_mut() {
|
||||
socket
|
||||
.close(None)
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
}
|
||||
*socket = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
|
|
@ -120,38 +118,38 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
|||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
})
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
async fn dial_upstream(
|
||||
model: &str,
|
||||
api_key: &str,
|
||||
api_base: Option<&str>,
|
||||
) -> CoreResult<ResponsesUpstreamWs> {
|
||||
) -> Result<ResponsesUpstreamWs, Error> {
|
||||
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||
.map_err(|error| CoreError::Auth(error.to_string()))?,
|
||||
.map_err(|error| Error::Auth(error.to_string()))?,
|
||||
);
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
|
||||
connect_async(request),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?;
|
||||
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
|
||||
result
|
||||
.map(|(socket, _)| socket)
|
||||
.map_err(|error| match error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
},
|
||||
other => CoreError::Network(other.to_string()),
|
||||
other => Error::Network(other.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -166,7 +164,7 @@ impl ResponsesWebSocketStreaming {
|
|||
observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
|
|
@ -193,7 +191,7 @@ pub(crate) async fn splice<In, Out>(
|
|||
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
|
|
@ -210,18 +208,18 @@ where
|
|||
.events
|
||||
{
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
|
||||
upstream_tx.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
message = upstream_rx.next() => {
|
||||
let Some(message) = message else { break };
|
||||
match message.map_err(|error| CoreError::Network(error.to_string()))? {
|
||||
match message.map_err(|error| Error::Network(error.to_string()))? {
|
||||
Message::Text(text) => {
|
||||
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
|
||||
observe(&event);
|
||||
for outbound in OPENAI_RESPONSES_WS_CONFIG
|
||||
.transform_ws_response(&event, model)?
|
||||
|
|
@ -229,7 +227,7 @@ where
|
|||
{
|
||||
client_out.send(outbound)
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
|
|
@ -252,7 +250,7 @@ pub async fn async_responses_websocket<In, Out>(
|
|||
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
|
|
@ -267,11 +265,11 @@ where
|
|||
.events
|
||||
{
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
|
||||
upstream_tx
|
||||
.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
ResponsesWebSocketStreaming::bidirectional_forward(
|
||||
|
|
@ -296,7 +294,7 @@ pub async fn responses_ws<In, Out>(
|
|||
observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
|
|
@ -514,7 +512,7 @@ mod tests {
|
|||
)
|
||||
.await
|
||||
.expect_err("status error");
|
||||
assert!(matches!(error, CoreError::Http { status: 401, .. }));
|
||||
assert!(matches!(error, Error::Http { status: 401, .. }));
|
||||
server.await.expect("server task");
|
||||
}
|
||||
|
||||
|
|
@ -543,7 +541,7 @@ mod tests {
|
|||
)
|
||||
.await
|
||||
.expect_err("status error");
|
||||
assert!(matches!(error, CoreError::Http { status: 500, .. }));
|
||||
assert!(matches!(error, Error::Http { status: 500, .. }));
|
||||
server.await.expect("server task");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,8 +3,7 @@ use std::time::{Duration, Instant};
|
|||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use reqwest::Url;
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -56,7 +55,7 @@ fn is_azure_document_intelligence_model(model: &str) -> bool {
|
|||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
|
|
@ -65,7 +64,7 @@ pub(super) fn string_headers(
|
|||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
Error::InvalidRequest(format!(
|
||||
"OCR extra_headers.{key} must be a string, got {}",
|
||||
litellm_core::error::json_type_name(&value)
|
||||
))
|
||||
|
|
@ -80,7 +79,7 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
|||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
|
||||
fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
|
||||
let Some(object) = document.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
|
@ -138,13 +137,13 @@ fn is_blocked_ip(ip: IpAddr) -> bool {
|
|||
}
|
||||
}
|
||||
|
||||
fn blocked_url_error(url: &Url) -> CoreError {
|
||||
CoreError::InvalidRequest(format!(
|
||||
fn blocked_url_error(url: &Url) -> Error {
|
||||
Error::InvalidRequest(format!(
|
||||
"OCR document URL rejected by SSRF protection: {url}"
|
||||
))
|
||||
}
|
||||
|
||||
async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
|
||||
async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> {
|
||||
if !matches!(url.scheme(), "http" | "https") {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
|
|
@ -162,7 +161,7 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
|
|||
.ok_or_else(|| blocked_url_error(url))?;
|
||||
let addresses = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let mut saw_address = false;
|
||||
for address in addresses {
|
||||
saw_address = true;
|
||||
|
|
@ -176,25 +175,25 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url> {
|
||||
fn redirect_location(response: &reqwest::Response, url: &Url) -> Result<Url, Error> {
|
||||
let location = response
|
||||
.headers()
|
||||
.get(reqwest::header::LOCATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("OCR document redirect missing Location header".to_string())
|
||||
Error::InvalidResponse("OCR document redirect missing Location header".to_string())
|
||||
})?;
|
||||
url.join(location)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}")))
|
||||
.map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}")))
|
||||
}
|
||||
|
||||
async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> {
|
||||
async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> {
|
||||
let client = reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let mut current_url = Url::parse(url)
|
||||
.map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
|
||||
.map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
|
||||
|
||||
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
|
||||
validate_safe_fetch_url(¤t_url).await?;
|
||||
|
|
@ -202,28 +201,28 @@ async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)
|
|||
.get(current_url.clone())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
if !response.status().is_redirection() {
|
||||
return Ok((current_url, response));
|
||||
}
|
||||
current_url = redirect_location(&response, ¤t_url)?;
|
||||
}
|
||||
|
||||
Err(CoreError::InvalidRequest(
|
||||
Err(Error::InvalidRequest(
|
||||
"Too many redirects while fetching OCR document URL".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> {
|
||||
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> {
|
||||
if max_bytes == 0 {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
return Err(Error::InvalidRequest(format!(
|
||||
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
|
||||
)));
|
||||
}
|
||||
if content_length > max_bytes {
|
||||
let size_mb = content_length as f64 / (1024.0 * 1024.0);
|
||||
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
return Err(Error::InvalidRequest(format!(
|
||||
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
|
||||
)));
|
||||
}
|
||||
|
|
@ -233,7 +232,7 @@ fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Core
|
|||
async fn read_response_with_limit(
|
||||
mut response: reqwest::Response,
|
||||
url: &Url,
|
||||
) -> CoreResult<Vec<u8>> {
|
||||
) -> Result<Vec<u8>, Error> {
|
||||
let max_bytes = max_document_download_bytes();
|
||||
if let Some(content_length) = response.content_length() {
|
||||
enforce_download_size(content_length, max_bytes, url)?;
|
||||
|
|
@ -246,7 +245,7 @@ async fn read_response_with_limit(
|
|||
while let Some(chunk) = response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?
|
||||
.map_err(|err| Error::Network(err.to_string()))?
|
||||
{
|
||||
bytes_downloaded += chunk.len() as u64;
|
||||
enforce_download_size(bytes_downloaded, max_bytes, url)?;
|
||||
|
|
@ -255,7 +254,7 @@ async fn read_response_with_limit(
|
|||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
|
||||
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
|
||||
let Some((field, url)) = document_url_field(&document)? else {
|
||||
return Ok(document);
|
||||
};
|
||||
|
|
@ -267,7 +266,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes
|
|||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(CoreError::Http {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&body),
|
||||
});
|
||||
|
|
@ -290,7 +289,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes
|
|||
let mut transformed = document
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?;
|
||||
.ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?;
|
||||
transformed.insert(field.to_string(), Value::String(data_uri));
|
||||
Ok(Value::Object(transformed))
|
||||
}
|
||||
|
|
@ -316,11 +315,11 @@ fn retry_after_secs(response: &reqwest::Response) -> u64 {
|
|||
.unwrap_or(2)
|
||||
}
|
||||
|
||||
fn operation_status(response_json: &Value) -> CoreResult<&str> {
|
||||
fn operation_status(response_json: &Value) -> Result<&str, Error> {
|
||||
let status = response_json
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
.ok_or(Error::MissingField("status"))?;
|
||||
match status {
|
||||
"succeeded" => Ok("succeeded"),
|
||||
"running" | "notStarted" => Ok("running"),
|
||||
|
|
@ -330,11 +329,11 @@ fn operation_status(response_json: &Value) -> CoreResult<&str> {
|
|||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("Unknown error");
|
||||
Err(CoreError::InvalidResponse(format!(
|
||||
Err(Error::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed: {message}"
|
||||
)))
|
||||
}
|
||||
other => Err(CoreError::InvalidResponse(format!(
|
||||
other => Err(Error::InvalidResponse(format!(
|
||||
"Unknown operation status: {other}"
|
||||
))),
|
||||
}
|
||||
|
|
@ -345,9 +344,9 @@ pub(super) async fn poll_document_intelligence(
|
|||
original_url: &str,
|
||||
headers: &[(String, String)],
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
) -> Result<Value, Error> {
|
||||
if !same_origin(operation_url, original_url) {
|
||||
return Err(CoreError::InvalidResponse(
|
||||
return Err(Error::InvalidResponse(
|
||||
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -358,7 +357,7 @@ pub(super) async fn poll_document_intelligence(
|
|||
));
|
||||
loop {
|
||||
if start.elapsed() > timeout {
|
||||
return Err(CoreError::Network(format!(
|
||||
return Err(Error::Network(format!(
|
||||
"Azure Document Intelligence operation polling timed out after {} seconds",
|
||||
timeout.as_secs()
|
||||
)));
|
||||
|
|
@ -373,21 +372,21 @@ pub(super) async fn poll_document_intelligence(
|
|||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let retry_after = retry_after_secs(&response);
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
})?;
|
||||
if operation_status(&response_json)? == "succeeded" {
|
||||
return Ok(response_json);
|
||||
|
|
@ -426,7 +425,7 @@ mod tests {
|
|||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CoreError::InvalidRequest(message)
|
||||
Error::InvalidRequest(message)
|
||||
if message.contains("SSRF protection")
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -7,7 +6,7 @@ use super::common_utils::{poll_document_intelligence, truncate_error_body};
|
|||
use super::types::ProviderOcrRequest;
|
||||
use crate::client::http_client;
|
||||
|
||||
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
|
||||
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result<Value, Error> {
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
|
|
@ -19,7 +18,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
|
|||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
|
|
@ -31,7 +30,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
|
|||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse(
|
||||
Error::InvalidResponse(
|
||||
"Azure Document Intelligence returned 202 but no Operation-Location header found"
|
||||
.to_string(),
|
||||
)
|
||||
|
|
@ -52,17 +51,17 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
|
|||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(request
|
||||
.config
|
||||
|
|
|
|||
|
|
@ -1,11 +1,9 @@
|
|||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::ocr::transformation::OcrAuthStrategy;
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use super::common_utils::{
|
||||
convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
|
||||
|
|
@ -27,7 +25,7 @@ pub(crate) struct OcrLifecycleHooks {
|
|||
request_metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
|
||||
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
|
||||
|
||||
impl OcrLifecycleHooks {
|
||||
|
|
@ -46,7 +44,7 @@ impl OcrLifecycleHooks {
|
|||
async fn run_pre_call_guardrails(
|
||||
&self,
|
||||
request: PreparedOcrRequest,
|
||||
) -> CoreResult<PreparedOcrRequest> {
|
||||
) -> Result<PreparedOcrRequest, Error> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
|
|
@ -74,9 +72,9 @@ impl OcrLifecycleHooks {
|
|||
async fn prepare_provider_request(
|
||||
&self,
|
||||
request: PreparedOcrRequest,
|
||||
) -> CoreResult<ProviderOcrRequest> {
|
||||
) -> Result<ProviderOcrRequest, Error> {
|
||||
let config = ocr_provider_config(&request.custom_llm_provider, &request.model)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
|
||||
.ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let headers = string_headers(request.extra_headers)?;
|
||||
let auth_strategy = config.auth_strategy();
|
||||
|
|
@ -89,7 +87,14 @@ impl OcrLifecycleHooks {
|
|||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_ocr_params(&request.optional_params);
|
||||
let supported_params = config.get_supported_ocr_params();
|
||||
let non_default_params = request
|
||||
.optional_params
|
||||
.iter()
|
||||
.filter(|(param, _)| supported_params.contains(¶m.as_str()))
|
||||
.map(|(param, value)| (param.clone(), value.clone()))
|
||||
.collect();
|
||||
let filtered_params = config.map_ocr_params(&non_default_params);
|
||||
let model = request.model.clone();
|
||||
let custom_llm_provider = request.custom_llm_provider.clone();
|
||||
let document = if config.requires_data_uri_document() {
|
||||
|
|
@ -120,7 +125,7 @@ impl OcrLifecycleHooks {
|
|||
custom_llm_provider: &str,
|
||||
url: &str,
|
||||
body: Value,
|
||||
) -> CoreResult<Value> {
|
||||
) -> Result<Value, Error> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(body);
|
||||
}
|
||||
|
|
@ -217,7 +222,7 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
error: &'a CoreError,
|
||||
error: &'a Error,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
|
|
@ -278,19 +283,19 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
|
|||
|
||||
fn parse_ocr_pre_call_guardrail_request(
|
||||
request: GuardrailRequest,
|
||||
) -> CoreResult<(Value, Map<String, Value>)> {
|
||||
) -> Result<(Value, Map<String, Value>), Error> {
|
||||
let Value::Object(mut data) = request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"OCR pre_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let document = data.remove("document").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string())
|
||||
Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
|
||||
})?;
|
||||
let optional_params = match data.remove("optional_params") {
|
||||
Some(Value::Object(params)) => params,
|
||||
Some(_) => {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"OCR pre_call guardrail optional_params must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -299,33 +304,31 @@ fn parse_ocr_pre_call_guardrail_request(
|
|||
Ok((document, optional_params))
|
||||
}
|
||||
|
||||
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult<Value> {
|
||||
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
|
||||
let Value::Object(mut data) = request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"OCR during_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
data.remove("body").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string())
|
||||
})
|
||||
data.remove("body")
|
||||
.ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
|
||||
}
|
||||
|
||||
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
|
||||
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
|
||||
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
|
||||
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
|
||||
}
|
||||
|
||||
fn core_error_kind(error: &CoreError) -> &'static str {
|
||||
fn core_error_kind(error: &Error) -> &'static str {
|
||||
match error {
|
||||
CoreError::Auth(_) => "AuthError",
|
||||
CoreError::InvalidProvider(_) => "InvalidProvider",
|
||||
CoreError::InvalidRequest(_) => "InvalidRequest",
|
||||
CoreError::InvalidType { .. } => "InvalidType",
|
||||
CoreError::MissingField(_) => "MissingField",
|
||||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
Error::Auth(_) => "AuthError",
|
||||
Error::InvalidProvider(_) => "InvalidProvider",
|
||||
Error::InvalidRequest(_) => "InvalidRequest",
|
||||
Error::InvalidType { .. } => "InvalidType",
|
||||
Error::MissingField(_) => "MissingField",
|
||||
Error::Http { .. } => "HttpError",
|
||||
Error::InvalidResponse(_) => "InvalidResponse",
|
||||
Error::Network(_) => "NetworkError",
|
||||
Error::Routing(_) => "RoutingError",
|
||||
Error::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::Error;
|
||||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -13,7 +13,7 @@ pub use types::OcrRequest;
|
|||
use handler::execute_ocr_provider_call;
|
||||
use prepare::{PreparedOcrCall, prepare_ocr_call};
|
||||
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
|
||||
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
|
||||
CallLifecycle::default()
|
||||
.run_request(request, &hooks, execute_ocr_provider_call)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use serde_json::{Map, Value, json};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
|
@ -18,6 +18,14 @@ use crate::integrations::custom_logger::{
|
|||
};
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
type LifecycleEvents = Arc<Mutex<Vec<&'static str>>>;
|
||||
|
||||
fn record_lifecycle_event(events: Option<&LifecycleEvents>, event: &'static str) {
|
||||
if let Some(events) = events {
|
||||
events.lock().unwrap().push(event);
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_http_headers(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
|
|
@ -80,9 +88,17 @@ struct RecordedLogEvent {
|
|||
#[derive(Default)]
|
||||
struct RecordingOcrLogger {
|
||||
events: Mutex<Vec<RecordedLogEvent>>,
|
||||
lifecycle_events: Option<LifecycleEvents>,
|
||||
}
|
||||
|
||||
impl RecordingOcrLogger {
|
||||
fn with_lifecycle_events(lifecycle_events: LifecycleEvents) -> Self {
|
||||
Self {
|
||||
events: Mutex::new(Vec::new()),
|
||||
lifecycle_events: Some(lifecycle_events),
|
||||
}
|
||||
}
|
||||
|
||||
fn events(&self) -> Vec<RecordedLogEvent> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
|
|
@ -96,6 +112,7 @@ impl CustomLogger for RecordingOcrLogger {
|
|||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_log_success_event");
|
||||
self.events.lock().unwrap().push(RecordedLogEvent {
|
||||
hook: "async_log_success_event",
|
||||
model: model_call_details.model.clone(),
|
||||
|
|
@ -115,6 +132,7 @@ impl CustomLogger for RecordingOcrLogger {
|
|||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_log_failure_event");
|
||||
self.events.lock().unwrap().push(RecordedLogEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: model_call_details.model.clone(),
|
||||
|
|
@ -135,22 +153,41 @@ struct RecordingOcrGuardrail {
|
|||
hooks: Vec<GuardrailEventHook>,
|
||||
events: Mutex<Vec<&'static str>>,
|
||||
block_pre_call: bool,
|
||||
block_during_call: bool,
|
||||
lifecycle_events: Option<LifecycleEvents>,
|
||||
}
|
||||
|
||||
impl RecordingOcrGuardrail {
|
||||
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
|
||||
fn with_lifecycle_events(
|
||||
hooks: Vec<GuardrailEventHook>,
|
||||
lifecycle_events: LifecycleEvents,
|
||||
) -> Self {
|
||||
Self {
|
||||
hooks,
|
||||
events: Mutex::new(Vec::new()),
|
||||
block_pre_call: false,
|
||||
block_during_call: false,
|
||||
lifecycle_events: Some(lifecycle_events),
|
||||
}
|
||||
}
|
||||
|
||||
fn blocking_pre_call() -> Self {
|
||||
fn blocking_pre_call(lifecycle_events: LifecycleEvents) -> Self {
|
||||
Self {
|
||||
hooks: vec![GuardrailEventHook::PreCall],
|
||||
events: Mutex::new(Vec::new()),
|
||||
block_pre_call: true,
|
||||
block_during_call: false,
|
||||
lifecycle_events: Some(lifecycle_events),
|
||||
}
|
||||
}
|
||||
|
||||
fn blocking_during_call(lifecycle_events: LifecycleEvents) -> Self {
|
||||
Self {
|
||||
hooks: vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall],
|
||||
events: Mutex::new(Vec::new()),
|
||||
block_pre_call: false,
|
||||
block_during_call: true,
|
||||
lifecycle_events: Some(lifecycle_events),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -174,6 +211,7 @@ impl CustomGuardrail for RecordingOcrGuardrail {
|
|||
mut request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_pre_call_hook");
|
||||
self.events.lock().unwrap().push("async_pre_call_hook");
|
||||
if self.block_pre_call {
|
||||
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
|
||||
|
|
@ -191,7 +229,13 @@ impl CustomGuardrail for RecordingOcrGuardrail {
|
|||
mut request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_moderation_hook");
|
||||
self.events.lock().unwrap().push("async_moderation_hook");
|
||||
if self.block_during_call {
|
||||
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
|
||||
"blocked during provider preparation",
|
||||
)));
|
||||
}
|
||||
request.data["body"]["guarded_during"] = json!(true);
|
||||
Ok(GuardrailDecision::Mask(request))
|
||||
})
|
||||
|
|
@ -242,7 +286,7 @@ fn ocr_dispatch_supports_migrated_providers() {
|
|||
assert!(
|
||||
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
|
||||
.expect("vertex deepseek config resolves")
|
||||
.supported_ocr_params()
|
||||
.get_supported_ocr_params()
|
||||
.contains(&"temperature")
|
||||
);
|
||||
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
|
||||
|
|
@ -281,14 +325,17 @@ fn auth_header_detection_is_case_insensitive() {
|
|||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
|
||||
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let provider_events = lifecycle_events.clone();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let request = read_http_request(&mut socket).await;
|
||||
record_lifecycle_event(Some(&provider_events), "provider_request_received");
|
||||
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
|
|
@ -302,11 +349,13 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
|
|||
request
|
||||
});
|
||||
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
|
||||
GuardrailEventHook::PreCall,
|
||||
GuardrailEventHook::DuringCall,
|
||||
]));
|
||||
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::with_lifecycle_events(
|
||||
vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall],
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
let response = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
|
|
@ -346,22 +395,81 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
|
|||
error_kind: None,
|
||||
}]
|
||||
);
|
||||
assert_eq!(
|
||||
lifecycle_events.lock().unwrap().as_slice(),
|
||||
[
|
||||
"async_pre_call_hook",
|
||||
"async_moderation_hook",
|
||||
"provider_request_received",
|
||||
"async_log_success_event",
|
||||
]
|
||||
);
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
|
||||
assert!(request.contains(r#""guarded_during":true"#), "{request}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_during_call_block_skips_provider_and_runs_failure_callback() {
|
||||
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
|
||||
let error = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("mistral"),
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_millis(100)),
|
||||
callbacks: vec![logger],
|
||||
guardrails: vec![guardrail],
|
||||
request_metadata: RequestMetadata::default(),
|
||||
litellm_call_id: Some("ocr-call-during-block"),
|
||||
})
|
||||
.await
|
||||
.expect_err("during-call guardrail blocks request");
|
||||
|
||||
assert!(matches!(error, Error::InvalidRequest(_)));
|
||||
assert_eq!(
|
||||
lifecycle_events.lock().unwrap().as_slice(),
|
||||
[
|
||||
"async_pre_call_hook",
|
||||
"async_moderation_hook",
|
||||
"async_log_failure_event",
|
||||
]
|
||||
);
|
||||
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
|
||||
assert!(accepted.is_err(), "provider socket should not be touched");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
|
||||
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let provider_events = lifecycle_events.clone();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let _request = read_http_request(&mut socket).await;
|
||||
record_lifecycle_event(Some(&provider_events), "provider_request_received");
|
||||
let response_body = "provider failed";
|
||||
let response = format!(
|
||||
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
|
|
@ -374,7 +482,9 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
|
|||
.expect("writes response");
|
||||
});
|
||||
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
let err = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
|
|
@ -395,7 +505,7 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
|
|||
.await
|
||||
.expect_err("provider error propagates");
|
||||
|
||||
assert!(matches!(err, CoreError::Http { status: 500, .. }));
|
||||
assert!(matches!(err, Error::Http { status: 500, .. }));
|
||||
server.await.expect("server task completes");
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
|
|
@ -408,16 +518,25 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
|
|||
error_kind: Some("HttpError".to_string()),
|
||||
}]
|
||||
);
|
||||
assert_eq!(
|
||||
lifecycle_events.lock().unwrap().as_slice(),
|
||||
["provider_request_received", "async_log_failure_event"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
|
||||
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
|
||||
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call(
|
||||
lifecycle_events.clone(),
|
||||
));
|
||||
|
||||
let err = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
|
|
@ -439,7 +558,7 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
|
|||
.await
|
||||
.expect_err("guardrail blocks request");
|
||||
|
||||
assert!(matches!(err, CoreError::InvalidRequest(_)));
|
||||
assert!(matches!(err, Error::InvalidRequest(_)));
|
||||
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
|
|
@ -452,6 +571,10 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
|
|||
error_kind: Some("InvalidRequest".to_string()),
|
||||
}]
|
||||
);
|
||||
assert_eq!(
|
||||
lifecycle_events.lock().unwrap().as_slice(),
|
||||
["async_pre_call_hook", "async_log_failure_event"]
|
||||
);
|
||||
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
|
||||
assert!(accepted.is_err(), "provider socket should not be touched");
|
||||
}
|
||||
|
|
@ -607,7 +730,7 @@ fn string_headers_rejects_non_string_values() {
|
|||
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
Error::InvalidRequest(
|
||||
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
|
|
|
|||
|
|
@ -6,33 +6,31 @@
|
|||
//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python.
|
||||
//!
|
||||
//! Compiled only under the `python-config` feature.
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::router::{Deployment, Router};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::gil;
|
||||
|
||||
/// Load the router's `model_list` from `config_path` via the Python reader.
|
||||
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
|
||||
pub fn load_router_from_config(config_path: &str) -> Result<Router, Error> {
|
||||
gil::record_acquisition();
|
||||
Python::attach(|py| {
|
||||
let model_list = py
|
||||
.import("litellm.proxy.read_model_list")
|
||||
.and_then(|module| module.getattr("read_model_list"))
|
||||
.and_then(|reader| reader.call1((config_path,)))
|
||||
.map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?;
|
||||
.map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?;
|
||||
|
||||
let model_list_json: String = py
|
||||
.import("json")
|
||||
.and_then(|json| json.getattr("dumps"))
|
||||
.and_then(|dumps| dumps.call1((model_list,)))
|
||||
.and_then(|encoded| encoded.extract())
|
||||
.map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?;
|
||||
.map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?;
|
||||
|
||||
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
|
||||
.map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?;
|
||||
.map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?;
|
||||
|
||||
Ok(Router::new(deployments))
|
||||
})
|
||||
|
|
|
|||
|
|
@ -4,21 +4,36 @@ mod service;
|
|||
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::extract::{Json, State};
|
||||
use axum::extract::{Json, Request, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use axum::middleware::{self, Next};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use litellm_core::CoreError;
|
||||
use litellm_core::Error;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH};
|
||||
use crate::state::AppState;
|
||||
|
||||
const CORE_ENGINE_HEADER: &str = "x-litellm-core";
|
||||
const RUST_CORE_ENGINE: &str = "rust";
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
|
||||
Router::new()
|
||||
.route(MESSAGES_ROUTE_PATH, post(handle))
|
||||
.route_layer(middleware::from_fn(core_engine_header))
|
||||
}
|
||||
|
||||
async fn core_engine_header(request: Request, next: Next) -> Response {
|
||||
let mut response = next.run(request).await;
|
||||
response.headers_mut().insert(
|
||||
CORE_ENGINE_HEADER,
|
||||
HeaderValue::from_static(RUST_CORE_ENGINE),
|
||||
);
|
||||
response
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
|
|
@ -46,7 +61,7 @@ fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRout
|
|||
let mut response = Response::builder()
|
||||
.status(
|
||||
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
|
||||
MessagesRouteError(CoreError::InvalidResponse(format!(
|
||||
MessagesRouteError(Error::InvalidResponse(format!(
|
||||
"invalid upstream response status: {error}"
|
||||
)))
|
||||
})?,
|
||||
|
|
@ -58,13 +73,13 @@ fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRout
|
|||
response
|
||||
.body(Body::from_stream(upstream.bytes_stream()))
|
||||
.map_err(|error| {
|
||||
MessagesRouteError(CoreError::InvalidResponse(format!(
|
||||
MessagesRouteError(Error::InvalidResponse(format!(
|
||||
"failed to build streaming response: {error}"
|
||||
)))
|
||||
})
|
||||
}
|
||||
|
||||
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
|
||||
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, Error> {
|
||||
let forwarded = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
|
|
@ -74,19 +89,19 @@ fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>,
|
|||
})
|
||||
.map(|(name, value)| {
|
||||
let value = value.to_str().map_err(|_| {
|
||||
CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str()))
|
||||
Error::InvalidRequest(format!("invalid value for header {}", name.as_str()))
|
||||
})?;
|
||||
Ok((name.to_string(), Value::String(value.to_string())))
|
||||
})
|
||||
.collect::<Result<Map<_, _>, CoreError>>()?;
|
||||
.collect::<Result<Map<_, _>, Error>>()?;
|
||||
Ok((!forwarded.is_empty()).then_some(forwarded))
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MessagesRouteError(CoreError);
|
||||
struct MessagesRouteError(Error);
|
||||
|
||||
impl From<CoreError> for MessagesRouteError {
|
||||
fn from(error: CoreError) -> Self {
|
||||
impl From<Error> for MessagesRouteError {
|
||||
fn from(error: Error) -> Self {
|
||||
Self(error)
|
||||
}
|
||||
}
|
||||
|
|
@ -94,28 +109,27 @@ impl From<CoreError> for MessagesRouteError {
|
|||
impl IntoResponse for MessagesRouteError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, message) = match self.0 {
|
||||
CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
|
||||
CoreError::InvalidProvider(_) | CoreError::Routing(_) => (
|
||||
Error::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
|
||||
Error::InvalidProvider(_) | Error::Routing(_) => (
|
||||
StatusCode::NOT_FOUND,
|
||||
"no messages deployment is configured for this model".to_string(),
|
||||
),
|
||||
CoreError::Auth(_) => (
|
||||
Error::Auth(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider authentication failed".to_string(),
|
||||
),
|
||||
CoreError::Http { .. }
|
||||
| CoreError::Network(_)
|
||||
| CoreError::Connect(_)
|
||||
| CoreError::InvalidResponse(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_) => (
|
||||
Error::Http { .. }
|
||||
| Error::Network(_)
|
||||
| Error::InvalidResponse(_)
|
||||
| Error::InvalidType { .. }
|
||||
| Error::MissingField(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider request failed".to_string(),
|
||||
),
|
||||
// The gateway has no Python implementation to decline to, so a
|
||||
// request the core cannot serve is reported to the caller. The
|
||||
// reason is a fixed internal string, never provider content.
|
||||
CoreError::Unsupported(reason) => (
|
||||
Error::Unsupported(reason) => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("messages request is not supported: {reason}"),
|
||||
),
|
||||
|
|
@ -143,6 +157,7 @@ mod tests {
|
|||
use tower::ServiceExt;
|
||||
|
||||
use super::super::app;
|
||||
use super::{CORE_ENGINE_HEADER, RUST_CORE_ENGINE};
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use crate::state::AppState;
|
||||
|
||||
|
|
@ -290,6 +305,7 @@ mod tests {
|
|||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()[CORE_ENGINE_HEADER], RUST_CORE_ENGINE);
|
||||
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body reads");
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_core::Error;
|
||||
use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use litellm_core::messages::types::MessagesRequest;
|
||||
use litellm_core::messages::{messages, messages_stream};
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(crate) enum MessagesResponse {
|
||||
|
|
@ -16,16 +16,16 @@ pub async fn run(
|
|||
router: &Arc<Router>,
|
||||
body: Value,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<MessagesResponse> {
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|model| !model.is_empty())
|
||||
.ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?;
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
.ok_or_else(|| Error::InvalidRequest("messages body requires a model".to_string()))?;
|
||||
let deployment = router
|
||||
.get_available_deployment(model)
|
||||
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
|
||||
let provider_model = deployment.litellm_params.model.as_str();
|
||||
let upstream_model = provider_model
|
||||
.split_once('/')
|
||||
|
|
@ -37,7 +37,7 @@ pub async fn run(
|
|||
};
|
||||
let mut body = body;
|
||||
body.as_object_mut()
|
||||
.ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))?
|
||||
.ok_or_else(|| Error::InvalidRequest("messages body must be an object".to_string()))?
|
||||
.insert(
|
||||
"model".to_string(),
|
||||
Value::String(upstream_model.to_string()),
|
||||
|
|
@ -60,6 +60,6 @@ pub async fn run(
|
|||
serde_json::to_value(response)
|
||||
.map(MessagesResponse::Json)
|
||||
.map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
|
||||
Error::InvalidResponse(format!("failed to serialize messages response: {err}"))
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,8 +11,7 @@ use std::time::Duration;
|
|||
|
||||
use crate::io::realtime_pool::{RealtimePool, upstream_key};
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router;
|
||||
|
||||
|
|
@ -29,15 +28,15 @@ pub async fn run<In, Out>(
|
|||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let deployment = router
|
||||
.get_available_deployment(model)
|
||||
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
|
||||
let params = &deployment.litellm_params;
|
||||
// Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model.
|
||||
let provider_model = params
|
||||
|
|
|
|||
|
|
@ -2,13 +2,13 @@ use std::sync::Arc;
|
|||
use std::time::Duration;
|
||||
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::Error;
|
||||
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
|
||||
use litellm_core::responses::instrumentation::{
|
||||
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
|
||||
ResponsesWsMetadata,
|
||||
};
|
||||
use litellm_core::responses::types::ResponsesWsEvent;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
|
||||
|
|
@ -26,22 +26,22 @@ pub async fn run<In, Out>(
|
|||
metadata: RequestMetadata,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
) -> Result<(), Error>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let deployment = router
|
||||
.get_available_deployment(model)
|
||||
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
|
||||
let params = &deployment.litellm_params;
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
if params.model.contains('/') && !params.model.starts_with("openai/") {
|
||||
return Err(CoreError::InvalidProvider(
|
||||
return Err(Error::InvalidProvider(
|
||||
"Responses WebSocket route supports OpenAI deployments only".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,11 +6,17 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
tracing.workspace = true
|
||||
base64.workspace = true
|
||||
futures-util.workspace = true
|
||||
bytes.workspace = true
|
||||
mime.workspace = true
|
||||
rand.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
url.workspace = true
|
||||
sha2.workspace = true
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
|
||||
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
|
||||
|
|
|
|||
14
litellm-rust/crates/core/src/audio_transcription/client.rs
Normal file
14
litellm-rust/crates/core/src/audio_transcription/client.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
|
||||
|
||||
pub(super) fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
})
|
||||
}
|
||||
91
litellm-rust/crates/core/src/audio_transcription/handler.rs
Normal file
91
litellm-rust/crates/core/src/audio_transcription/handler.rs
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::Error;
|
||||
use crate::http_utils::truncate_error_body;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::types::ProviderAudioTranscriptionRequest;
|
||||
|
||||
pub async fn execute_audio_transcription_provider_call(
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> Result<Value, Error> {
|
||||
let body = serde_json::to_vec(&request.body)
|
||||
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
|
||||
let headers = signed_headers(&request, &body).await?;
|
||||
let mut request_builder = http_client().post(&request.url).body(body);
|
||||
for (key, value) in headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json = serde_json::from_str(&text)
|
||||
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
|
||||
Ok(request
|
||||
.config
|
||||
.transform_transcription_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
async fn signed_headers(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
body: &[u8],
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use crate::providers::bedrock::audio_transcription::aws_auth_config;
|
||||
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
|
||||
|
||||
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
|
||||
return Ok(request.upstream_headers.clone());
|
||||
};
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let credentials = resolve_credentials(
|
||||
aws_auth_config(&request.optional_params, &env_lookup),
|
||||
&env_lookup,
|
||||
)
|
||||
.await?;
|
||||
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
|
||||
let signature = sign_bedrock_post(
|
||||
&request.url,
|
||||
body,
|
||||
&unsigned,
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
Ok(unsigned.into_iter().chain(signature).collect())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "bedrock-auth"))]
|
||||
async fn signed_headers(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
_body: &[u8],
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
|
||||
match request.auth {
|
||||
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
|
||||
"AWS SigV4 requires the bedrock-auth feature",
|
||||
)),
|
||||
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,20 @@
|
|||
use crate::Error;
|
||||
mod client;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
pub use handler::execute_audio_transcription_provider_call;
|
||||
pub use prepare::prepare_audio_transcription_provider_call;
|
||||
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
|
||||
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
|
|||
72
litellm-rust/crates/core/src/audio_transcription/prepare.rs
Normal file
72
litellm-rust/crates/core/src/audio_transcription/prepare.rs
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
use crate::error::Error;
|
||||
use crate::http_utils::{has_header, string_headers};
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
|
||||
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
|
||||
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
if provider == "bedrock" {
|
||||
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
|
||||
}
|
||||
let _ = provider;
|
||||
None
|
||||
}
|
||||
|
||||
pub fn prepare_audio_transcription_provider_call(
|
||||
request: AudioTranscriptionRequest<'_>,
|
||||
) -> Result<ProviderAudioTranscriptionRequest, Error> {
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.or_else(|| {
|
||||
request
|
||||
.custom_llm_provider
|
||||
.map(|provider| CustomLlmProvider {
|
||||
model: request.model,
|
||||
custom_llm_provider: provider,
|
||||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for audio transcription request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let model = provider_info.model.to_string();
|
||||
let config = provider_config(provider_info.custom_llm_provider)
|
||||
.ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let mut headers = string_headers("audio transcription", request.extra_headers)?;
|
||||
let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?;
|
||||
if matches!(auth, AudioTranscriptionAuth::Bearer)
|
||||
&& !has_header(&headers, "authorization")
|
||||
&& let Some(api_key) = request.api_key
|
||||
{
|
||||
headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
|
||||
}
|
||||
if !has_header(&headers, "content-type") {
|
||||
headers.push(("Content-Type".to_string(), "application/json".to_string()));
|
||||
}
|
||||
let url = config.complete_url(
|
||||
request.api_base,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_transcription_params(&request.optional_params);
|
||||
let transformed =
|
||||
config.transform_transcription_request(&model, request.audio, filtered_params)?;
|
||||
Ok(ProviderAudioTranscriptionRequest {
|
||||
model,
|
||||
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
|
||||
config,
|
||||
url,
|
||||
body: transformed.body,
|
||||
upstream_headers: headers,
|
||||
auth,
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
optional_params: request.optional_params,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
50
litellm-rust/crates/core/src/audio_transcription/tests.rs
Normal file
50
litellm-rust/crates/core/src/audio_transcription/tests.rs
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::thread;
|
||||
|
||||
use serde_json::{Map, json};
|
||||
|
||||
use super::audio_transcription;
|
||||
use super::types::AudioTranscriptionRequest;
|
||||
|
||||
#[tokio::test]
|
||||
async fn bedrock_request_is_signed_and_contains_audio() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
|
||||
let address = listener.local_addr().expect("address");
|
||||
let server = thread::spawn(move || {
|
||||
let (mut stream, _) = listener.accept().expect("connection");
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 16_384];
|
||||
let count = stream.read(&mut buffer).expect("request");
|
||||
request.extend_from_slice(&buffer[..count]);
|
||||
let request = String::from_utf8_lossy(&request);
|
||||
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
|
||||
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
|
||||
assert!(request.contains("x-amz-date:"));
|
||||
assert!(request.contains("\"bytes\":\"AQI=\""));
|
||||
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
|
||||
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
|
||||
stream.write_all(response).expect("response");
|
||||
});
|
||||
|
||||
let optional_params = Map::from_iter([
|
||||
("aws_access_key_id".to_string(), json!("access-key")),
|
||||
("aws_secret_access_key".to_string(), json!("secret-key")),
|
||||
("aws_region_name".to_string(), json!("us-east-1")),
|
||||
]);
|
||||
let api_base = format!("http://{address}");
|
||||
let response = audio_transcription(AudioTranscriptionRequest {
|
||||
model: "mistral.voxtral-mini-3b-2507",
|
||||
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
|
||||
api_key: None,
|
||||
api_base: Some(&api_base),
|
||||
custom_llm_provider: Some("bedrock"),
|
||||
extra_headers: None,
|
||||
optional_params,
|
||||
timeout: None,
|
||||
})
|
||||
.await
|
||||
.expect("transcription");
|
||||
assert_eq!(response, json!({"text": "hello"}));
|
||||
server.join().expect("server");
|
||||
}
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
use crate::Error;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::CoreResult;
|
||||
|
||||
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
|
|
@ -32,13 +31,13 @@ pub trait AudioTranscriptionProviderConfig: Sync {
|
|||
model: &str,
|
||||
audio: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<AudioTranscriptionRequestData>;
|
||||
) -> Result<AudioTranscriptionRequestData, Error>;
|
||||
|
||||
fn transform_transcription_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<AudioTranscriptionResponseData>;
|
||||
) -> Result<AudioTranscriptionResponseData, Error>;
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
|
|
@ -46,12 +45,12 @@ pub trait AudioTranscriptionProviderConfig: Sync {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn auth_strategy(
|
||||
&self,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<AudioTranscriptionAuth>;
|
||||
) -> Result<AudioTranscriptionAuth, Error>;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,56 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
|
||||
|
||||
pub struct AudioTranscriptionRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub audio: Value,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProviderAudioTranscriptionRequest {
|
||||
pub(super) model: String,
|
||||
pub(super) custom_llm_provider: String,
|
||||
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) auth: AudioTranscriptionAuth,
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
pub(super) optional_params: Map<String, Value>,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl ProviderAudioTranscriptionRequest {
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
pub fn custom_llm_provider(&self) -> &str {
|
||||
&self.custom_llm_provider
|
||||
}
|
||||
|
||||
pub fn url(&self) -> &str {
|
||||
&self.url
|
||||
}
|
||||
|
||||
pub fn body(&self) -> &Value {
|
||||
&self.body
|
||||
}
|
||||
|
||||
pub fn with_body(self, body: Value) -> Self {
|
||||
Self { body, ..self }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AudioTranscriptionRequestData {
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::future::Future;
|
||||
use std::time::{Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::{CoreError, CoreResult};
|
||||
use crate::Error;
|
||||
|
||||
pub mod types;
|
||||
|
||||
|
|
@ -11,14 +11,14 @@ pub use types::{
|
|||
};
|
||||
|
||||
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
|
||||
type PreCallFuture<'a>: Future<Output = CoreResult<InitialReq>> + Send + 'a
|
||||
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
|
||||
where
|
||||
Self: 'a,
|
||||
InitialReq: 'a,
|
||||
ProviderReq: 'a,
|
||||
Resp: 'a;
|
||||
|
||||
type DuringCallFuture<'a>: Future<Output = CoreResult<ProviderReq>> + Send + 'a
|
||||
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Error>> + Send + 'a
|
||||
where
|
||||
Self: 'a,
|
||||
InitialReq: 'a,
|
||||
|
|
@ -56,7 +56,7 @@ pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
error: &'a CoreError,
|
||||
error: &'a Error,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a>;
|
||||
}
|
||||
|
|
@ -86,12 +86,12 @@ impl<'a> CallLifecycle<'a> {
|
|||
request: InitialReq,
|
||||
hooks: &Hooks,
|
||||
provider_call: ProviderCall,
|
||||
) -> CoreResult<Resp>
|
||||
) -> Result<Resp, Error>
|
||||
where
|
||||
InitialReq: CallLifecycleRequest,
|
||||
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
|
||||
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
|
||||
ProviderFuture: Future<Output = CoreResult<Resp>>,
|
||||
ProviderFuture: Future<Output = Result<Resp, Error>>,
|
||||
{
|
||||
let context = request.lifecycle_context();
|
||||
self.run(context, request, hooks, provider_call).await
|
||||
|
|
@ -103,11 +103,11 @@ impl<'a> CallLifecycle<'a> {
|
|||
request: InitialReq,
|
||||
hooks: &Hooks,
|
||||
provider_call: ProviderCall,
|
||||
) -> CoreResult<Resp>
|
||||
) -> Result<Resp, Error>
|
||||
where
|
||||
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
|
||||
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
|
||||
ProviderFuture: Future<Output = CoreResult<Resp>>,
|
||||
ProviderFuture: Future<Output = Result<Resp, Error>>,
|
||||
{
|
||||
let call_start = epoch_seconds();
|
||||
let mut phases = Vec::new();
|
||||
|
|
@ -166,7 +166,7 @@ impl<'a> CallLifecycle<'a> {
|
|||
&self,
|
||||
context: &CallLifecycleContext,
|
||||
hooks: &Hooks,
|
||||
error: &CoreError,
|
||||
error: &Error,
|
||||
call_start: f64,
|
||||
phases: &mut Vec<CallLifecyclePhaseTiming>,
|
||||
) where
|
||||
|
|
@ -251,8 +251,8 @@ mod tests {
|
|||
}
|
||||
|
||||
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
|
||||
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type PreCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
|
||||
type SuccessFuture<'a> = BoxFuture<'a, ()>;
|
||||
type FailureFuture<'a> = BoxFuture<'a, ()>;
|
||||
|
||||
|
|
@ -294,7 +294,7 @@ mod tests {
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_error: &'a CoreError,
|
||||
_error: &'a Error,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
|
|
@ -304,8 +304,8 @@ mod tests {
|
|||
}
|
||||
|
||||
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
|
||||
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<RecordingRequest>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, Error>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
|
||||
type SuccessFuture<'a> = BoxFuture<'a, ()>;
|
||||
type FailureFuture<'a> = BoxFuture<'a, ()>;
|
||||
|
||||
|
|
@ -345,7 +345,7 @@ mod tests {
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_error: &'a CoreError,
|
||||
_error: &'a Error,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
|
|
@ -383,13 +383,13 @@ mod tests {
|
|||
"request".to_string(),
|
||||
&hooks,
|
||||
|_request| async move {
|
||||
Err::<String, CoreError>(CoreError::Network("provider down".to_string()))
|
||||
Err::<String, Error>(Error::Network("provider down".to_string()))
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("call fails");
|
||||
|
||||
assert_eq!(error, CoreError::Network("provider down".to_string()));
|
||||
assert_eq!(error, Error::Network("provider down".to_string()));
|
||||
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::ChatCompletionsProviderConfig;
|
||||
|
||||
|
|
@ -23,6 +22,6 @@ pub(super) fn chat_completions_provider_config(
|
|||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::http_utils::truncate_error_body;
|
||||
use crate::error::{Error, as_response_error};
|
||||
use crate::http_utils::{classify_send_error, truncate_error_body};
|
||||
|
||||
use super::client::http_client;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
|
|
@ -11,9 +11,9 @@ use super::types::{
|
|||
|
||||
pub(super) async fn execute_chat_completions_provider_call(
|
||||
request: ProviderChatCompletionsRequest,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let body = serde_json::to_vec(&request.body).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
Error::InvalidRequest(format!(
|
||||
"failed to serialize chat completions request: {err}"
|
||||
))
|
||||
})?;
|
||||
|
|
@ -27,32 +27,23 @@ pub(super) async fn execute_chat_completions_provider_call(
|
|||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request_builder.send().await.map_err(|err| {
|
||||
// Failing to establish the connection means the request never went out,
|
||||
// so the host can still serve it. Everything else here, a timeout
|
||||
// above all, may have reached the provider and been answered.
|
||||
if err.is_connect() || err.is_builder() {
|
||||
CoreError::Connect(err.to_string())
|
||||
} else {
|
||||
CoreError::Network(err.to_string())
|
||||
}
|
||||
})?;
|
||||
let response = request_builder.send().await.map_err(classify_send_error)?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let body: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
|
||||
Error::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
|
||||
})?;
|
||||
request
|
||||
.config
|
||||
|
|
@ -60,27 +51,11 @@ pub(super) async fn execute_chat_completions_provider_call(
|
|||
.map_err(as_response_error)
|
||||
}
|
||||
|
||||
/// Re-tag an error raised while normalizing a response the provider already
|
||||
/// returned.
|
||||
///
|
||||
/// A config reports the same variants on either side of the call: a missing
|
||||
/// field or an unsupported block can mean "this request cannot be translated"
|
||||
/// during prepare and "this response cannot be normalized" here. Only the
|
||||
/// second kind has already been billed, and a host that keeps a reference
|
||||
/// implementation must not retry those, so collapse them to one variant that
|
||||
/// can only mean the provider was already called.
|
||||
pub(super) fn as_response_error(err: CoreError) -> CoreError {
|
||||
match err {
|
||||
already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already,
|
||||
other => CoreError::InvalidResponse(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
|
|
@ -101,7 +76,7 @@ pub(super) async fn signed_headers(
|
|||
.iter()
|
||||
.any(|(name, _)| is_sigv4_computed_header(name))
|
||||
{
|
||||
return Err(CoreError::Unsupported(
|
||||
return Err(Error::Unsupported(
|
||||
"request forwards a header AWS SigV4 computes",
|
||||
));
|
||||
}
|
||||
|
|
@ -137,9 +112,9 @@ pub(super) async fn signed_headers(
|
|||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
_body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
match &request.auth {
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported(
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
|
||||
"AWS SigV4 requires the bedrock-auth feature",
|
||||
)),
|
||||
_ => Ok(request.upstream_headers.clone()),
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@
|
|||
//! credentials, and it resolves the provider, translates the conversation,
|
||||
//! calls the provider, and returns a typed OpenAI-shaped response.
|
||||
|
||||
use crate::Error;
|
||||
mod client;
|
||||
mod common_utils;
|
||||
pub mod conversation;
|
||||
|
|
@ -15,17 +16,18 @@ pub mod response_utils;
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use crate::streaming::OpenedStream;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use handler::execute_chat_completions_provider_call;
|
||||
use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config};
|
||||
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
|
||||
use types::{
|
||||
ChatCompletionsRequest, ChatCompletionsResponse, ChatCompletionsStreamRequest, ChatStreamEvent,
|
||||
};
|
||||
|
||||
pub async fn chat_completions(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await
|
||||
}
|
||||
|
||||
|
|
@ -55,5 +57,51 @@ pub fn chat_completions_decline_reason(
|
|||
.map(|reason| reason.0)
|
||||
}
|
||||
|
||||
pub async fn chat_completions_stream(
|
||||
_request: ChatCompletionsStreamRequest,
|
||||
) -> Result<OpenedStream<ChatStreamEvent>, Error> {
|
||||
Err(crate::Error::Unsupported(
|
||||
"chat completions streaming provider registration",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
#[cfg(test)]
|
||||
mod stream_entrypoint_tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::Error;
|
||||
use crate::streaming::{
|
||||
ProviderCredentials, StreamProviderId, StreamTarget, StreamTransportOptions,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_stream_declines_until_a_provider_is_registered() {
|
||||
let body = serde_json::from_value(json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"stream": true
|
||||
}))
|
||||
.expect("valid chat stream request");
|
||||
let result = chat_completions_stream(ChatCompletionsStreamRequest {
|
||||
body,
|
||||
target: StreamTarget::new(
|
||||
StreamProviderId::Anthropic,
|
||||
ProviderCredentials::default(),
|
||||
None,
|
||||
),
|
||||
transport: StreamTransportOptions::default(),
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Unsupported(
|
||||
"chat completions streaming provider registration"
|
||||
))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use crate::http_utils::has_header;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
|
|
@ -11,7 +11,7 @@ use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsR
|
|||
pub(super) fn resolve_provider_config<'a>(
|
||||
model: &'a str,
|
||||
custom_llm_provider: Option<&'a str>,
|
||||
) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> {
|
||||
) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> {
|
||||
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
|
||||
.or_else(|| {
|
||||
custom_llm_provider.map(|provider| CustomLlmProvider {
|
||||
|
|
@ -20,35 +20,34 @@ pub(super) fn resolve_provider_config<'a>(
|
|||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidProvider(
|
||||
Error::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for chat completions request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let config = chat_completions_provider_config(provider_info.custom_llm_provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
|
||||
.ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
|
||||
Ok((provider_info.model.to_string(), config))
|
||||
}
|
||||
|
||||
pub(super) fn parse_messages(messages: Value) -> CoreResult<Vec<ChatMessage>> {
|
||||
serde_json::from_value(messages).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!("invalid chat completions messages: {err}"))
|
||||
})
|
||||
pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error> {
|
||||
serde_json::from_value(messages)
|
||||
.map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}")))
|
||||
}
|
||||
|
||||
pub(super) fn prepare_chat_completions_call(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ProviderChatCompletionsRequest> {
|
||||
) -> Result<ProviderChatCompletionsRequest, Error> {
|
||||
let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let messages = parse_messages(request.messages)?;
|
||||
if messages.is_empty() {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"chat completions requires at least one message".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
|
||||
return Err(CoreError::Unsupported(reason.0));
|
||||
return Err(Error::Unsupported(reason.0));
|
||||
}
|
||||
|
||||
let mut headers = string_headers(request.extra_headers)?;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::error::CoreError;
|
||||
use crate::error::Error;
|
||||
|
||||
use super::prepare::prepare_chat_completions_call;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
|
|
@ -29,7 +29,7 @@ fn request<'a>(
|
|||
|
||||
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
|
||||
/// carry resolved credentials), so unwrap the failure case by hand.
|
||||
fn decline(request: ChatCompletionsRequest<'_>) -> CoreError {
|
||||
fn decline(request: ChatCompletionsRequest<'_>) -> Error {
|
||||
match prepare_chat_completions_call(request) {
|
||||
Err(error) => error,
|
||||
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
|
||||
|
|
@ -196,7 +196,7 @@ fn declines_an_unsupported_request_before_resolving_credentials() {
|
|||
call.api_key = None;
|
||||
// No api_key is set and no env is consulted: the gate must run first, so the
|
||||
// error is the decline rather than a missing-credential error.
|
||||
assert_eq!(decline(call), CoreError::Unsupported("streaming"));
|
||||
assert_eq!(decline(call), Error::Unsupported("streaming"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -208,7 +208,7 @@ fn rejects_an_unknown_provider() {
|
|||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider("openai".to_string())
|
||||
Error::InvalidProvider("openai".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -221,7 +221,7 @@ fn rejects_a_model_with_no_resolvable_provider() {
|
|||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider(_)
|
||||
Error::InvalidProvider(_)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -234,7 +234,7 @@ fn rejects_an_empty_or_malformed_message_list() {
|
|||
json!([]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest("chat completions requires at least one message".to_string())
|
||||
Error::InvalidRequest("chat completions requires at least one message".to_string())
|
||||
);
|
||||
assert!(matches!(
|
||||
decline(request(
|
||||
|
|
@ -243,7 +243,7 @@ fn rejects_an_empty_or_malformed_message_list() {
|
|||
json!("not a list"),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest(_)
|
||||
Error::InvalidRequest(_)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -258,7 +258,7 @@ fn rejects_non_string_extra_headers() {
|
|||
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
|
||||
assert_eq!(
|
||||
decline(call),
|
||||
CoreError::InvalidRequest(
|
||||
Error::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
|
|
@ -374,7 +374,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
|
|||
.await
|
||||
.expect_err("{forwarded} should decline instead of being signed");
|
||||
assert!(
|
||||
matches!(error, CoreError::Unsupported(_)),
|
||||
matches!(error, Error::Unsupported(_)),
|
||||
"{forwarded} declined as {error:?}, which the host would not fall back on"
|
||||
);
|
||||
}
|
||||
|
|
@ -727,7 +727,7 @@ mod round_trip {
|
|||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
matches!(err, Error::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
|
@ -745,7 +745,7 @@ mod round_trip {
|
|||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
matches!(err, Error::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
|
@ -763,17 +763,13 @@ mod round_trip {
|
|||
.expect_err("upstream rejects");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::Http { status: 429, .. }),
|
||||
matches!(err, Error::Http { status: 429, .. }),
|
||||
"expected a 429, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_connection_that_is_never_established_declines_instead_of_failing() {
|
||||
// Nothing was sent, so nothing was billed and the host can still serve
|
||||
// the request. Classing this with the post-send failures would turn a
|
||||
// recoverable fallback into a user-facing error on exactly the
|
||||
// deployments whose transport is configured only on the Python client.
|
||||
async fn a_connection_that_is_never_established_is_still_terminal() {
|
||||
let port = {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
listener.local_addr().expect("has an address").port()
|
||||
|
|
@ -787,34 +783,8 @@ mod round_trip {
|
|||
.await
|
||||
.expect_err("nothing is listening");
|
||||
assert!(
|
||||
matches!(err, CoreError::Connect(_)),
|
||||
"expected a pre-send connect failure, got {err:?}"
|
||||
matches!(err, Error::Network(_)),
|
||||
"expected a terminal network failure, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
|
||||
use crate::chat_completions::handler::as_response_error;
|
||||
|
||||
for original in [
|
||||
CoreError::MissingField("usage"),
|
||||
CoreError::Unsupported("non-text response content block"),
|
||||
CoreError::InvalidRequest("whatever".to_string()),
|
||||
CoreError::Auth("whatever".to_string()),
|
||||
] {
|
||||
let label = format!("{original:?}");
|
||||
assert!(
|
||||
matches!(as_response_error(original), CoreError::InvalidResponse(_)),
|
||||
"{label} must not stay retryable once the provider has answered"
|
||||
);
|
||||
}
|
||||
// An upstream status is already unambiguous, so it survives intact.
|
||||
assert!(matches!(
|
||||
as_response_error(CoreError::Http {
|
||||
status: 500,
|
||||
body: "boom".to_string()
|
||||
}),
|
||||
CoreError::Http { status: 500, .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,22 @@
|
|||
use crate::Error;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use super::types::{
|
||||
ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
use super::types::{ChatCompletionsStreamRequest, ChatStreamEvent};
|
||||
use crate::streaming::StreamProvider;
|
||||
|
||||
pub trait ChatCompletionsStreamProvider:
|
||||
StreamProvider<ChatCompletionsStreamRequest, ChatStreamEvent>
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> ChatCompletionsStreamProvider for T where
|
||||
T: StreamProvider<ChatCompletionsStreamRequest, ChatStreamEvent>
|
||||
{
|
||||
}
|
||||
|
||||
/// How the upstream call is authenticated. API-key strategies are resolved in
|
||||
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
|
||||
|
|
@ -39,7 +50,7 @@ pub trait ChatCompletionsProviderConfig: Sync {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
|
|
@ -47,7 +58,7 @@ pub trait ChatCompletionsProviderConfig: Sync {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth>;
|
||||
) -> Result<ChatCompletionsAuth, Error>;
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[("content-type", "application/json")]
|
||||
|
|
@ -91,13 +102,13 @@ pub trait ChatCompletionsProviderConfig: Sync {
|
|||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData>;
|
||||
) -> Result<ProviderChatRequestData, Error>;
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse>;
|
||||
) -> Result<ChatCompletionsResponse, Error>;
|
||||
}
|
||||
|
||||
pub fn unsupported_param(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
|
||||
use crate::streaming::{JsonObject, StreamTarget, StreamTransportOptions};
|
||||
|
||||
/// A `/chat/completions` call as it crosses into the core.
|
||||
///
|
||||
|
|
@ -110,3 +111,183 @@ pub struct ChatCompletionsResponse {
|
|||
pub choices: Vec<ChatCompletionsChoice>,
|
||||
pub usage: ChatCompletionsUsage,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ChatStreamRole {
|
||||
Assistant,
|
||||
Developer,
|
||||
Function,
|
||||
System,
|
||||
Tool,
|
||||
User,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatStreamMessageContent {
|
||||
Text(String),
|
||||
Parts(Vec<JsonObject>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatStreamMessage {
|
||||
pub role: ChatStreamRole,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<ChatStreamMessageContent>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatStreamStop {
|
||||
One(String),
|
||||
Many(Vec<String>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatStreamStringOrObject {
|
||||
Name(String),
|
||||
Definition(JsonObject),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsStreamParameters {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_completion_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stop: Option<ChatStreamStop>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stream_options: Option<JsonObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<JsonObject>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<ChatStreamStringOrObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_format: Option<JsonObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_effort: Option<ChatStreamStringOrObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<JsonObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub user: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsStreamRequestBody {
|
||||
pub model: String,
|
||||
pub messages: Vec<ChatStreamMessage>,
|
||||
#[serde(flatten)]
|
||||
pub parameters: ChatCompletionsStreamParameters,
|
||||
}
|
||||
|
||||
pub struct ChatCompletionsStreamRequest {
|
||||
pub body: ChatCompletionsStreamRequestBody,
|
||||
pub target: StreamTarget,
|
||||
pub transport: StreamTransportOptions,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatStreamToolFunctionChunk {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
pub arguments: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<JsonObject>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatStreamToolCallChunk {
|
||||
pub id: Option<String>,
|
||||
#[serde(rename = "type")]
|
||||
pub tool_type: String,
|
||||
pub function: ChatStreamToolFunctionChunk,
|
||||
pub index: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatStreamUsage {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_tokens_details: Option<JsonObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub completion_tokens_details: Option<JsonObject>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatStreamEvent {
|
||||
pub text: String,
|
||||
pub tool_use: Option<ChatStreamToolCallChunk>,
|
||||
pub is_finished: bool,
|
||||
pub finish_reason: String,
|
||||
pub usage: Option<ChatStreamUsage>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub index: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<JsonObject>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod stream_contract_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn request_uses_public_chat_completion_parameter_names() {
|
||||
let request: ChatCompletionsStreamRequestBody = serde_json::from_value(serde_json::json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 32,
|
||||
"stream": true,
|
||||
"tool_choice": "auto"
|
||||
}))
|
||||
.expect("public request shape");
|
||||
|
||||
assert_eq!(request.parameters.max_tokens, Some(32));
|
||||
assert_eq!(request.parameters.stream, Some(true));
|
||||
assert!(matches!(
|
||||
request.parameters.tool_choice,
|
||||
Some(ChatStreamStringOrObject::Name(ref value)) if value == "auto"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn event_matches_python_generic_streaming_chunk_shape() {
|
||||
let event = ChatStreamEvent {
|
||||
text: "hello".to_string(),
|
||||
tool_use: None,
|
||||
is_finished: false,
|
||||
finish_reason: String::new(),
|
||||
usage: None,
|
||||
index: Some(0),
|
||||
provider_specific_fields: None,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(event).expect("serializable event"),
|
||||
serde_json::json!({
|
||||
"text": "hello",
|
||||
"tool_use": null,
|
||||
"is_finished": false,
|
||||
"finish_reason": "",
|
||||
"usage": null,
|
||||
"index": 0
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,8 @@ pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
|
|||
/// Connect timeout for chat completions provider calls, in seconds.
|
||||
pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// `object` field every non-streaming chat completion response carries.
|
||||
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
use thiserror::Error;
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
pub type CoreResult<T> = Result<T, CoreError>;
|
||||
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum CoreError {
|
||||
#[derive(Debug, ThisError, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("expected {expected}, got {actual}")]
|
||||
InvalidType {
|
||||
expected: &'static str,
|
||||
|
|
@ -23,13 +21,6 @@ pub enum CoreError {
|
|||
Http { status: u16, body: String },
|
||||
#[error("upstream network error: {0}")]
|
||||
Network(String),
|
||||
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
|
||||
/// before any byte of the request went out. Nothing was billed, so a host
|
||||
/// that keeps a reference implementation can serve the request itself.
|
||||
/// A timeout is deliberately not this, since the provider may have received
|
||||
/// and answered the request already.
|
||||
#[error("could not reach the provider: {0}")]
|
||||
Connect(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
/// The request is outside the surface this route covers in Rust. Hosts that
|
||||
|
|
@ -38,6 +29,14 @@ pub enum CoreError {
|
|||
Unsupported(&'static str),
|
||||
}
|
||||
|
||||
/// Re-tag an error raised after the provider has already returned a response.
|
||||
pub(crate) fn as_response_error(err: Error) -> Error {
|
||||
match err {
|
||||
already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already,
|
||||
other => Error::InvalidResponse(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
||||
match value {
|
||||
serde_json::Value::Null => "null",
|
||||
|
|
@ -48,3 +47,34 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
|||
serde_json::Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn response_errors_collapse_to_one_non_retryable_variant() {
|
||||
for original in [
|
||||
Error::MissingField("usage"),
|
||||
Error::Unsupported("non-text response content block"),
|
||||
Error::InvalidRequest("whatever".to_string()),
|
||||
Error::Auth("whatever".to_string()),
|
||||
] {
|
||||
assert!(matches!(
|
||||
as_response_error(original),
|
||||
Error::InvalidResponse(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_errors_preserve_an_upstream_status() {
|
||||
assert!(matches!(
|
||||
as_response_error(Error::Http {
|
||||
status: 500,
|
||||
body: "boom".to_string()
|
||||
}),
|
||||
Error::Http { status: 500, .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,11 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::{Error, json_type_name};
|
||||
|
||||
pub(crate) fn classify_send_error(error: reqwest::Error) -> Error {
|
||||
Error::Network(error.to_string())
|
||||
}
|
||||
|
||||
/// Bound an upstream error body before it crosses a host boundary, so provider
|
||||
/// bodies stay data-minimized.
|
||||
|
|
@ -18,7 +22,7 @@ pub fn truncate_error_body(body: &str) -> String {
|
|||
pub fn string_headers(
|
||||
context: &'static str,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
|
|
@ -27,7 +31,7 @@ pub fn string_headers(
|
|||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
Error::InvalidRequest(format!(
|
||||
"{context} extra_headers.{key} must be a string, got {}",
|
||||
json_type_name(&value)
|
||||
))
|
||||
|
|
@ -81,7 +85,7 @@ mod tests {
|
|||
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
Error::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
|
|
|
|||
|
|
@ -12,5 +12,6 @@ pub mod realtime;
|
|||
pub mod responses;
|
||||
pub mod router;
|
||||
pub mod routing_utils;
|
||||
pub mod streaming;
|
||||
|
||||
pub use error::{CoreError, CoreResult};
|
||||
pub use error::Error;
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
|
|
@ -23,6 +22,6 @@ pub(super) fn messages_provider_config(
|
|||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::{Error, as_response_error};
|
||||
use crate::http_utils::classify_send_error;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::common_utils::truncate_error_body;
|
||||
|
|
@ -7,7 +8,7 @@ use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest};
|
|||
|
||||
pub(super) async fn execute_messages_provider_call(
|
||||
request: ProviderMessagesRequest,
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
|
|
@ -16,35 +17,34 @@ pub(super) async fn execute_messages_provider_call(
|
|||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let response = request_builder.send().await.map_err(classify_send_error)?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid messages response JSON: {err}"))
|
||||
})?;
|
||||
request.config.transform_response(&request.model, response)
|
||||
let response = serde_json::from_str(&text)
|
||||
.map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?;
|
||||
request
|
||||
.config
|
||||
.transform_response(&request.model, response)
|
||||
.map_err(as_response_error)
|
||||
}
|
||||
|
||||
pub(super) async fn execute_messages_provider_stream(
|
||||
request: ProviderMessagesRequest,
|
||||
) -> CoreResult<reqwest::Response> {
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"streaming messages is not supported for this provider".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -60,17 +60,169 @@ pub(super) async fn execute_messages_provider_stream(
|
|||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
return Err(CoreError::Http {
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use super::*;
|
||||
use crate::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
struct RejectingResponseConfig;
|
||||
|
||||
impl AnthropicMessagesProviderConfig for RejectingResponseConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
_api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
_api_key: Option<&str>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
_model: &str,
|
||||
_response: AnthropicMessagesResponse,
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
Err(Error::MissingField("normalized_content"))
|
||||
}
|
||||
}
|
||||
|
||||
static REJECTING_RESPONSE_CONFIG: RejectingResponseConfig = RejectingResponseConfig;
|
||||
|
||||
fn request(url: String, timeout: Duration) -> ProviderMessagesRequest {
|
||||
ProviderMessagesRequest {
|
||||
provider: "anthropic".to_string(),
|
||||
model: "claude-test".to_string(),
|
||||
config: &REJECTING_RESPONSE_CONFIG,
|
||||
url,
|
||||
body: json!({}),
|
||||
upstream_headers: Vec::new(),
|
||||
timeout: Some(timeout),
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
loop {
|
||||
let read = socket.read(&mut buffer).await.expect("reads request");
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..read]);
|
||||
if request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn post_response_transform_errors_are_non_retryable() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let _ = read_http_request(&mut socket).await;
|
||||
let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
});
|
||||
|
||||
let error = execute_messages_provider_call(request(
|
||||
format!("http://{addr}/v1/messages"),
|
||||
Duration::from_secs(5),
|
||||
))
|
||||
.await
|
||||
.expect_err("response transform should fail");
|
||||
|
||||
server.await.expect("server task completes");
|
||||
assert!(
|
||||
matches!(error, Error::InvalidResponse(message) if message.contains("normalized_content"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refused_connections_are_terminal_network_errors() {
|
||||
let port = {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
listener.local_addr().expect("has an address").port()
|
||||
};
|
||||
let error = execute_messages_provider_call(request(
|
||||
format!("http://127.0.0.1:{port}"),
|
||||
Duration::from_secs(1),
|
||||
))
|
||||
.await
|
||||
.expect_err("nothing is listening");
|
||||
|
||||
assert!(matches!(error, Error::Network(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn established_request_timeouts_are_network_errors() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("has an address");
|
||||
let (request_received_tx, request_received_rx) = tokio::sync::oneshot::channel();
|
||||
let (release_server_tx, release_server_rx) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let received = read_http_request(&mut socket).await;
|
||||
request_received_tx.send(received).expect("reports request");
|
||||
release_server_rx.await.expect("server is released");
|
||||
});
|
||||
|
||||
let error = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
execute_messages_provider_call(request(
|
||||
format!("http://{addr}"),
|
||||
Duration::from_millis(100),
|
||||
)),
|
||||
)
|
||||
.await
|
||||
.expect("client call completes")
|
||||
.expect_err("established request times out");
|
||||
|
||||
let received = tokio::time::timeout(Duration::from_secs(2), request_received_rx)
|
||||
.await
|
||||
.expect("server observes request")
|
||||
.expect("server reports request");
|
||||
assert!(received.starts_with("POST / "), "{received}");
|
||||
release_server_tx.send(()).expect("releases server");
|
||||
server.await.expect("server task completes");
|
||||
assert!(matches!(error, Error::Network(_)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,26 +7,75 @@
|
|||
//! is the streaming variant; it hands the raw upstream response back so a host
|
||||
//! can splice the event stream to its own caller.
|
||||
|
||||
use crate::Error;
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use crate::error::CoreResult;
|
||||
use crate::streaming::OpenedStream;
|
||||
|
||||
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
|
||||
use prepare::prepare_messages_call;
|
||||
use types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
use types::{
|
||||
AnthropicMessagesResponse, MessagesRequest, MessagesStreamEvent, MessagesStreamRequest,
|
||||
};
|
||||
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<AnthropicMessagesResponse> {
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
|
||||
execute_messages_provider_call(prepare_messages_call(request)?).await
|
||||
}
|
||||
|
||||
pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult<reqwest::Response> {
|
||||
pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
|
||||
execute_messages_provider_stream(prepare_messages_call(request)?).await
|
||||
}
|
||||
|
||||
pub async fn messages_event_stream(
|
||||
_request: MessagesStreamRequest,
|
||||
) -> Result<OpenedStream<MessagesStreamEvent>, Error> {
|
||||
Err(crate::Error::Unsupported(
|
||||
"messages event streaming provider registration",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
#[cfg(test)]
|
||||
mod stream_entrypoint_tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::Error;
|
||||
use crate::streaming::{
|
||||
ProviderCredentials, StreamProviderId, StreamTarget, StreamTransportOptions,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_event_stream_declines_until_a_provider_is_registered() {
|
||||
let body = serde_json::from_value(json!({
|
||||
"model": "claude-sonnet",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 32,
|
||||
"stream": true
|
||||
}))
|
||||
.expect("valid Messages stream request");
|
||||
let result = messages_event_stream(MessagesStreamRequest {
|
||||
body,
|
||||
target: StreamTarget::new(
|
||||
StreamProviderId::Anthropic,
|
||||
ProviderCredentials::default(),
|
||||
None,
|
||||
),
|
||||
transport: StreamTransportOptions::default(),
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Unsupported(
|
||||
"messages event streaming provider registration"
|
||||
))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
|
||||
|
|
@ -7,7 +7,7 @@ use super::types::{MessagesRequest, ProviderMessagesRequest};
|
|||
|
||||
pub(super) fn prepare_messages_call(
|
||||
request: MessagesRequest<'_>,
|
||||
) -> CoreResult<ProviderMessagesRequest> {
|
||||
) -> Result<ProviderMessagesRequest, Error> {
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.or_else(|| {
|
||||
request
|
||||
|
|
@ -18,15 +18,16 @@ pub(super) fn prepare_messages_call(
|
|||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidProvider(
|
||||
Error::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for messages request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let model = provider_info.model.to_string();
|
||||
let provider = provider_info.custom_llm_provider;
|
||||
|
||||
let config = messages_provider_config(provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?;
|
||||
let config = messages_provider_config(provider).ok_or(Error::Unsupported(
|
||||
"messages provider is not registered in the Rust bridge",
|
||||
))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let mut headers = string_headers(request.extra_headers)?;
|
||||
|
|
@ -53,11 +54,11 @@ pub(super) fn prepare_messages_call(
|
|||
|
||||
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
|
||||
let typed_request = serde_json::from_value(request.body).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
|
||||
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
|
||||
})?;
|
||||
let transformed = config.transform_request(typed_request)?;
|
||||
let body = serde_json::to_value(transformed).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
Error::InvalidRequest(format!(
|
||||
"failed to serialize Anthropic messages request: {err}"
|
||||
))
|
||||
})?;
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ use serde_json::{Map, Value, json};
|
|||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::error::CoreError;
|
||||
use crate::error::Error;
|
||||
|
||||
use super::common_utils::{
|
||||
has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
|
||||
|
|
@ -77,7 +77,7 @@ fn truncate_error_body_caps_long_payloads() {
|
|||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
|
||||
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
|
||||
assert!(matches!(err, CoreError::InvalidRequest(_)));
|
||||
assert!(matches!(err, Error::InvalidRequest(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -341,7 +341,7 @@ async fn messages_requires_auth_when_no_key_and_no_header() {
|
|||
.await
|
||||
.expect_err("missing auth errors");
|
||||
|
||||
assert!(matches!(err, CoreError::Auth(_)));
|
||||
assert!(matches!(err, Error::Auth(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -420,7 +420,7 @@ async fn messages_maps_provider_error_status_to_http_error() {
|
|||
.await
|
||||
.expect_err("provider error propagates");
|
||||
|
||||
assert!(matches!(err, CoreError::Http { status: 401, .. }));
|
||||
assert!(matches!(err, Error::Http { status: 401, .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -437,5 +437,8 @@ async fn messages_rejects_unsupported_provider() {
|
|||
.await
|
||||
.expect_err("unsupported provider errors");
|
||||
|
||||
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai"));
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::Unsupported("messages provider is not registered in the Rust bridge")
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,19 @@
|
|||
use crate::error::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::streaming::StreamProvider;
|
||||
|
||||
use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
|
||||
use super::types::{
|
||||
AnthropicMessagesRequest, AnthropicMessagesResponse, MessagesStreamEvent, MessagesStreamRequest,
|
||||
};
|
||||
|
||||
pub trait MessagesStreamProvider:
|
||||
StreamProvider<MessagesStreamRequest, MessagesStreamEvent>
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> MessagesStreamProvider for T where
|
||||
T: StreamProvider<MessagesStreamRequest, MessagesStreamEvent>
|
||||
{
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum MessagesAuthStrategy {
|
||||
|
|
@ -23,13 +36,13 @@ pub trait AnthropicMessagesProviderConfig: Sync {
|
|||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn auth_strategy(&self) -> MessagesAuthStrategy {
|
||||
MessagesAuthStrategy::Header("x-api-key")
|
||||
|
|
@ -49,7 +62,7 @@ pub trait AnthropicMessagesProviderConfig: Sync {
|
|||
fn transform_request(
|
||||
&self,
|
||||
request: AnthropicMessagesRequest,
|
||||
) -> CoreResult<AnthropicMessagesRequest> {
|
||||
) -> Result<AnthropicMessagesRequest, Error> {
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
|
|
@ -57,7 +70,7 @@ pub trait AnthropicMessagesProviderConfig: Sync {
|
|||
&self,
|
||||
_model: &str,
|
||||
response: AnthropicMessagesResponse,
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
use crate::streaming::{JsonObject, StreamTarget, StreamTransportOptions};
|
||||
|
||||
pub struct MessagesRequest<'a> {
|
||||
pub model: &'a str,
|
||||
|
|
@ -47,6 +48,76 @@ pub struct ContentBlock {
|
|||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub struct MessagesStreamRequest {
|
||||
pub body: AnthropicMessagesRequest,
|
||||
pub target: StreamTarget,
|
||||
pub transport: StreamTransportOptions,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum MessagesStreamEvent {
|
||||
MessageStart {
|
||||
message: AnthropicMessagesResponse,
|
||||
},
|
||||
ContentBlockStart {
|
||||
index: u64,
|
||||
content_block: JsonObject,
|
||||
},
|
||||
ContentBlockDelta {
|
||||
index: u64,
|
||||
delta: JsonObject,
|
||||
},
|
||||
ContentBlockStop {
|
||||
index: u64,
|
||||
},
|
||||
MessageDelta {
|
||||
delta: JsonObject,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
usage: Option<JsonObject>,
|
||||
},
|
||||
MessageStop,
|
||||
Ping,
|
||||
Error {
|
||||
error: JsonObject,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod stream_contract_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn message_stop_serializes_as_anthropic_event() {
|
||||
assert_eq!(
|
||||
serde_json::to_value(MessagesStreamEvent::MessageStop).expect("serializable event"),
|
||||
serde_json::json!({"type": "message_stop"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_delta_keeps_typed_event_fields() {
|
||||
let event = MessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: JsonObject(
|
||||
serde_json::json!({"type": "text_delta", "text": "hello"})
|
||||
.as_object()
|
||||
.expect("object")
|
||||
.clone(),
|
||||
),
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(event).expect("serializable event"),
|
||||
serde_json::json!({
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "text_delta", "text": "hello"}
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheControl {
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
|
|
|
|||
234
litellm-rust/crates/core/src/ocr/canonical.rs
Normal file
234
litellm-rust/crates/core/src/ocr/canonical.rs
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use bytes::Bytes;
|
||||
use mime::Mime;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
use super::policy::OcrCanonicalField;
|
||||
use super::types::{Field, OcrDialectId};
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum CanonicalOcrError {
|
||||
#[error("invalid OCR parameter {field}: {reason}")]
|
||||
InvalidParameter {
|
||||
field: &'static str,
|
||||
reason: &'static str,
|
||||
},
|
||||
#[error("provider extra collides with canonical OCR field: {0}")]
|
||||
ExtraCollidesWithCanonical(String),
|
||||
#[error("LiteLLM control cannot enter OCR provider extras: {0}")]
|
||||
ReservedExtra(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum DocumentKind {
|
||||
Image,
|
||||
Pdf,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum OcrDocument {
|
||||
RemoteUrl {
|
||||
kind: DocumentKind,
|
||||
url: Url,
|
||||
},
|
||||
Inline {
|
||||
kind: DocumentKind,
|
||||
media_type: Mime,
|
||||
bytes: Bytes,
|
||||
},
|
||||
ProviderReference {
|
||||
provider: OcrDialectId,
|
||||
id: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct PageSelection(Vec<u32>);
|
||||
|
||||
impl PageSelection {
|
||||
pub fn new(pages: impl IntoIterator<Item = u32>) -> Self {
|
||||
Self(pages.into_iter().collect())
|
||||
}
|
||||
|
||||
pub fn pages(&self) -> &[u32] {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AnnotationFormat(Value);
|
||||
|
||||
impl AnnotationFormat {
|
||||
pub fn new(schema: Value) -> Result<Self, CanonicalOcrError> {
|
||||
if !schema.is_object() {
|
||||
return Err(CanonicalOcrError::InvalidParameter {
|
||||
field: "annotation_format",
|
||||
reason: "must be a JSON object",
|
||||
});
|
||||
}
|
||||
Ok(Self(schema))
|
||||
}
|
||||
|
||||
pub fn as_value(&self) -> &Value {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TableFormat {
|
||||
Html,
|
||||
Markdown,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ConfidenceScoresGranularity {
|
||||
Page,
|
||||
Word,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct OcrOutputOptions {
|
||||
pub include_image_base64: Field<bool>,
|
||||
pub image_limit: Field<u32>,
|
||||
pub image_min_size: Field<u32>,
|
||||
pub bbox_annotation_format: Field<AnnotationFormat>,
|
||||
pub document_annotation_format: Field<AnnotationFormat>,
|
||||
pub document_annotation_prompt: Field<String>,
|
||||
pub extract_header: Field<bool>,
|
||||
pub extract_footer: Field<bool>,
|
||||
pub table_format: Field<TableFormat>,
|
||||
pub confidence_scores_granularity: Field<ConfidenceScoresGranularity>,
|
||||
pub include_blocks: Field<bool>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct OcrRequestId(String);
|
||||
|
||||
impl OcrRequestId {
|
||||
pub fn new(value: impl Into<String>) -> Result<Self, CanonicalOcrError> {
|
||||
let value = value.into();
|
||||
if value.trim().is_empty() {
|
||||
return Err(CanonicalOcrError::InvalidParameter {
|
||||
field: "id",
|
||||
reason: "must not be blank",
|
||||
});
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct ExplicitProviderExtras {
|
||||
dialect: OcrDialectId,
|
||||
fields: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
impl ExplicitProviderExtras {
|
||||
pub fn try_new(
|
||||
dialect: OcrDialectId,
|
||||
fields: BTreeMap<String, Value>,
|
||||
) -> Result<Self, CanonicalOcrError> {
|
||||
if let Some(field) = fields.keys().find(|field| {
|
||||
OcrCanonicalField::from_wire_name(field).is_some() || is_litellm_control(field)
|
||||
}) {
|
||||
return Err(if OcrCanonicalField::from_wire_name(field).is_some() {
|
||||
CanonicalOcrError::ExtraCollidesWithCanonical(field.clone())
|
||||
} else {
|
||||
CanonicalOcrError::ReservedExtra(field.clone())
|
||||
});
|
||||
}
|
||||
Ok(Self { dialect, fields })
|
||||
}
|
||||
|
||||
pub fn dialect(&self) -> OcrDialectId {
|
||||
self.dialect
|
||||
}
|
||||
|
||||
pub fn fields(&self) -> &BTreeMap<String, Value> {
|
||||
&self.fields
|
||||
}
|
||||
}
|
||||
|
||||
fn is_litellm_control(field: &str) -> bool {
|
||||
field.starts_with("litellm_")
|
||||
|| matches!(
|
||||
field,
|
||||
"api_base"
|
||||
| "api_key"
|
||||
| "custom_llm_provider"
|
||||
| "fallbacks"
|
||||
| "metadata"
|
||||
| "mock_response"
|
||||
| "num_retries"
|
||||
| "request_timeout"
|
||||
| "retry_policy"
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct CanonicalOcrRequest {
|
||||
pub model: String,
|
||||
pub document: OcrDocument,
|
||||
pub pages: Field<PageSelection>,
|
||||
pub output: OcrOutputOptions,
|
||||
pub request_id: Field<OcrRequestId>,
|
||||
pub provider_extras: ExplicitProviderExtras,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn extras_reject_canonical_field_collisions() {
|
||||
let error = ExplicitProviderExtras::try_new(
|
||||
OcrDialectId::Mistral,
|
||||
BTreeMap::from([("pages".to_string(), json!([0]))]),
|
||||
)
|
||||
.expect_err("canonical fields cannot enter extras");
|
||||
assert_eq!(
|
||||
error,
|
||||
CanonicalOcrError::ExtraCollidesWithCanonical("pages".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extras_reject_litellm_controls() {
|
||||
for field in ["request_timeout", "litellm_future_control"] {
|
||||
let error = ExplicitProviderExtras::try_new(
|
||||
OcrDialectId::Mistral,
|
||||
BTreeMap::from([(field.to_string(), json!(true))]),
|
||||
)
|
||||
.expect_err("LiteLLM controls cannot enter extras");
|
||||
assert_eq!(error, CanonicalOcrError::ReservedExtra(field.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extras_remain_bound_to_one_dialect() {
|
||||
let extras = ExplicitProviderExtras::try_new(
|
||||
OcrDialectId::ReductoV3,
|
||||
BTreeMap::from([("chunking".to_string(), json!({"size": 1}))]),
|
||||
)
|
||||
.expect("provider field is accepted");
|
||||
assert_eq!(extras.dialect(), OcrDialectId::ReductoV3);
|
||||
assert_eq!(extras.fields()["chunking"], json!({"size": 1}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn canonical_fields_preserve_absent_null_and_value() {
|
||||
assert_ne!(Field::<PageSelection>::Absent, Field::Null);
|
||||
assert_ne!(
|
||||
Field::Null,
|
||||
Field::Value(PageSelection::new([0_u32, 2_u32]))
|
||||
);
|
||||
}
|
||||
}
|
||||
218
litellm-rust/crates/core/src/ocr/compiler.rs
Normal file
218
litellm-rust/crates/core/src/ocr/compiler.rs
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
use super::canonical::{CanonicalOcrRequest, DocumentKind, OcrDocument};
|
||||
use super::plan::{CompletionPlan, DocumentPlan};
|
||||
use super::policy::OcrParameterPolicy;
|
||||
use super::response::NormalizedOcr;
|
||||
use super::types::OcrDialectId;
|
||||
pub use super::wire::{MultipartBodyPlan, MultipartPart, OcrJsonValue, OcrWireBody, OcrWireError};
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum CompileError {
|
||||
#[error("invalid OCR parameter {field}: {reason}")]
|
||||
InvalidParameter {
|
||||
field: &'static str,
|
||||
reason: &'static str,
|
||||
},
|
||||
#[error("OCR dialect is not compiled yet: {0:?}")]
|
||||
UnsupportedDialect(OcrDialectId),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum NormalizeError {
|
||||
#[error("invalid terminal OCR response: {0}")]
|
||||
InvalidPayload(&'static str),
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, PartialEq, Eq)]
|
||||
pub struct OcrCredentials {
|
||||
api_key: Option<String>,
|
||||
oauth_token: Option<String>,
|
||||
}
|
||||
|
||||
impl OcrCredentials {
|
||||
pub fn new(api_key: Option<String>, oauth_token: Option<String>) -> Self {
|
||||
Self {
|
||||
api_key,
|
||||
oauth_token,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn api_key(&self) -> Option<&str> {
|
||||
self.api_key.as_deref()
|
||||
}
|
||||
|
||||
pub fn oauth_token(&self) -> Option<&str> {
|
||||
self.oauth_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct ResolvedOcrTarget {
|
||||
dialect: OcrDialectId,
|
||||
api_base: Url,
|
||||
credentials: OcrCredentials,
|
||||
}
|
||||
|
||||
impl ResolvedOcrTarget {
|
||||
pub fn new(dialect: OcrDialectId, api_base: Url, credentials: OcrCredentials) -> Self {
|
||||
Self {
|
||||
dialect,
|
||||
api_base,
|
||||
credentials,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dialect(&self) -> OcrDialectId {
|
||||
self.dialect
|
||||
}
|
||||
|
||||
pub fn api_base(&self) -> &Url {
|
||||
&self.api_base
|
||||
}
|
||||
|
||||
pub fn credentials(&self) -> &OcrCredentials {
|
||||
&self.credentials
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum ProviderDocument {
|
||||
RemoteUrl {
|
||||
kind: DocumentKind,
|
||||
url: Url,
|
||||
},
|
||||
Inline {
|
||||
kind: DocumentKind,
|
||||
media_type: mime::Mime,
|
||||
bytes: bytes::Bytes,
|
||||
},
|
||||
Reference {
|
||||
id: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum HttpMethod {
|
||||
Get,
|
||||
Post,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct CompiledHttpRequest {
|
||||
pub method: HttpMethod,
|
||||
pub url: Url,
|
||||
pub headers: BTreeMap<String, String>,
|
||||
pub body: OcrWireBody,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct ProviderPayload(Value);
|
||||
|
||||
impl ProviderPayload {
|
||||
pub fn new(value: Value) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
pub fn as_value(&self) -> &Value {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> Value {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrDocumentPolicy {
|
||||
Ready,
|
||||
FetchRemoteUrlAndInline,
|
||||
UploadUnlessProviderReference,
|
||||
}
|
||||
|
||||
pub trait OcrDialectCompiler: Send + Sync {
|
||||
fn parameter_policy(&self) -> &'static OcrParameterPolicy;
|
||||
|
||||
fn prepare_document(
|
||||
&self,
|
||||
document: &OcrDocument,
|
||||
target: &ResolvedOcrTarget,
|
||||
) -> Result<DocumentPlan, CompileError>;
|
||||
|
||||
fn compile_submit(
|
||||
&self,
|
||||
request: &CanonicalOcrRequest,
|
||||
document: ProviderDocument,
|
||||
target: &ResolvedOcrTarget,
|
||||
) -> Result<CompiledHttpRequest, CompileError>;
|
||||
|
||||
fn completion_plan(&self) -> CompletionPlan;
|
||||
|
||||
fn normalize(
|
||||
&self,
|
||||
terminal_response: ProviderPayload,
|
||||
) -> Result<NormalizedOcr, NormalizeError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use bytes::Bytes;
|
||||
use mime::Mime;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn compilation_preserves_the_inline_media_allocation() {
|
||||
let source = Bytes::from_static(b"pdf payload");
|
||||
let source_pointer = source.as_ptr();
|
||||
let canonical_document = OcrDocument::Inline {
|
||||
kind: DocumentKind::Pdf,
|
||||
media_type: "application/pdf".parse::<Mime>().expect("valid MIME type"),
|
||||
bytes: source.clone(),
|
||||
};
|
||||
let OcrDocument::Inline {
|
||||
kind,
|
||||
media_type,
|
||||
bytes,
|
||||
} = &canonical_document
|
||||
else {
|
||||
panic!("inline canonical document expected");
|
||||
};
|
||||
let provider_document = ProviderDocument::Inline {
|
||||
kind: *kind,
|
||||
media_type: media_type.clone(),
|
||||
bytes: bytes.clone(),
|
||||
};
|
||||
let ProviderDocument::Inline {
|
||||
media_type, bytes, ..
|
||||
} = provider_document
|
||||
else {
|
||||
panic!("inline document expected");
|
||||
};
|
||||
let request = CompiledHttpRequest {
|
||||
method: HttpMethod::Post,
|
||||
url: Url::parse("https://example.com/ocr").expect("valid URL"),
|
||||
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
||||
body: OcrWireBody::JsonWithMedia(OcrJsonValue::Object(BTreeMap::from([
|
||||
(
|
||||
"document".to_string(),
|
||||
OcrJsonValue::InlineDataUri { media_type, bytes },
|
||||
),
|
||||
("model".to_string(), OcrJsonValue::Value(json!("ocr-model"))),
|
||||
]))),
|
||||
};
|
||||
|
||||
let OcrWireBody::JsonWithMedia(OcrJsonValue::Object(fields)) = &request.body else {
|
||||
panic!("media JSON body expected");
|
||||
};
|
||||
let OcrJsonValue::InlineDataUri { bytes, .. } = &fields["document"] else {
|
||||
panic!("inline media expected");
|
||||
};
|
||||
assert_eq!(bytes.as_ptr(), source_pointer);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,9 @@
|
|||
pub mod canonical;
|
||||
pub mod compiler;
|
||||
pub mod plan;
|
||||
pub mod policy;
|
||||
pub mod profile;
|
||||
pub mod response;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
pub mod wire;
|
||||
|
|
|
|||
40
litellm-rust/crates/core/src/ocr/plan.rs
Normal file
40
litellm-rust/crates/core/src/ocr/plan.rs
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use url::Url;
|
||||
|
||||
use super::canonical::{DocumentKind, OcrDocument};
|
||||
use super::compiler::ProviderDocument;
|
||||
use super::types::OcrDialectId;
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum DocumentPlan {
|
||||
Ready(ProviderDocument),
|
||||
FetchAndInline(FetchPlan),
|
||||
Upload(UploadPlan),
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct FetchPlan {
|
||||
pub kind: DocumentKind,
|
||||
pub url: Url,
|
||||
pub max_bytes: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct UploadPlan {
|
||||
pub dialect: OcrDialectId,
|
||||
pub document: OcrDocument,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum CompletionPlan {
|
||||
Immediate,
|
||||
Poll(PollPlan),
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct PollPlan {
|
||||
pub operation_location_header: &'static str,
|
||||
pub interval: Duration,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
106
litellm-rust/crates/core/src/ocr/policy.rs
Normal file
106
litellm-rust/crates/core/src/ocr/policy.rs
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ParameterDisposition {
|
||||
Forward,
|
||||
Rename(&'static str),
|
||||
Transform,
|
||||
ConsumeAsConfiguration,
|
||||
Reject,
|
||||
}
|
||||
|
||||
macro_rules! ocr_parameter_schema {
|
||||
($(($variant:ident, $field:ident, $wire_name:literal)),+ $(,)?) => {
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrCanonicalField {
|
||||
$($variant),+
|
||||
}
|
||||
|
||||
impl OcrCanonicalField {
|
||||
pub const ALL: [Self; ocr_parameter_schema!(@count $($variant),+)] = [
|
||||
$(Self::$variant),+
|
||||
];
|
||||
|
||||
pub const fn wire_name(self) -> &'static str {
|
||||
match self {
|
||||
$(Self::$variant => $wire_name),+
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_wire_name(value: &str) -> Option<Self> {
|
||||
match value {
|
||||
$($wire_name => Some(Self::$variant)),+,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct OcrParameterPolicy {
|
||||
$(pub $field: ParameterDisposition),+
|
||||
}
|
||||
|
||||
impl OcrParameterPolicy {
|
||||
pub const fn disposition(self, field: OcrCanonicalField) -> ParameterDisposition {
|
||||
match field {
|
||||
$(OcrCanonicalField::$variant => self.$field),+
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
(@count $($item:ident),+) => {
|
||||
<[()]>::len(&[$(ocr_parameter_schema!(@replace $item ())),+])
|
||||
};
|
||||
(@replace $_item:ident $sub:expr) => { $sub };
|
||||
}
|
||||
|
||||
ocr_parameter_schema!(
|
||||
(Pages, pages, "pages"),
|
||||
(
|
||||
IncludeImageBase64,
|
||||
include_image_base64,
|
||||
"include_image_base64"
|
||||
),
|
||||
(ImageLimit, image_limit, "image_limit"),
|
||||
(ImageMinSize, image_min_size, "image_min_size"),
|
||||
(
|
||||
BboxAnnotationFormat,
|
||||
bbox_annotation_format,
|
||||
"bbox_annotation_format"
|
||||
),
|
||||
(
|
||||
DocumentAnnotationFormat,
|
||||
document_annotation_format,
|
||||
"document_annotation_format"
|
||||
),
|
||||
(
|
||||
DocumentAnnotationPrompt,
|
||||
document_annotation_prompt,
|
||||
"document_annotation_prompt"
|
||||
),
|
||||
(ExtractHeader, extract_header, "extract_header"),
|
||||
(ExtractFooter, extract_footer, "extract_footer"),
|
||||
(TableFormat, table_format, "table_format"),
|
||||
(
|
||||
ConfidenceScoresGranularity,
|
||||
confidence_scores_granularity,
|
||||
"confidence_scores_granularity"
|
||||
),
|
||||
(IncludeBlocks, include_blocks, "include_blocks"),
|
||||
(RequestId, request_id, "id"),
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn every_canonical_field_round_trips_through_its_wire_name() {
|
||||
assert_eq!(OcrCanonicalField::ALL.len(), 13);
|
||||
for field in OcrCanonicalField::ALL {
|
||||
assert_eq!(
|
||||
OcrCanonicalField::from_wire_name(field.wire_name()),
|
||||
Some(field)
|
||||
);
|
||||
}
|
||||
assert_eq!(OcrCanonicalField::from_wire_name("provider_private"), None);
|
||||
}
|
||||
}
|
||||
189
litellm-rust/crates/core/src/ocr/profile.rs
Normal file
189
litellm-rust/crates/core/src/ocr/profile.rs
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
use super::compiler::OcrDocumentPolicy;
|
||||
use super::policy::{OcrParameterPolicy, ParameterDisposition};
|
||||
use super::types::OcrDialectId;
|
||||
|
||||
pub const MISTRAL_OCR_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
|
||||
pages: ParameterDisposition::Forward,
|
||||
include_image_base64: ParameterDisposition::Forward,
|
||||
image_limit: ParameterDisposition::Forward,
|
||||
image_min_size: ParameterDisposition::Forward,
|
||||
bbox_annotation_format: ParameterDisposition::Forward,
|
||||
document_annotation_format: ParameterDisposition::Forward,
|
||||
document_annotation_prompt: ParameterDisposition::Forward,
|
||||
extract_header: ParameterDisposition::Forward,
|
||||
extract_footer: ParameterDisposition::Forward,
|
||||
table_format: ParameterDisposition::Forward,
|
||||
confidence_scores_granularity: ParameterDisposition::Forward,
|
||||
include_blocks: ParameterDisposition::Forward,
|
||||
request_id: ParameterDisposition::Forward,
|
||||
};
|
||||
|
||||
pub const AZURE_DOCUMENT_INTELLIGENCE_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
|
||||
pages: ParameterDisposition::Transform,
|
||||
include_image_base64: ParameterDisposition::Reject,
|
||||
image_limit: ParameterDisposition::Reject,
|
||||
image_min_size: ParameterDisposition::Reject,
|
||||
bbox_annotation_format: ParameterDisposition::Reject,
|
||||
document_annotation_format: ParameterDisposition::Reject,
|
||||
document_annotation_prompt: ParameterDisposition::Reject,
|
||||
extract_header: ParameterDisposition::Reject,
|
||||
extract_footer: ParameterDisposition::Reject,
|
||||
table_format: ParameterDisposition::Reject,
|
||||
confidence_scores_granularity: ParameterDisposition::Reject,
|
||||
include_blocks: ParameterDisposition::Reject,
|
||||
request_id: ParameterDisposition::Reject,
|
||||
};
|
||||
|
||||
pub const REJECT_CANONICAL_OCR_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
|
||||
pages: ParameterDisposition::Reject,
|
||||
include_image_base64: ParameterDisposition::Reject,
|
||||
image_limit: ParameterDisposition::Reject,
|
||||
image_min_size: ParameterDisposition::Reject,
|
||||
bbox_annotation_format: ParameterDisposition::Reject,
|
||||
document_annotation_format: ParameterDisposition::Reject,
|
||||
document_annotation_prompt: ParameterDisposition::Reject,
|
||||
extract_header: ParameterDisposition::Reject,
|
||||
extract_footer: ParameterDisposition::Reject,
|
||||
table_format: ParameterDisposition::Reject,
|
||||
confidence_scores_granularity: ParameterDisposition::Reject,
|
||||
include_blocks: ParameterDisposition::Reject,
|
||||
request_id: ParameterDisposition::Reject,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct OcrPollingProfile {
|
||||
pub operation_location_header: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct OcrDialectProfile {
|
||||
pub dialect: OcrDialectId,
|
||||
pub parameter_policy: &'static OcrParameterPolicy,
|
||||
pub document_policy: OcrDocumentPolicy,
|
||||
pub polling: Option<OcrPollingProfile>,
|
||||
}
|
||||
|
||||
pub const OCR_DIALECT_PROFILES: [OcrDialectProfile; 7] = [
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::Mistral,
|
||||
parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::Ready,
|
||||
polling: None,
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::AzureFoundryMistral,
|
||||
parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::FetchRemoteUrlAndInline,
|
||||
polling: None,
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::AzureDocumentIntelligence,
|
||||
parameter_policy: &AZURE_DOCUMENT_INTELLIGENCE_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::Ready,
|
||||
polling: Some(OcrPollingProfile {
|
||||
operation_location_header: "operation-location",
|
||||
}),
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::VertexMistral,
|
||||
parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::FetchRemoteUrlAndInline,
|
||||
polling: None,
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::VertexDeepSeek,
|
||||
parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::Ready,
|
||||
polling: None,
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::ReductoV3,
|
||||
parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::UploadUnlessProviderReference,
|
||||
polling: None,
|
||||
},
|
||||
OcrDialectProfile {
|
||||
dialect: OcrDialectId::ReductoLegacy,
|
||||
parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
|
||||
document_policy: OcrDocumentPolicy::UploadUnlessProviderReference,
|
||||
polling: None,
|
||||
},
|
||||
];
|
||||
|
||||
pub const fn ocr_dialect_profile(dialect: OcrDialectId) -> &'static OcrDialectProfile {
|
||||
match dialect {
|
||||
OcrDialectId::Mistral => &OCR_DIALECT_PROFILES[0],
|
||||
OcrDialectId::AzureFoundryMistral => &OCR_DIALECT_PROFILES[1],
|
||||
OcrDialectId::AzureDocumentIntelligence => &OCR_DIALECT_PROFILES[2],
|
||||
OcrDialectId::VertexMistral => &OCR_DIALECT_PROFILES[3],
|
||||
OcrDialectId::VertexDeepSeek => &OCR_DIALECT_PROFILES[4],
|
||||
OcrDialectId::ReductoV3 => &OCR_DIALECT_PROFILES[5],
|
||||
OcrDialectId::ReductoLegacy => &OCR_DIALECT_PROFILES[6],
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::ocr::policy::{OcrCanonicalField, ParameterDisposition};
|
||||
|
||||
#[test]
|
||||
fn mistral_compatible_dialects_share_parameter_rules_but_not_document_rules() {
|
||||
let mistral = ocr_dialect_profile(OcrDialectId::Mistral);
|
||||
let foundry = ocr_dialect_profile(OcrDialectId::AzureFoundryMistral);
|
||||
let vertex = ocr_dialect_profile(OcrDialectId::VertexMistral);
|
||||
|
||||
for field in OcrCanonicalField::ALL {
|
||||
let expected = mistral.parameter_policy.disposition(field);
|
||||
assert_eq!(foundry.parameter_policy.disposition(field), expected);
|
||||
assert_eq!(vertex.parameter_policy.disposition(field), expected);
|
||||
}
|
||||
assert_eq!(mistral.document_policy, OcrDocumentPolicy::Ready);
|
||||
assert_eq!(
|
||||
foundry.document_policy,
|
||||
OcrDocumentPolicy::FetchRemoteUrlAndInline
|
||||
);
|
||||
assert_eq!(
|
||||
vertex.document_policy,
|
||||
OcrDocumentPolicy::FetchRemoteUrlAndInline
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_mistral_profiles_preserve_provider_specific_boundaries() {
|
||||
let azure = ocr_dialect_profile(OcrDialectId::AzureDocumentIntelligence);
|
||||
assert_eq!(
|
||||
azure.parameter_policy.disposition(OcrCanonicalField::Pages),
|
||||
ParameterDisposition::Transform
|
||||
);
|
||||
assert_eq!(
|
||||
azure.polling,
|
||||
Some(OcrPollingProfile {
|
||||
operation_location_header: "operation-location"
|
||||
})
|
||||
);
|
||||
|
||||
for dialect in [OcrDialectId::ReductoV3, OcrDialectId::ReductoLegacy] {
|
||||
let reducto = ocr_dialect_profile(dialect);
|
||||
assert_eq!(
|
||||
reducto.document_policy,
|
||||
OcrDocumentPolicy::UploadUnlessProviderReference
|
||||
);
|
||||
assert!(OcrCanonicalField::ALL.iter().all(|field| {
|
||||
reducto.parameter_policy.disposition(*field) == ParameterDisposition::Reject
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_dialect_has_exactly_one_profile() {
|
||||
for (index, profile) in OCR_DIALECT_PROFILES.iter().enumerate() {
|
||||
assert_eq!(ocr_dialect_profile(profile.dialect), profile);
|
||||
assert!(
|
||||
OCR_DIALECT_PROFILES[index + 1..]
|
||||
.iter()
|
||||
.all(|other| other.dialect != profile.dialect)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
79
litellm-rust/crates/core/src/ocr/response.rs
Normal file
79
litellm-rust/crates/core/src/ocr/response.rs
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct OcrPage {
|
||||
pub index: u32,
|
||||
pub markdown: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrUsage {
|
||||
pub pages_processed: Option<u32>,
|
||||
pub credits: Option<f64>,
|
||||
pub document_size_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct NormalizedOcr {
|
||||
pub pages: Vec<OcrPage>,
|
||||
pub model: String,
|
||||
pub document_annotation: Option<Value>,
|
||||
pub content: Option<String>,
|
||||
pub tables: Option<Vec<Value>>,
|
||||
pub key_value_pairs: Option<Vec<Value>>,
|
||||
pub usage: OcrUsage,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct NativeOcrPayload(Value);
|
||||
|
||||
impl NativeOcrPayload {
|
||||
pub fn new(value: Value) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
pub fn as_value(&self) -> &Value {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> Value {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct OcrOutcome {
|
||||
pub normalized: NormalizedOcr,
|
||||
pub native: NativeOcrPayload,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn normalized_serialization_excludes_the_native_payload() {
|
||||
let outcome = OcrOutcome {
|
||||
normalized: NormalizedOcr {
|
||||
pages: vec![OcrPage {
|
||||
index: 0,
|
||||
markdown: "portable".to_string(),
|
||||
}],
|
||||
model: "ocr-model".to_string(),
|
||||
document_annotation: None,
|
||||
content: None,
|
||||
tables: None,
|
||||
key_value_pairs: None,
|
||||
usage: OcrUsage::default(),
|
||||
},
|
||||
native: NativeOcrPayload::new(json!({"provider_secret_field": "native"})),
|
||||
};
|
||||
let public =
|
||||
serde_json::to_value(&outcome.normalized).expect("normalized output serializes");
|
||||
assert_eq!(public["pages"][0]["markdown"], "portable");
|
||||
assert!(public.get("provider_secret_field").is_none());
|
||||
assert_eq!(outcome.native.as_value()["provider_secret_field"], "native");
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
use crate::Error;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::CoreResult;
|
||||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
|
|
@ -26,16 +25,16 @@ pub enum OcrResponseHandling {
|
|||
}
|
||||
|
||||
pub trait OcrProviderConfig: Sync {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str];
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str];
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
let mut mapped_params = Map::new();
|
||||
for (param, value) in non_default_params {
|
||||
if self.supported_ocr_params().contains(¶m.as_str()) {
|
||||
mapped_params.insert(param.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
mapped_params
|
||||
let supported_params = self.get_supported_ocr_params();
|
||||
non_default_params
|
||||
.iter()
|
||||
.filter(|(param, _)| supported_params.contains(¶m.as_str()))
|
||||
.map(|(param, value)| (param.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
|
|
@ -43,13 +42,13 @@ pub trait OcrProviderConfig: Sync {
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData>;
|
||||
) -> Result<OcrRequestData, Error>;
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
) -> Result<OcrResponseData, Error>;
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
|
|
@ -57,13 +56,13 @@ pub trait OcrProviderConfig: Sync {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Bearer
|
||||
|
|
|
|||
|
|
@ -1,6 +1,34 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Field<T> {
|
||||
Absent,
|
||||
Null,
|
||||
Value(T),
|
||||
}
|
||||
|
||||
impl<T> Field<T> {
|
||||
pub fn as_ref(&self) -> Field<&T> {
|
||||
match self {
|
||||
Self::Absent => Field::Absent,
|
||||
Self::Null => Field::Null,
|
||||
Self::Value(value) => Field::Value(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrDialectId {
|
||||
Mistral,
|
||||
AzureFoundryMistral,
|
||||
AzureDocumentIntelligence,
|
||||
VertexMistral,
|
||||
VertexDeepSeek,
|
||||
ReductoV3,
|
||||
ReductoLegacy,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrRequestData {
|
||||
pub data: Value,
|
||||
|
|
|
|||
358
litellm-rust/crates/core/src/ocr/wire.rs
Normal file
358
litellm-rust/crates/core/src/ocr/wire.rs
Normal file
|
|
@ -0,0 +1,358 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::io::{self, Write};
|
||||
use std::str::Utf8Error;
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use bytes::Bytes;
|
||||
use mime::Mime;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
const BASE64_INPUT_CHUNK_SIZE: usize = 48 * 1024;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum OcrWireError {
|
||||
#[error("failed to encode OCR JSON body: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("failed to write OCR body: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("encoded OCR data URI is not UTF-8: {0}")]
|
||||
InvalidDataUri(#[from] Utf8Error),
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum OcrJsonValue {
|
||||
Value(Value),
|
||||
Array(Vec<Self>),
|
||||
Object(BTreeMap<String, Self>),
|
||||
InlineDataUri { media_type: Mime, bytes: Bytes },
|
||||
EncodedDataUri(Bytes),
|
||||
}
|
||||
|
||||
impl OcrJsonValue {
|
||||
fn write_to(&self, writer: &mut impl Write) -> Result<(), OcrWireError> {
|
||||
match self {
|
||||
Self::Value(value) => serde_json::to_writer(writer, value).map_err(Into::into),
|
||||
Self::Array(values) => {
|
||||
writer.write_all(b"[")?;
|
||||
for (index, value) in values.iter().enumerate() {
|
||||
if index != 0 {
|
||||
writer.write_all(b",")?;
|
||||
}
|
||||
value.write_to(writer)?;
|
||||
}
|
||||
writer.write_all(b"]")?;
|
||||
Ok(())
|
||||
}
|
||||
Self::Object(fields) => {
|
||||
writer.write_all(b"{")?;
|
||||
for (index, (key, value)) in fields.iter().enumerate() {
|
||||
if index != 0 {
|
||||
writer.write_all(b",")?;
|
||||
}
|
||||
serde_json::to_writer(&mut *writer, key)?;
|
||||
writer.write_all(b":")?;
|
||||
value.write_to(writer)?;
|
||||
}
|
||||
writer.write_all(b"}")?;
|
||||
Ok(())
|
||||
}
|
||||
Self::InlineDataUri { media_type, bytes } => {
|
||||
writer.write_all(b"\"data:")?;
|
||||
writer.write_all(media_type.as_ref().as_bytes())?;
|
||||
writer.write_all(b";base64,")?;
|
||||
for chunk in bytes.chunks(BASE64_INPUT_CHUNK_SIZE) {
|
||||
let encoded = STANDARD.encode(chunk);
|
||||
writer.write_all(encoded.as_bytes())?;
|
||||
}
|
||||
writer.write_all(b"\"")?;
|
||||
Ok(())
|
||||
}
|
||||
Self::EncodedDataUri(data_uri) => {
|
||||
let data_uri = std::str::from_utf8(data_uri)?;
|
||||
serde_json::to_writer(writer, data_uri).map_err(Into::into)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum MultipartPart {
|
||||
Text {
|
||||
name: String,
|
||||
value: String,
|
||||
},
|
||||
Json {
|
||||
name: String,
|
||||
value: Value,
|
||||
},
|
||||
File {
|
||||
name: String,
|
||||
file_name: String,
|
||||
media_type: Mime,
|
||||
bytes: Bytes,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct MultipartBodyPlan {
|
||||
boundary: String,
|
||||
parts: Vec<MultipartPart>,
|
||||
}
|
||||
|
||||
impl MultipartBodyPlan {
|
||||
pub fn new(boundary: impl Into<String>, parts: Vec<MultipartPart>) -> Self {
|
||||
Self {
|
||||
boundary: boundary.into(),
|
||||
parts,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn boundary(&self) -> &str {
|
||||
&self.boundary
|
||||
}
|
||||
|
||||
pub fn parts(&self) -> &[MultipartPart] {
|
||||
&self.parts
|
||||
}
|
||||
|
||||
fn write_to(&self, writer: &mut impl Write) -> Result<(), OcrWireError> {
|
||||
for part in &self.parts {
|
||||
write!(writer, "--{}\r\n", self.boundary)?;
|
||||
match part {
|
||||
MultipartPart::Text { name, value } => {
|
||||
write!(
|
||||
writer,
|
||||
"Content-Disposition: form-data; name=\"{name}\"\r\n\r\n{value}\r\n"
|
||||
)?;
|
||||
}
|
||||
MultipartPart::Json { name, value } => {
|
||||
write!(
|
||||
writer,
|
||||
"Content-Disposition: form-data; name=\"{name}\"\r\nContent-Type: application/json\r\n\r\n"
|
||||
)?;
|
||||
serde_json::to_writer(&mut *writer, value)?;
|
||||
writer.write_all(b"\r\n")?;
|
||||
}
|
||||
MultipartPart::File {
|
||||
name,
|
||||
file_name,
|
||||
media_type,
|
||||
bytes,
|
||||
} => {
|
||||
write!(
|
||||
writer,
|
||||
"Content-Disposition: form-data; name=\"{name}\"; filename=\"{file_name}\"\r\nContent-Type: {media_type}\r\n\r\n"
|
||||
)?;
|
||||
writer.write_all(bytes)?;
|
||||
writer.write_all(b"\r\n")?;
|
||||
}
|
||||
}
|
||||
}
|
||||
write!(writer, "--{}--\r\n", self.boundary)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum OcrWireBody {
|
||||
Json(Value),
|
||||
JsonWithMedia(OcrJsonValue),
|
||||
Multipart(MultipartBodyPlan),
|
||||
}
|
||||
|
||||
impl OcrWireBody {
|
||||
pub fn content_type(&self) -> String {
|
||||
match self {
|
||||
Self::Json(_) | Self::JsonWithMedia(_) => "application/json".to_string(),
|
||||
Self::Multipart(plan) => {
|
||||
format!("multipart/form-data; boundary={}", plan.boundary())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write_to(&self, mut writer: impl Write) -> Result<(), OcrWireError> {
|
||||
match self {
|
||||
Self::Json(value) => serde_json::to_writer(writer, value).map_err(Into::into),
|
||||
Self::JsonWithMedia(value) => value.write_to(&mut writer),
|
||||
Self::Multipart(plan) => plan.write_to(&mut writer),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn ordinary_json_body_writes_without_a_media_plan() {
|
||||
let body = OcrWireBody::Json(json!({"model": "ocr-model", "pages": [0, 2]}));
|
||||
let mut encoded = Vec::new();
|
||||
|
||||
body.write_to(&mut encoded).expect("body writes");
|
||||
|
||||
assert_eq!(body.content_type(), "application/json");
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded).expect("valid JSON"),
|
||||
json!({"model": "ocr-model", "pages": [0, 2]})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_with_media_streams_raw_bytes_as_a_data_uri() {
|
||||
let owner: Arc<[u8]> = vec![b'x'; BASE64_INPUT_CHUNK_SIZE + 1].into();
|
||||
let bytes = Bytes::from_owner(Arc::clone(&owner));
|
||||
let source_pointer = bytes.as_ptr();
|
||||
let body = OcrWireBody::JsonWithMedia(OcrJsonValue::Object(BTreeMap::from([
|
||||
(
|
||||
"document".to_string(),
|
||||
OcrJsonValue::InlineDataUri {
|
||||
media_type: "application/pdf".parse().expect("valid MIME type"),
|
||||
bytes,
|
||||
},
|
||||
),
|
||||
("model".to_string(), OcrJsonValue::Value(json!("ocr-model"))),
|
||||
])));
|
||||
|
||||
let OcrWireBody::JsonWithMedia(OcrJsonValue::Object(fields)) = &body else {
|
||||
panic!("media JSON body must retain its typed representation");
|
||||
};
|
||||
let OcrJsonValue::InlineDataUri { bytes, .. } = &fields["document"] else {
|
||||
panic!("document must remain shared binary media");
|
||||
};
|
||||
assert_eq!(bytes.as_ptr(), source_pointer);
|
||||
assert_eq!(Arc::strong_count(&owner), 2);
|
||||
|
||||
let mut encoded = Vec::new();
|
||||
body.write_to(&mut encoded).expect("body writes");
|
||||
let expected_data_uri = format!(
|
||||
"data:application/pdf;base64,{}",
|
||||
STANDARD.encode(owner.as_ref())
|
||||
);
|
||||
let expected = json!({"document": expected_data_uri, "model": "ocr-model"});
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded).expect("valid JSON"),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_media_encoding_uses_bounded_writes() {
|
||||
let bytes = Bytes::from(vec![b'x'; BASE64_INPUT_CHUNK_SIZE * 3 + 1]);
|
||||
let body = OcrWireBody::JsonWithMedia(OcrJsonValue::InlineDataUri {
|
||||
media_type: "application/pdf".parse().expect("valid MIME type"),
|
||||
bytes,
|
||||
});
|
||||
let mut sink = BoundedSink {
|
||||
maximum_write: BASE64_INPUT_CHUNK_SIZE * 4 / 3,
|
||||
written: 0,
|
||||
};
|
||||
|
||||
body.write_to(&mut sink).expect("writes remain bounded");
|
||||
|
||||
assert!(sink.written > BASE64_INPUT_CHUNK_SIZE * 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoded_data_uri_is_retained_without_decoding_or_copying() {
|
||||
let data_uri = Bytes::from_static(b"data:image/png;base64,aGVsbG8=");
|
||||
let source_pointer = data_uri.as_ptr();
|
||||
let body = OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(data_uri));
|
||||
|
||||
let OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(retained)) = &body else {
|
||||
panic!("encoded data URI must remain bytes");
|
||||
};
|
||||
assert_eq!(retained.as_ptr(), source_pointer);
|
||||
|
||||
let mut encoded = Vec::new();
|
||||
body.write_to(&mut encoded).expect("body writes");
|
||||
assert_eq!(encoded, br#""data:image/png;base64,aGVsbG8=""#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoded_data_uri_is_json_escaped_without_changing_its_allocation() {
|
||||
let data_uri = Bytes::from_static(b"data:text/plain,quoted%20\"value\"");
|
||||
let source_pointer = data_uri.as_ptr();
|
||||
let body = OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(data_uri));
|
||||
|
||||
let mut encoded = Vec::new();
|
||||
body.write_to(&mut encoded).expect("body writes");
|
||||
|
||||
let OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(retained)) = &body else {
|
||||
panic!("encoded data URI expected");
|
||||
};
|
||||
assert_eq!(retained.as_ptr(), source_pointer);
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<String>(&encoded).expect("valid JSON string"),
|
||||
"data:text/plain,quoted%20\"value\""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multipart_file_is_replayable_and_retains_shared_bytes() {
|
||||
let file = Bytes::from_static(b"large-pdf-payload");
|
||||
let source_pointer = file.as_ptr();
|
||||
let body = OcrWireBody::Multipart(MultipartBodyPlan::new(
|
||||
"ocr-boundary",
|
||||
vec![MultipartPart::File {
|
||||
name: "file".to_string(),
|
||||
file_name: "document.pdf".to_string(),
|
||||
media_type: "application/pdf".parse().expect("valid MIME type"),
|
||||
bytes: file,
|
||||
}],
|
||||
));
|
||||
|
||||
let OcrWireBody::Multipart(plan) = &body else {
|
||||
panic!("multipart body expected");
|
||||
};
|
||||
let MultipartPart::File { bytes, .. } = &plan.parts()[0] else {
|
||||
panic!("file part expected");
|
||||
};
|
||||
assert_eq!(bytes.as_ptr(), source_pointer);
|
||||
|
||||
let mut first = Vec::new();
|
||||
let mut retry = Vec::new();
|
||||
body.write_to(&mut first).expect("first write succeeds");
|
||||
body.write_to(&mut retry).expect("retry write succeeds");
|
||||
assert_eq!(first, retry);
|
||||
assert!(
|
||||
first
|
||||
.windows(file_name_marker().len())
|
||||
.any(|window| window == file_name_marker())
|
||||
);
|
||||
assert!(
|
||||
first
|
||||
.windows(bytes.len())
|
||||
.any(|window| window == bytes.as_ref())
|
||||
);
|
||||
}
|
||||
|
||||
fn file_name_marker() -> &'static [u8] {
|
||||
b"filename=\"document.pdf\""
|
||||
}
|
||||
|
||||
struct BoundedSink {
|
||||
maximum_write: usize,
|
||||
written: usize,
|
||||
}
|
||||
|
||||
impl Write for BoundedSink {
|
||||
fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
|
||||
if buffer.len() > self.maximum_write {
|
||||
return Err(io::Error::other("write exceeded bound"));
|
||||
}
|
||||
self.written += buffer.len();
|
||||
Ok(buffer.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
use super::*;
|
||||
use crate::Error;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
|
|
@ -19,7 +20,7 @@ fn transform(model: &str, msgs: Value, opts: Value) -> Value {
|
|||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
fn transform_response(body: Value) -> Result<ChatCompletionsResponse, Error> {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_response("claude-sonnet-4-5", ProviderChatResponseData { body })
|
||||
}
|
||||
|
|
@ -390,29 +391,26 @@ fn declines_a_response_carrying_a_non_text_block() {
|
|||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect_err("non-text block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
assert_eq!(err, Error::Unsupported("non-text response content block"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("messages response is not an object".to_string())
|
||||
Error::InvalidResponse("messages response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"),
|
||||
CoreError::MissingField("content")
|
||||
Error::MissingField("content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "content": []})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
Error::MissingField("usage")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"content": [], "usage": {}})).expect_err("no model"),
|
||||
CoreError::MissingField("model")
|
||||
Error::MissingField("model")
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ use crate::chat_completions::types::{
|
|||
ProviderChatRequestData, ProviderChatResponseData,
|
||||
};
|
||||
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use crate::providers::anthropic::messages::transformation::{
|
||||
complete_anthropic_url, resolve_anthropic_api_key,
|
||||
};
|
||||
|
|
@ -74,7 +74,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
Ok(complete_anthropic_url(api_base, env_lookup))
|
||||
}
|
||||
|
||||
|
|
@ -84,7 +84,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
) -> Result<ChatCompletionsAuth, Error> {
|
||||
Ok(ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
value: resolve_anthropic_api_key(api_key, env_lookup)?,
|
||||
|
|
@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
) -> Result<ProviderChatRequestData, Error> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: anthropic_body(model, &build_conversation(&messages), optional_params),
|
||||
})
|
||||
|
|
@ -147,15 +147,16 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
&self,
|
||||
_model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("messages response is not an object".into())
|
||||
})?;
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let body = response
|
||||
.body
|
||||
.as_object()
|
||||
.ok_or_else(|| Error::InvalidResponse("messages response is not an object".into()))?;
|
||||
|
||||
let content = body
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("content"))?;
|
||||
.ok_or(Error::MissingField("content"))?;
|
||||
// The route declines tool and thinking requests, so a non-text block
|
||||
// means the response carries something this path never asked for.
|
||||
// Decline rather than silently dropping it; the host falls back.
|
||||
|
|
@ -163,7 +164,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
.iter()
|
||||
.any(|block| block.get("type").and_then(Value::as_str) != Some("text"))
|
||||
{
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
return Err(Error::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
|
|
@ -173,7 +174,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
.ok_or(Error::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
|
|
@ -181,7 +182,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
|||
model: body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("model"))?
|
||||
.ok_or(Error::MissingField("model"))?
|
||||
.to_string(),
|
||||
choices: vec![ChatCompletionsChoice {
|
||||
index: 0,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
|
||||
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
|
||||
|
|
@ -17,12 +17,12 @@ pub fn non_empty(value: Option<&str>) -> Option<&str> {
|
|||
pub fn resolve_anthropic_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
Error::Auth(
|
||||
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \
|
||||
environment variable"
|
||||
.to_string(),
|
||||
|
|
@ -52,7 +52,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
|
|||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
Ok(complete_anthropic_url(api_base, env_lookup))
|
||||
}
|
||||
|
||||
|
|
@ -60,7 +60,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_anthropic_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -121,7 +121,7 @@ mod tests {
|
|||
);
|
||||
assert!(matches!(
|
||||
resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"),
|
||||
CoreError::Auth(_)
|
||||
Error::Auth(_)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
use crate::messages::types::{
|
||||
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
|
||||
|
|
@ -28,12 +28,12 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
|
|||
pub fn resolve_azure_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
Error::Auth(
|
||||
"Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable"
|
||||
.to_string(),
|
||||
)
|
||||
|
|
@ -43,12 +43,12 @@ pub fn resolve_azure_api_key(
|
|||
pub fn complete_azure_anthropic_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let api_base = non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
Error::Auth(
|
||||
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \
|
||||
Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
|
||||
.to_string(),
|
||||
|
|
@ -147,7 +147,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
|
|||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
complete_azure_anthropic_url(api_base, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -155,7 +155,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_azure_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -174,7 +174,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
|
|||
fn transform_request(
|
||||
&self,
|
||||
request: AnthropicMessagesRequest,
|
||||
) -> CoreResult<AnthropicMessagesRequest> {
|
||||
) -> Result<AnthropicMessagesRequest, Error> {
|
||||
let mut request = fold_system_role_messages(request);
|
||||
if let Some(system) = request.system.as_mut() {
|
||||
strip_scope_from_system(system);
|
||||
|
|
@ -190,7 +190,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response: AnthropicMessagesResponse,
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
self.anthropic.transform_response(model, response)
|
||||
}
|
||||
}
|
||||
|
|
@ -268,7 +268,7 @@ mod tests {
|
|||
"https://env.services.ai.azure.com/anthropic/v1/messages"
|
||||
);
|
||||
let err = complete_azure_anthropic_url(Some(" "), &|_| None).expect_err("missing base");
|
||||
assert!(matches!(err, CoreError::Auth(_)));
|
||||
assert!(matches!(err, Error::Auth(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -284,7 +284,7 @@ mod tests {
|
|||
);
|
||||
assert!(matches!(
|
||||
resolve_azure_api_key(None, &|_| None).expect_err("missing key"),
|
||||
CoreError::Auth(_)
|
||||
Error::Auth(_)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::{Error, json_type_name};
|
||||
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
|
@ -32,17 +32,17 @@ fn resolve_value(
|
|||
env_name: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
missing_message: &str,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
non_empty(explicit)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(missing_message.to_string()))
|
||||
.ok_or_else(|| Error::Auth(missing_message.to_string()))
|
||||
}
|
||||
|
||||
pub fn resolve_azure_ai_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_AI_API_KEY_ENV,
|
||||
|
|
@ -54,7 +54,7 @@ pub fn resolve_azure_ai_api_key(
|
|||
pub fn resolve_azure_ai_api_base(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_AI_API_BASE_ENV,
|
||||
|
|
@ -66,7 +66,7 @@ pub fn resolve_azure_ai_api_base(
|
|||
pub fn complete_azure_ai_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let base = resolve_azure_ai_api_base(api_base, env_lookup)?;
|
||||
Ok(format!(
|
||||
"{}/providers/mistral/azure/ocr",
|
||||
|
|
@ -77,7 +77,7 @@ pub fn complete_azure_ai_url(
|
|||
pub fn resolve_document_intelligence_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV,
|
||||
|
|
@ -89,7 +89,7 @@ pub fn resolve_document_intelligence_api_key(
|
|||
pub fn resolve_document_intelligence_endpoint(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV,
|
||||
|
|
@ -127,7 +127,7 @@ fn pages_token_is_valid(token: &str) -> bool {
|
|||
}
|
||||
}
|
||||
|
||||
fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
||||
fn normalize_pages_param(pages: &Value) -> Result<Option<String>, Error> {
|
||||
match pages {
|
||||
Value::String(value) => {
|
||||
let normalized = value
|
||||
|
|
@ -138,7 +138,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
|||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
Ok(Some(normalized))
|
||||
} else {
|
||||
Err(CoreError::InvalidRequest(format!(
|
||||
Err(Error::InvalidRequest(format!(
|
||||
"Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'."
|
||||
)))
|
||||
}
|
||||
|
|
@ -152,7 +152,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
|||
for value in values {
|
||||
let page = value.as_i64().expect("checked is_i64");
|
||||
if page < 0 {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
return Err(Error::InvalidRequest(
|
||||
"`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(),
|
||||
));
|
||||
}
|
||||
|
|
@ -176,16 +176,16 @@ fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
|||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
return Ok(Some(normalized));
|
||||
}
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
return Err(Error::InvalidRequest(format!(
|
||||
"Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'."
|
||||
)));
|
||||
}
|
||||
Err(CoreError::InvalidRequest(
|
||||
Err(Error::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
))
|
||||
}
|
||||
_ => Err(CoreError::InvalidRequest(
|
||||
_ => Err(Error::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
)),
|
||||
|
|
@ -199,13 +199,13 @@ fn feature_token_is_valid(token: &str) -> bool {
|
|||
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
|
||||
}
|
||||
|
||||
fn invalid_features_error(features: &Value) -> CoreError {
|
||||
CoreError::InvalidRequest(format!(
|
||||
fn invalid_features_error(features: &Value) -> Error {
|
||||
Error::InvalidRequest(format!(
|
||||
"Invalid `features` for Azure Document Intelligence: {features:?}. Expected a list of feature names or a comma-separated string like 'keyValuePairs' or 'keyValuePairs,languages'."
|
||||
))
|
||||
}
|
||||
|
||||
fn normalize_features_param(features: &Value) -> CoreResult<Option<String>> {
|
||||
fn normalize_features_param(features: &Value) -> Result<Option<String>, Error> {
|
||||
let normalized = match features {
|
||||
Value::String(value) => value
|
||||
.split(',')
|
||||
|
|
@ -237,7 +237,7 @@ pub fn complete_document_intelligence_url(
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?;
|
||||
let mut url = format!(
|
||||
"{}/documentintelligence/documentModels/{}:analyze?api-version={}",
|
||||
|
|
@ -263,20 +263,20 @@ pub fn complete_document_intelligence_url(
|
|||
Ok(url)
|
||||
}
|
||||
|
||||
fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> {
|
||||
let object = document.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
.ok_or(Error::MissingField("document.type"))?;
|
||||
let field_name = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
"image_url" => "image_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
return Err(Error::InvalidRequest(format!(
|
||||
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
|
||||
)));
|
||||
}
|
||||
|
|
@ -285,7 +285,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
|
|||
.get(field_name)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(field_name))
|
||||
.ok_or(Error::MissingField(field_name))
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_uri(data_uri: &str) -> &str {
|
||||
|
|
@ -328,8 +328,8 @@ fn page_dimensions(page: &Map<String, Value>) -> Value {
|
|||
}
|
||||
|
||||
impl OcrProviderConfig for AzureAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
|
|
@ -337,7 +337,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
|
|
@ -345,7 +345,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
) -> Result<OcrResponseData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
|
|
@ -355,7 +355,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
complete_azure_ai_url(api_base, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -363,7 +363,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_azure_ai_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -373,7 +373,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
}
|
||||
|
||||
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
|
|
@ -382,7 +382,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
_model: &str,
|
||||
document: Value,
|
||||
_optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
let document_url = document_url_from_mistral_document(&document)?;
|
||||
let mut data = Map::new();
|
||||
if document_url.starts_with("data:") {
|
||||
|
|
@ -406,19 +406,19 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
) -> Result<OcrResponseData, Error> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
.ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
let status = response
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
.ok_or(Error::MissingField("status"))?;
|
||||
if status != "succeeded" {
|
||||
return Err(CoreError::InvalidResponse(format!(
|
||||
return Err(Error::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed with status: {status}"
|
||||
)));
|
||||
}
|
||||
|
|
@ -461,7 +461,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
complete_document_intelligence_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -469,7 +469,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_document_intelligence_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -600,7 +600,7 @@ mod tests {
|
|||
.expect_err("invalid features must fail");
|
||||
|
||||
assert!(
|
||||
matches!(error, CoreError::InvalidRequest(message) if message.contains("Invalid `features`")),
|
||||
matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `features`")),
|
||||
"features={features:?}"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use crate::audio_transcription::transformation::{
|
|||
use crate::audio_transcription::types::{
|
||||
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
|
||||
};
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::{Error, json_type_name};
|
||||
|
||||
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
|
||||
|
|
@ -18,8 +18,8 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
|
|||
|
||||
pub struct BedrockAudioTranscriptionConfig;
|
||||
|
||||
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
|
||||
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
fn audio_fields(audio: Value) -> Result<(String, String), Error> {
|
||||
let object = audio.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&audio),
|
||||
})?;
|
||||
|
|
@ -27,13 +27,13 @@ fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
|
|||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField("audio.data"))?;
|
||||
.ok_or(Error::MissingField("audio.data"))?;
|
||||
let format = object
|
||||
.get("format")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg"))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string())
|
||||
Error::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string())
|
||||
})?;
|
||||
Ok((data.to_string(), format.to_string()))
|
||||
}
|
||||
|
|
@ -55,7 +55,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
_model: &str,
|
||||
audio: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<AudioTranscriptionRequestData> {
|
||||
) -> Result<AudioTranscriptionRequestData, Error> {
|
||||
let (data, format) = audio_fields(audio)?;
|
||||
let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string();
|
||||
if let Some(language) = optional_string(&optional_params, "language") {
|
||||
|
|
@ -87,14 +87,14 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
&self,
|
||||
_model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<AudioTranscriptionResponseData> {
|
||||
) -> Result<AudioTranscriptionResponseData, Error> {
|
||||
let content = response_json
|
||||
.get("output")
|
||||
.and_then(|value| value.get("message"))
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(Value::as_array)
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("Bedrock response has no output content".to_string())
|
||||
Error::InvalidResponse("Bedrock response has no output content".to_string())
|
||||
})?;
|
||||
let mut text = String::new();
|
||||
for block in content {
|
||||
|
|
@ -111,7 +111,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let (model_id, model_region) = bedrock_model_id_and_region(model);
|
||||
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
|
||||
let endpoint = optional_params
|
||||
|
|
@ -133,7 +133,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<AudioTranscriptionAuth> {
|
||||
) -> Result<AudioTranscriptionAuth, Error> {
|
||||
let (_, model_region) = bedrock_model_id_and_region(model);
|
||||
Ok(AudioTranscriptionAuth::AwsSigV4 {
|
||||
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ use std::time::Duration;
|
|||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::caching::in_memory_cache::InMemoryCache;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
use aws_sigv4::http_request::{
|
||||
|
|
@ -197,7 +197,7 @@ pub fn classify_auth(
|
|||
pub async fn resolve_credentials(
|
||||
config: AwsAuthConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> CoreResult<Credentials> {
|
||||
) -> Result<Credentials, Error> {
|
||||
let resolved = config.clone().with_environment(env_lookup);
|
||||
let flow = classify_auth(config, env_lookup);
|
||||
match flow {
|
||||
|
|
@ -244,9 +244,10 @@ pub async fn resolve_credentials(
|
|||
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
|
||||
.profile_name(name)
|
||||
.build();
|
||||
provider.provide_credentials().await.map_err(|error| {
|
||||
CoreError::Auth(format!("AWS profile credentials failed: {error}"))
|
||||
})
|
||||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::Auth(format!("AWS profile credentials failed: {error}")))
|
||||
}
|
||||
AwsAuthFlow::AssumeRole { role, session_name } => {
|
||||
if is_already_running_as_role(&role, &resolved).await? {
|
||||
|
|
@ -260,7 +261,7 @@ pub async fn resolve_credentials(
|
|||
.build()
|
||||
.await;
|
||||
let credentials = provider.provide_credentials().await.map_err(|error| {
|
||||
CoreError::Auth(format!("AWS default credentials failed: {error}"))
|
||||
Error::Auth(format!("AWS default credentials failed: {error}"))
|
||||
})?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
|
|
@ -301,7 +302,7 @@ pub async fn resolve_credentials(
|
|||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}")))
|
||||
.map_err(|error| Error::Auth(format!("AWS role credentials failed: {error}")))
|
||||
}
|
||||
AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
|
|
@ -325,13 +326,13 @@ pub async fn resolve_credentials(
|
|||
.send()
|
||||
.await
|
||||
.map_err(|error| {
|
||||
CoreError::Auth(format!("AWS web identity credentials failed: {error}"))
|
||||
Error::Auth(format!("AWS web identity credentials failed: {error}"))
|
||||
})?;
|
||||
let credentials = response.credentials().ok_or_else(|| {
|
||||
CoreError::Auth("AWS web identity response had no credentials".to_string())
|
||||
Error::Auth("AWS web identity response had no credentials".to_string())
|
||||
})?;
|
||||
let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| {
|
||||
CoreError::Auth(format!("AWS web identity expiration was invalid: {error}"))
|
||||
Error::Auth(format!("AWS web identity expiration was invalid: {error}"))
|
||||
})?;
|
||||
Ok(Credentials::new(
|
||||
credentials.access_key_id(),
|
||||
|
|
@ -350,9 +351,10 @@ pub async fn resolve_credentials(
|
|||
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
|
||||
.build()
|
||||
.await;
|
||||
let credentials = provider.provide_credentials().await.map_err(|error| {
|
||||
CoreError::Auth(format!("AWS default credentials failed: {error}"))
|
||||
})?;
|
||||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::Auth(format!("AWS default credentials failed: {error}")))?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
|
|
@ -363,7 +365,7 @@ pub async fn resolve_credentials(
|
|||
}
|
||||
}
|
||||
|
||||
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult<bool> {
|
||||
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result<bool, Error> {
|
||||
if role_identity(role).is_none() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
|
@ -437,7 +439,7 @@ pub fn sign_bedrock_post(
|
|||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> CoreResult<BTreeMap<String, String>> {
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let identity: Identity = credentials.clone().into();
|
||||
let params = v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
|
|
@ -447,14 +449,14 @@ pub fn sign_bedrock_post(
|
|||
.settings(SigningSettings::default())
|
||||
.build()
|
||||
.map(SigningParams::from)
|
||||
.map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?;
|
||||
.map_err(|error| Error::Auth(format!("AWS signing parameters failed: {error}")))?;
|
||||
let header_refs = headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()));
|
||||
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
|
||||
.map_err(|error| CoreError::Auth(format!("AWS signable request failed: {error}")))?;
|
||||
.map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?;
|
||||
let (instructions, _) = sign(request, ¶ms)
|
||||
.map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))?
|
||||
.map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))?
|
||||
.into_parts();
|
||||
Ok(instructions
|
||||
.headers()
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use super::*;
|
||||
use crate::Error;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
|
|
@ -23,7 +24,7 @@ fn transform(msgs: Value, opts: Value) -> Value {
|
|||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
fn transform_response(body: Value) -> Result<ChatCompletionsResponse, Error> {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response(
|
||||
"anthropic.claude-sonnet-4-5-v1:0",
|
||||
ProviderChatResponseData { body },
|
||||
|
|
@ -478,25 +479,22 @@ fn declines_a_response_carrying_a_tool_use_block() {
|
|||
"usage": {"inputTokens": 1, "outputTokens": 1}
|
||||
}))
|
||||
.expect_err("tool use block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
assert_eq!(err, Error::Unsupported("non-text response content block"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("converse response is not an object".to_string())
|
||||
Error::InvalidResponse("converse response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"usage": {}})).expect_err("no output"),
|
||||
CoreError::MissingField("output.message.content")
|
||||
Error::MissingField("output.message.content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
Error::MissingField("usage")
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ use crate::chat_completions::types::{
|
|||
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::error::Error;
|
||||
|
||||
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
|
||||
|
|
@ -110,7 +110,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let (model_id, model_region) = bedrock_model_id_and_region(model);
|
||||
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
|
||||
let endpoint = optional_params
|
||||
|
|
@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
) -> Result<ChatCompletionsAuth, Error> {
|
||||
// Python reads `api_key` as the Bedrock bearer token and consults the
|
||||
// env only when the caller passed none, so a caller-supplied empty key
|
||||
// falls through to SigV4 without reaching for the environment. An
|
||||
|
|
@ -208,7 +208,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
_model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
) -> Result<ProviderChatRequestData, Error> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: converse_body(&build_conversation(&messages), &optional_params),
|
||||
})
|
||||
|
|
@ -218,17 +218,18 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("converse response is not an object".into())
|
||||
})?;
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let body = response
|
||||
.body
|
||||
.as_object()
|
||||
.ok_or_else(|| Error::InvalidResponse("converse response is not an object".into()))?;
|
||||
|
||||
let content = body
|
||||
.get("output")
|
||||
.and_then(|output| output.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("output.message.content"))?;
|
||||
.ok_or(Error::MissingField("output.message.content"))?;
|
||||
// The route declines tool requests, so anything other than a text block
|
||||
// is something this path never asked for. Decline; the host falls back.
|
||||
if content.iter().any(|block| {
|
||||
|
|
@ -236,7 +237,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
.as_object()
|
||||
.is_none_or(|block| block.len() != 1 || !block.contains_key("text"))
|
||||
}) {
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
return Err(Error::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
|
|
@ -246,7 +247,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
|||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
.ok_or(Error::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
let computed = usage_from_parts(
|
||||
field("inputTokens"),
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::{Error, json_type_name};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -47,7 +47,7 @@ pub fn complete_url(api_base: Option<&str>) -> String {
|
|||
|
||||
/// Resolve the Mistral API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth`
|
||||
/// Blank/whitespace values are treated as absent. Returns `Error::Auth`
|
||||
/// when no usable key is available.
|
||||
///
|
||||
/// Note: the env fallback only reads the process environment. Secret-manager
|
||||
|
|
@ -56,13 +56,13 @@ pub fn complete_url(api_base: Option<&str>) -> String {
|
|||
pub fn resolve_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
pub struct MistralOcrConfig;
|
||||
|
|
@ -70,18 +70,20 @@ pub struct MistralOcrConfig;
|
|||
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
|
||||
|
||||
impl OcrProviderConfig for MistralOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
if !document.is_object() {
|
||||
return Err(CoreError::InvalidType {
|
||||
return Err(Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&document),
|
||||
});
|
||||
|
|
@ -100,14 +102,15 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
})
|
||||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
) -> Result<OcrResponseData, Error> {
|
||||
let response_object = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
.ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
|
|
@ -140,7 +143,7 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
Ok(complete_url(api_base))
|
||||
}
|
||||
|
||||
|
|
@ -148,13 +151,13 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn supported_ocr_params() -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
pub fn get_supported_ocr_params() -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
|
||||
}
|
||||
|
||||
pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
|
|
@ -165,11 +168,11 @@ pub fn transform_ocr_request(
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult<OcrResponseData> {
|
||||
pub fn transform_ocr_response(model: &str, response_json: Value) -> Result<OcrResponseData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
|
|
@ -181,7 +184,7 @@ mod tests {
|
|||
#[test]
|
||||
fn supported_params_match_python_mistral_ocr_config() {
|
||||
assert_eq!(
|
||||
supported_ocr_params(),
|
||||
get_supported_ocr_params(),
|
||||
&[
|
||||
"pages",
|
||||
"include_image_base64",
|
||||
|
|
@ -250,7 +253,7 @@ mod tests {
|
|||
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidType {
|
||||
Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: "string",
|
||||
}
|
||||
|
|
@ -307,6 +310,6 @@ mod tests {
|
|||
#[test]
|
||||
fn resolve_api_key_errors_when_absent() {
|
||||
let err = resolve_api_key(None, &|_| None).expect_err("missing key should error");
|
||||
assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string()));
|
||||
assert_eq!(err, Error::Auth(MISSING_KEY_MESSAGE.to_string()));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::realtime::transformation::RealtimeProviderConfig;
|
||||
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
|
||||
|
|
@ -72,7 +72,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig {
|
|||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
) -> Result<RealtimeTransformResult, Error> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
|
||||
|
|
@ -80,7 +80,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig {
|
|||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
) -> Result<RealtimeTransformResult, Error> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
}
|
||||
|
|
@ -88,14 +88,14 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig {
|
|||
pub fn transform_realtime_request(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
) -> Result<RealtimeTransformResult, Error> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model)
|
||||
}
|
||||
|
||||
pub fn transform_realtime_response(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
) -> Result<RealtimeTransformResult, Error> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult};
|
||||
use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model};
|
||||
|
||||
|
|
@ -15,7 +15,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig {
|
|||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<ResponsesWsTransformResult> {
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(enforce_model(
|
||||
event, model,
|
||||
)))
|
||||
|
|
@ -25,7 +25,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig {
|
|||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<ResponsesWsTransformResult> {
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::{Error, json_type_name};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
|
@ -43,7 +43,7 @@ pub fn is_deepseek_model(model: &str) -> bool {
|
|||
pub fn resolve_vertex_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
|
|
@ -51,7 +51,7 @@ pub fn resolve_vertex_api_key(
|
|||
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
Error::Auth(
|
||||
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
|
||||
.to_string(),
|
||||
)
|
||||
|
|
@ -61,12 +61,12 @@ pub fn resolve_vertex_api_key(
|
|||
fn vertex_project(
|
||||
params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
string_param(params, &["vertex_project", "vertex_ai_project"])
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(
|
||||
Error::InvalidRequest(
|
||||
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
|
||||
.to_string(),
|
||||
)
|
||||
|
|
@ -99,7 +99,7 @@ pub fn complete_vertex_mistral_url(
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = vertex_mistral_api_base(api_base, &location);
|
||||
|
|
@ -112,7 +112,7 @@ pub fn complete_vertex_deepseek_url(
|
|||
api_base: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = api_base
|
||||
|
|
@ -125,20 +125,20 @@ pub fn complete_vertex_deepseek_url(
|
|||
))
|
||||
}
|
||||
|
||||
fn document_content_item(document: &Value) -> CoreResult<Value> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
fn document_content_item(document: &Value) -> Result<Value, Error> {
|
||||
let object = document.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
.ok_or(Error::MissingField("document.type"))?;
|
||||
let url_field = match doc_type {
|
||||
"image_url" => "image_url",
|
||||
"document_url" => "document_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
return Err(Error::InvalidRequest(format!(
|
||||
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
|
||||
)));
|
||||
}
|
||||
|
|
@ -147,7 +147,7 @@ fn document_content_item(document: &Value) -> CoreResult<Value> {
|
|||
.get(url_field)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(url_field))?;
|
||||
.ok_or(Error::MissingField(url_field))?;
|
||||
|
||||
Ok(json!({
|
||||
"type": "image_url",
|
||||
|
|
@ -163,7 +163,7 @@ fn deepseek_model_name(model: &str) -> String {
|
|||
}
|
||||
}
|
||||
|
||||
fn first_choice_content(response: &Value) -> CoreResult<Value> {
|
||||
fn first_choice_content(response: &Value) -> Result<Value, Error> {
|
||||
response
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
|
|
@ -176,9 +176,7 @@ fn first_choice_content(response: &Value) -> CoreResult<Value> {
|
|||
Value::Object(_) => true,
|
||||
_ => false,
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string())
|
||||
})
|
||||
.ok_or_else(|| Error::InvalidResponse("No content in DeepSeek OCR response".to_string()))
|
||||
}
|
||||
|
||||
fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> Value {
|
||||
|
|
@ -210,8 +208,8 @@ fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> V
|
|||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
|
|
@ -219,7 +217,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
|
|
@ -227,7 +225,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
) -> Result<OcrResponseData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
|
|
@ -237,7 +235,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
complete_vertex_mistral_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -245,7 +243,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -255,7 +253,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
DEEPSEEK_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
|
|
@ -264,7 +262,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
"model".to_string(),
|
||||
|
|
@ -289,10 +287,10 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
) -> Result<OcrResponseData, Error> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
.ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
|
|
@ -314,7 +312,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
});
|
||||
}
|
||||
|
||||
let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
let object = ocr_data.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&ocr_data),
|
||||
})?;
|
||||
|
|
@ -346,7 +344,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
_model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
|
|
@ -354,7 +352,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
) -> Result<String, Error> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
|
||||
pub trait RealtimeProviderConfig {
|
||||
|
|
@ -11,12 +11,12 @@ pub trait RealtimeProviderConfig {
|
|||
&self,
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult>;
|
||||
) -> Result<RealtimeTransformResult, Error>;
|
||||
|
||||
/// Transform a backend → client event before it is forwarded downstream.
|
||||
fn transform_realtime_response(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult>;
|
||||
) -> Result<RealtimeTransformResult, Error>;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,9 +5,9 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
|||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::Error;
|
||||
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType};
|
||||
use crate::{CoreError, CoreResult};
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ResponsesWsUsage {
|
||||
|
|
@ -205,7 +205,7 @@ impl ResponsesWsInstrumentation {
|
|||
}
|
||||
}
|
||||
|
||||
type LifecycleFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
|
||||
type LifecycleFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation {
|
||||
type PreCallFuture<'a> = LifecycleFuture<'a, ()>;
|
||||
|
|
@ -246,7 +246,7 @@ impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation {
|
|||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_error: &'a CoreError,
|
||||
_error: &'a Error,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
|
|
@ -342,7 +342,7 @@ mod tests {
|
|||
),
|
||||
(),
|
||||
&instrumentation,
|
||||
|_| async { Ok::<(), CoreError>(()) },
|
||||
|_| async { Ok::<(), Error>(()) },
|
||||
)
|
||||
.await;
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,81 @@
|
|||
pub mod instrumentation;
|
||||
pub mod types;
|
||||
pub mod websocket;
|
||||
|
||||
use crate::error::Error;
|
||||
use crate::streaming::OpenedStream;
|
||||
use types::{ResponsesStreamEvent, ResponsesStreamRequest, ResponsesWebSocketRequest};
|
||||
use websocket::TypedResponsesWebSocketSession;
|
||||
|
||||
pub async fn responses_stream(
|
||||
_request: ResponsesStreamRequest,
|
||||
) -> Result<OpenedStream<ResponsesStreamEvent>, Error> {
|
||||
Err(Error::Unsupported(
|
||||
"responses HTTP streaming provider registration",
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn responses_websocket(
|
||||
_request: ResponsesWebSocketRequest,
|
||||
) -> Result<Box<dyn TypedResponsesWebSocketSession>, Error> {
|
||||
Err(Error::Unsupported(
|
||||
"responses WebSocket streaming provider registration",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod stream_entrypoint_tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::streaming::{
|
||||
ProviderCredentials, StreamProviderId, StreamTarget, StreamTransportOptions,
|
||||
};
|
||||
|
||||
fn target() -> StreamTarget {
|
||||
StreamTarget::new(
|
||||
StreamProviderId::OpenAi,
|
||||
ProviderCredentials::default(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_http_stream_declines_until_a_provider_is_registered() {
|
||||
let body = serde_json::from_value(json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
"stream": true
|
||||
}))
|
||||
.expect("valid Responses stream request");
|
||||
let result = responses_stream(ResponsesStreamRequest {
|
||||
body,
|
||||
target: target(),
|
||||
transport: StreamTransportOptions::default(),
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Unsupported(
|
||||
"responses HTTP streaming provider registration"
|
||||
))
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_websocket_declines_until_a_provider_is_registered() {
|
||||
let result = responses_websocket(ResponsesWebSocketRequest {
|
||||
target: target(),
|
||||
transport: StreamTransportOptions::default(),
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Unsupported(
|
||||
"responses WebSocket streaming provider registration"
|
||||
))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,13 +1,46 @@
|
|||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::streaming::{JsonObject, StreamTarget, StreamTransportOptions};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ResponsesWsEventType {
|
||||
ResponseCreate,
|
||||
ResponseCreated,
|
||||
ResponseInProgress,
|
||||
ResponseReasoningSummaryPartAdded,
|
||||
ResponseReasoningSummaryTextDelta,
|
||||
ResponseReasoningSummaryTextDone,
|
||||
ResponseReasoningSummaryPartDone,
|
||||
ResponseOutputItemAdded,
|
||||
ResponseOutputTextDelta,
|
||||
ResponseOutputTextAnnotationAdded,
|
||||
ResponseOutputTextDone,
|
||||
ResponseRefusalDelta,
|
||||
ResponseRefusalDone,
|
||||
ResponseFunctionCallArgumentsDelta,
|
||||
ResponseFunctionCallArgumentsDone,
|
||||
ResponseFileSearchCallInProgress,
|
||||
ResponseFileSearchCallSearching,
|
||||
ResponseFileSearchCallCompleted,
|
||||
ResponseWebSearchCallInProgress,
|
||||
ResponseWebSearchCallSearching,
|
||||
ResponseWebSearchCallCompleted,
|
||||
ResponseMcpListToolsInProgress,
|
||||
ResponseMcpListToolsCompleted,
|
||||
ResponseMcpListToolsFailed,
|
||||
ResponseMcpCallInProgress,
|
||||
ResponseMcpCallArgumentsDelta,
|
||||
ResponseMcpCallArgumentsDone,
|
||||
ResponseMcpCallCompleted,
|
||||
ResponseMcpCallFailed,
|
||||
ResponseContentPartAdded,
|
||||
ResponseContentPartDone,
|
||||
ResponseOutputItemDone,
|
||||
ResponseCompleted,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
ImageGenerationPartialImage,
|
||||
Error,
|
||||
Other(String),
|
||||
}
|
||||
|
|
@ -17,9 +50,40 @@ impl ResponsesWsEventType {
|
|||
match self {
|
||||
Self::ResponseCreate => "response.create",
|
||||
Self::ResponseCreated => "response.created",
|
||||
Self::ResponseInProgress => "response.in_progress",
|
||||
Self::ResponseReasoningSummaryPartAdded => "response.reasoning_summary_part.added",
|
||||
Self::ResponseReasoningSummaryTextDelta => "response.reasoning_summary_text.delta",
|
||||
Self::ResponseReasoningSummaryTextDone => "response.reasoning_summary_text.done",
|
||||
Self::ResponseReasoningSummaryPartDone => "response.reasoning_summary_part.done",
|
||||
Self::ResponseOutputItemAdded => "response.output_item.added",
|
||||
Self::ResponseOutputTextDelta => "response.output_text.delta",
|
||||
Self::ResponseOutputTextAnnotationAdded => "response.output_text.annotation.added",
|
||||
Self::ResponseOutputTextDone => "response.output_text.done",
|
||||
Self::ResponseRefusalDelta => "response.refusal.delta",
|
||||
Self::ResponseRefusalDone => "response.refusal.done",
|
||||
Self::ResponseFunctionCallArgumentsDelta => "response.function_call_arguments.delta",
|
||||
Self::ResponseFunctionCallArgumentsDone => "response.function_call_arguments.done",
|
||||
Self::ResponseFileSearchCallInProgress => "response.file_search_call.in_progress",
|
||||
Self::ResponseFileSearchCallSearching => "response.file_search_call.searching",
|
||||
Self::ResponseFileSearchCallCompleted => "response.file_search_call.completed",
|
||||
Self::ResponseWebSearchCallInProgress => "response.web_search_call.in_progress",
|
||||
Self::ResponseWebSearchCallSearching => "response.web_search_call.searching",
|
||||
Self::ResponseWebSearchCallCompleted => "response.web_search_call.completed",
|
||||
Self::ResponseMcpListToolsInProgress => "response.mcp_list_tools.in_progress",
|
||||
Self::ResponseMcpListToolsCompleted => "response.mcp_list_tools.completed",
|
||||
Self::ResponseMcpListToolsFailed => "response.mcp_list_tools.failed",
|
||||
Self::ResponseMcpCallInProgress => "response.mcp_call.in_progress",
|
||||
Self::ResponseMcpCallArgumentsDelta => "response.mcp_call_arguments.delta",
|
||||
Self::ResponseMcpCallArgumentsDone => "response.mcp_call_arguments.done",
|
||||
Self::ResponseMcpCallCompleted => "response.mcp_call.completed",
|
||||
Self::ResponseMcpCallFailed => "response.mcp_call.failed",
|
||||
Self::ResponseContentPartAdded => "response.content_part.added",
|
||||
Self::ResponseContentPartDone => "response.content_part.done",
|
||||
Self::ResponseOutputItemDone => "response.output_item.done",
|
||||
Self::ResponseCompleted => "response.completed",
|
||||
Self::ResponseFailed => "response.failed",
|
||||
Self::ResponseIncomplete => "response.incomplete",
|
||||
Self::ImageGenerationPartialImage => "image_generation.partial_image",
|
||||
Self::Error => "error",
|
||||
Self::Other(value) => value,
|
||||
}
|
||||
|
|
@ -44,9 +108,40 @@ impl<'de> Deserialize<'de> for ResponsesWsEventType {
|
|||
Ok(match value.as_str() {
|
||||
"response.create" => Self::ResponseCreate,
|
||||
"response.created" => Self::ResponseCreated,
|
||||
"response.in_progress" => Self::ResponseInProgress,
|
||||
"response.reasoning_summary_part.added" => Self::ResponseReasoningSummaryPartAdded,
|
||||
"response.reasoning_summary_text.delta" => Self::ResponseReasoningSummaryTextDelta,
|
||||
"response.reasoning_summary_text.done" => Self::ResponseReasoningSummaryTextDone,
|
||||
"response.reasoning_summary_part.done" => Self::ResponseReasoningSummaryPartDone,
|
||||
"response.output_item.added" => Self::ResponseOutputItemAdded,
|
||||
"response.output_text.delta" => Self::ResponseOutputTextDelta,
|
||||
"response.output_text.annotation.added" => Self::ResponseOutputTextAnnotationAdded,
|
||||
"response.output_text.done" => Self::ResponseOutputTextDone,
|
||||
"response.refusal.delta" => Self::ResponseRefusalDelta,
|
||||
"response.refusal.done" => Self::ResponseRefusalDone,
|
||||
"response.function_call_arguments.delta" => Self::ResponseFunctionCallArgumentsDelta,
|
||||
"response.function_call_arguments.done" => Self::ResponseFunctionCallArgumentsDone,
|
||||
"response.file_search_call.in_progress" => Self::ResponseFileSearchCallInProgress,
|
||||
"response.file_search_call.searching" => Self::ResponseFileSearchCallSearching,
|
||||
"response.file_search_call.completed" => Self::ResponseFileSearchCallCompleted,
|
||||
"response.web_search_call.in_progress" => Self::ResponseWebSearchCallInProgress,
|
||||
"response.web_search_call.searching" => Self::ResponseWebSearchCallSearching,
|
||||
"response.web_search_call.completed" => Self::ResponseWebSearchCallCompleted,
|
||||
"response.mcp_list_tools.in_progress" => Self::ResponseMcpListToolsInProgress,
|
||||
"response.mcp_list_tools.completed" => Self::ResponseMcpListToolsCompleted,
|
||||
"response.mcp_list_tools.failed" => Self::ResponseMcpListToolsFailed,
|
||||
"response.mcp_call.in_progress" => Self::ResponseMcpCallInProgress,
|
||||
"response.mcp_call_arguments.delta" => Self::ResponseMcpCallArgumentsDelta,
|
||||
"response.mcp_call_arguments.done" => Self::ResponseMcpCallArgumentsDone,
|
||||
"response.mcp_call.completed" => Self::ResponseMcpCallCompleted,
|
||||
"response.mcp_call.failed" => Self::ResponseMcpCallFailed,
|
||||
"response.content_part.added" => Self::ResponseContentPartAdded,
|
||||
"response.content_part.done" => Self::ResponseContentPartDone,
|
||||
"response.output_item.done" => Self::ResponseOutputItemDone,
|
||||
"response.completed" => Self::ResponseCompleted,
|
||||
"response.failed" => Self::ResponseFailed,
|
||||
"response.incomplete" => Self::ResponseIncomplete,
|
||||
"image_generation.partial_image" => Self::ImageGenerationPartialImage,
|
||||
"error" => Self::Error,
|
||||
_ => Self::Other(value),
|
||||
})
|
||||
|
|
@ -84,6 +179,60 @@ pub struct ResponsesWsTransformResult {
|
|||
pub events: Vec<ResponsesWsEvent>,
|
||||
}
|
||||
|
||||
pub type ResponsesStreamEvent = ResponsesWsEvent;
|
||||
pub type ResponseCommand = ResponsesWsEvent;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesInput {
|
||||
Text(String),
|
||||
Items(Vec<JsonObject>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesToolChoice {
|
||||
Name(String),
|
||||
Definition(JsonObject),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ResponsesStreamRequestBody {
|
||||
pub model: String,
|
||||
pub input: ResponsesInput,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub instructions: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_output_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub previous_response_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub store: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<JsonObject>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<ResponsesToolChoice>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub metadata: Option<JsonObject>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub include: Option<Vec<String>>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub struct ResponsesStreamRequest {
|
||||
pub body: ResponsesStreamRequestBody,
|
||||
pub target: StreamTarget,
|
||||
pub transport: StreamTransportOptions,
|
||||
}
|
||||
|
||||
pub struct ResponsesWebSocketRequest {
|
||||
pub target: StreamTarget,
|
||||
pub transport: StreamTransportOptions,
|
||||
}
|
||||
|
||||
impl ResponsesWsTransformResult {
|
||||
pub fn passthrough(event: ResponsesWsEvent) -> Self {
|
||||
Self {
|
||||
|
|
@ -127,11 +276,14 @@ mod tests {
|
|||
let known: ResponsesWsEventType =
|
||||
serde_json::from_str("\"response.completed\"").expect("valid event type");
|
||||
assert_eq!(known, ResponsesWsEventType::ResponseCompleted);
|
||||
let unknown: ResponsesWsEventType =
|
||||
let output_delta: ResponsesWsEventType =
|
||||
serde_json::from_str("\"response.output_text.delta\"").expect("valid event type");
|
||||
assert_eq!(output_delta, ResponsesWsEventType::ResponseOutputTextDelta);
|
||||
let unknown: ResponsesWsEventType =
|
||||
serde_json::from_str("\"response.future_event\"").expect("valid event type");
|
||||
assert_eq!(
|
||||
unknown,
|
||||
ResponsesWsEventType::Other("response.output_text.delta".to_string())
|
||||
ResponsesWsEventType::Other("response.future_event".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -163,4 +315,42 @@ mod tests {
|
|||
assert_eq!(flat.model(), Some("gpt-5"));
|
||||
assert_eq!(nested.model(), Some("gpt-5-mini"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_request_deserializes_the_public_responses_shape() {
|
||||
let request: ResponsesStreamRequestBody = serde_json::from_value(serde_json::json!({
|
||||
"model": "gpt-5",
|
||||
"input": "hello",
|
||||
"stream": true,
|
||||
"max_output_tokens": 32
|
||||
}))
|
||||
.expect("public request shape");
|
||||
|
||||
assert_eq!(request.model, "gpt-5");
|
||||
assert_eq!(request.stream, Some(true));
|
||||
assert!(matches!(request.input, ResponsesInput::Text(ref text) if text == "hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_stream_events_round_trip_for_forward_compatibility() {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(serde_json::json!({
|
||||
"type": "response.future_event",
|
||||
"sequence_number": 7,
|
||||
"future_field": "value"
|
||||
}))
|
||||
.expect("unknown event");
|
||||
|
||||
assert_eq!(
|
||||
event.event_type,
|
||||
ResponsesWsEventType::Other("response.future_event".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(event).expect("serializable event"),
|
||||
serde_json::json!({
|
||||
"type": "response.future_event",
|
||||
"sequence_number": 7,
|
||||
"future_field": "value"
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,18 @@
|
|||
use crate::CoreResult;
|
||||
use crate::Error;
|
||||
use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH};
|
||||
use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult};
|
||||
use crate::responses::types::{
|
||||
ResponseCommand, ResponsesStreamEvent, ResponsesWsEvent, ResponsesWsEventType,
|
||||
ResponsesWsTransformResult,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
|
||||
pub trait TypedResponsesWebSocketSession: Send + Sync {
|
||||
fn send(&self, command: ResponseCommand) -> BoxFuture<'_, Result<(), Error>>;
|
||||
|
||||
fn recv(&self) -> BoxFuture<'_, Result<Option<ResponsesStreamEvent>, Error>>;
|
||||
|
||||
fn close(&self) -> BoxFuture<'_, Result<(), Error>>;
|
||||
}
|
||||
|
||||
pub trait ResponsesWebSocketProviderConfig: Sync {
|
||||
fn supports_native_websocket(&self) -> bool {
|
||||
|
|
@ -19,13 +31,13 @@ pub trait ResponsesWebSocketProviderConfig: Sync {
|
|||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<ResponsesWsTransformResult>;
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
|
||||
fn transform_ws_response(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<ResponsesWsTransformResult>;
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
}
|
||||
|
||||
pub fn complete_websocket_url(
|
||||
|
|
|
|||
636
litellm-rust/crates/core/src/streaming.rs
Normal file
636
litellm-rust/crates/core/src/streaming.rs
Normal file
|
|
@ -0,0 +1,636 @@
|
|||
use std::collections::VecDeque;
|
||||
use std::marker::PhantomData;
|
||||
use std::pin::Pin;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::error::Error;
|
||||
use futures_util::future::BoxFuture;
|
||||
use futures_util::{Stream, StreamExt, stream};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub type EventStream<E> = Pin<Box<dyn Stream<Item = Result<E, Error>> + Send + 'static>>;
|
||||
pub type ProviderChunkStream =
|
||||
Pin<Box<dyn Stream<Item = Result<ProviderStreamChunk, Error>> + Send + 'static>>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StreamTransport {
|
||||
Http,
|
||||
WebSocket,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StreamProviderId {
|
||||
Anthropic,
|
||||
AzureAi,
|
||||
BedrockConverse,
|
||||
OpenAi,
|
||||
}
|
||||
|
||||
impl TryFrom<&str> for StreamProviderId {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(value: &str) -> Result<Self, Self::Error> {
|
||||
match value {
|
||||
"anthropic" => Ok(Self::Anthropic),
|
||||
"azure_ai" => Ok(Self::AzureAi),
|
||||
"bedrock" | "bedrock_converse" => Ok(Self::BedrockConverse),
|
||||
"openai" => Ok(Self::OpenAi),
|
||||
_ => Err(Error::InvalidProvider(value.to_string())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct JsonObject(pub Map<String, Value>);
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Header {
|
||||
pub name: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, PartialEq)]
|
||||
pub struct ProviderCredentials {
|
||||
api_key: Option<String>,
|
||||
aws_access_key_id: Option<String>,
|
||||
aws_secret_access_key: Option<String>,
|
||||
aws_session_token: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderCredentials {
|
||||
pub fn new(
|
||||
api_key: Option<String>,
|
||||
aws_access_key_id: Option<String>,
|
||||
aws_secret_access_key: Option<String>,
|
||||
aws_session_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
api_key,
|
||||
aws_access_key_id,
|
||||
aws_secret_access_key,
|
||||
aws_session_token,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn api_key(&self) -> Option<&str> {
|
||||
self.api_key.as_deref()
|
||||
}
|
||||
|
||||
pub fn aws_access_key_id(&self) -> Option<&str> {
|
||||
self.aws_access_key_id.as_deref()
|
||||
}
|
||||
|
||||
pub fn aws_secret_access_key(&self) -> Option<&str> {
|
||||
self.aws_secret_access_key.as_deref()
|
||||
}
|
||||
|
||||
pub fn aws_session_token(&self) -> Option<&str> {
|
||||
self.aws_session_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// ```compile_fail
|
||||
/// fn assert_serialize<T: serde::Serialize>() {}
|
||||
/// assert_serialize::<litellm_core::streaming::ProviderCredentials>();
|
||||
/// assert_serialize::<litellm_core::streaming::StreamTarget>();
|
||||
/// assert_serialize::<litellm_core::streaming::StreamTransportOptions>();
|
||||
/// ```
|
||||
///
|
||||
/// ```compile_fail
|
||||
/// use litellm_core::streaming::{JsonObject, ProviderCredentials, StreamProviderId, StreamTarget};
|
||||
/// let mut target = StreamTarget::new(
|
||||
/// StreamProviderId::OpenAi,
|
||||
/// ProviderCredentials::default(),
|
||||
/// None,
|
||||
/// );
|
||||
/// target.metadata = JsonObject::default();
|
||||
/// ```
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct StreamTarget {
|
||||
provider: StreamProviderId,
|
||||
credentials: ProviderCredentials,
|
||||
api_base: Option<String>,
|
||||
}
|
||||
|
||||
impl StreamTarget {
|
||||
pub fn new(
|
||||
provider: StreamProviderId,
|
||||
credentials: ProviderCredentials,
|
||||
api_base: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
credentials,
|
||||
api_base,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn provider(&self) -> StreamProviderId {
|
||||
self.provider
|
||||
}
|
||||
|
||||
pub fn credentials(&self) -> &ProviderCredentials {
|
||||
&self.credentials
|
||||
}
|
||||
|
||||
pub fn api_base(&self) -> Option<&str> {
|
||||
self.api_base.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, PartialEq)]
|
||||
pub struct StreamTransportOptions {
|
||||
forwarded_headers: Vec<Header>,
|
||||
timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl StreamTransportOptions {
|
||||
pub fn new(forwarded_headers: Vec<Header>, timeout: Option<Duration>) -> Self {
|
||||
Self {
|
||||
forwarded_headers,
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forwarded_headers(&self) -> &[Header] {
|
||||
&self.forwarded_headers
|
||||
}
|
||||
|
||||
pub fn timeout(&self) -> Option<Duration> {
|
||||
self.timeout
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct StreamMetadata {
|
||||
pub status_code: u16,
|
||||
pub provider: StreamProviderId,
|
||||
pub transport: StreamTransport,
|
||||
pub response_headers: Vec<Header>,
|
||||
}
|
||||
|
||||
pub struct OpenedStream<E> {
|
||||
pub metadata: StreamMetadata,
|
||||
pub events: EventStream<E>,
|
||||
}
|
||||
|
||||
pub struct OpenedWireStream {
|
||||
pub metadata: StreamMetadata,
|
||||
pub chunks: ProviderChunkStream,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ProviderStreamChunk(Vec<u8>);
|
||||
|
||||
impl ProviderStreamChunk {
|
||||
pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
|
||||
Self(bytes.into())
|
||||
}
|
||||
|
||||
pub fn as_bytes(&self) -> &[u8] {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
pub trait StreamDecoder: Send + 'static {
|
||||
type WireEvent: Send + 'static;
|
||||
|
||||
fn push(&mut self, chunk: ProviderStreamChunk) -> Result<Vec<Self::WireEvent>, Error>;
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<Self::WireEvent>, Error> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait StreamProvider<R, E>: Send + Sync + 'static {
|
||||
type PreparedRequest: Send + 'static;
|
||||
type WireEvent: Send + 'static;
|
||||
type Decoder: StreamDecoder<WireEvent = Self::WireEvent>;
|
||||
|
||||
fn transform_request(&self, request: R) -> Result<Self::PreparedRequest, Error>;
|
||||
|
||||
fn call(
|
||||
&'static self,
|
||||
request: Self::PreparedRequest,
|
||||
) -> BoxFuture<'static, Result<OpenedWireStream, Error>>;
|
||||
|
||||
fn decoder(&self) -> Self::Decoder;
|
||||
|
||||
fn normalize(&self, event: Self::WireEvent) -> Result<Vec<E>, Error>;
|
||||
}
|
||||
|
||||
struct PipelineState<P: 'static, D, E, R>
|
||||
where
|
||||
D: StreamDecoder,
|
||||
{
|
||||
provider: &'static P,
|
||||
decoder: D,
|
||||
chunks: ProviderChunkStream,
|
||||
pending: VecDeque<Result<E, Error>>,
|
||||
finished: bool,
|
||||
request: PhantomData<fn(R)>,
|
||||
}
|
||||
|
||||
pub async fn open_provider_stream<P, R, E>(
|
||||
provider: &'static P,
|
||||
request: R,
|
||||
) -> Result<OpenedStream<E>, Error>
|
||||
where
|
||||
P: StreamProvider<R, E>,
|
||||
R: Send + 'static,
|
||||
E: Send + 'static,
|
||||
{
|
||||
let prepared = provider.transform_request(request)?;
|
||||
let opened = provider.call(prepared).await?;
|
||||
let state = PipelineState {
|
||||
provider,
|
||||
decoder: provider.decoder(),
|
||||
chunks: opened.chunks,
|
||||
pending: VecDeque::new(),
|
||||
finished: false,
|
||||
request: PhantomData,
|
||||
};
|
||||
let events = stream::unfold(state, |mut state| async move {
|
||||
loop {
|
||||
if let Some(event) = state.pending.pop_front() {
|
||||
return Some((event, state));
|
||||
}
|
||||
if state.finished {
|
||||
return None;
|
||||
}
|
||||
match state.chunks.next().await {
|
||||
Some(Ok(chunk)) => match state.decoder.push(chunk) {
|
||||
Ok(events) => queue_normalized(&mut state, events),
|
||||
Err(error) => {
|
||||
state.finished = true;
|
||||
return Some((Err(error), state));
|
||||
}
|
||||
},
|
||||
Some(Err(error)) => {
|
||||
state.finished = true;
|
||||
return Some((Err(error), state));
|
||||
}
|
||||
None => {
|
||||
state.finished = true;
|
||||
match state.decoder.finish() {
|
||||
Ok(events) => queue_normalized(&mut state, events),
|
||||
Err(error) => return Some((Err(error), state)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(OpenedStream {
|
||||
metadata: opened.metadata,
|
||||
events: Box::pin(events),
|
||||
})
|
||||
}
|
||||
|
||||
fn queue_normalized<P, D, E, R>(state: &mut PipelineState<P, D, E, R>, events: Vec<D::WireEvent>)
|
||||
where
|
||||
P: StreamProvider<R, E, Decoder = D>,
|
||||
D: StreamDecoder<WireEvent = <P as StreamProvider<R, E>>::WireEvent>,
|
||||
{
|
||||
for event in events {
|
||||
match state.provider.normalize(event) {
|
||||
Ok(normalized) => state.pending.extend(normalized.into_iter().map(Ok)),
|
||||
Err(error) => {
|
||||
state.pending.push_back(Err(error));
|
||||
state.finished = true;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use futures_util::future::FutureExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct FakeRequest;
|
||||
struct PreparedRequest;
|
||||
|
||||
struct FakeProvider {
|
||||
calls: Arc<Mutex<Vec<&'static str>>>,
|
||||
}
|
||||
|
||||
struct FakeDecoder {
|
||||
calls: Arc<Mutex<Vec<&'static str>>>,
|
||||
pending: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum FailurePoint {
|
||||
Transform,
|
||||
Call,
|
||||
Chunk,
|
||||
Push,
|
||||
Finish,
|
||||
Normalize,
|
||||
}
|
||||
|
||||
struct FailureProvider(FailurePoint);
|
||||
|
||||
struct FailureDecoder(FailurePoint);
|
||||
|
||||
impl StreamDecoder for FailureDecoder {
|
||||
type WireEvent = String;
|
||||
|
||||
fn push(&mut self, chunk: ProviderStreamChunk) -> Result<Vec<Self::WireEvent>, Error> {
|
||||
if matches!(self.0, FailurePoint::Push) {
|
||||
return Err(Error::InvalidResponse("decoder push failed".to_string()));
|
||||
}
|
||||
Ok(vec![
|
||||
String::from_utf8(chunk.0).expect("test chunk should be UTF-8"),
|
||||
])
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<Self::WireEvent>, Error> {
|
||||
if matches!(self.0, FailurePoint::Finish) {
|
||||
return Err(Error::InvalidResponse("decoder finish failed".to_string()));
|
||||
}
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamProvider<FakeRequest, String> for FailureProvider {
|
||||
type PreparedRequest = PreparedRequest;
|
||||
type WireEvent = String;
|
||||
type Decoder = FailureDecoder;
|
||||
|
||||
fn transform_request(&self, _request: FakeRequest) -> Result<Self::PreparedRequest, Error> {
|
||||
if matches!(self.0, FailurePoint::Transform) {
|
||||
return Err(Error::InvalidRequest("transform failed".to_string()));
|
||||
}
|
||||
Ok(PreparedRequest)
|
||||
}
|
||||
|
||||
fn call(
|
||||
&'static self,
|
||||
_request: Self::PreparedRequest,
|
||||
) -> BoxFuture<'static, Result<OpenedWireStream, Error>> {
|
||||
async move {
|
||||
if matches!(self.0, FailurePoint::Call) {
|
||||
return Err(Error::Network("open failed".to_string()));
|
||||
}
|
||||
let chunks: ProviderChunkStream = match self.0 {
|
||||
FailurePoint::Chunk => Box::pin(stream::iter([
|
||||
Err(Error::Network("source failed".to_string())),
|
||||
Ok(ProviderStreamChunk::new("ignored")),
|
||||
])),
|
||||
FailurePoint::Finish => Box::pin(stream::empty()),
|
||||
_ => Box::pin(stream::iter([Ok(ProviderStreamChunk::new("event"))])),
|
||||
};
|
||||
Ok(OpenedWireStream {
|
||||
metadata: StreamMetadata {
|
||||
status_code: 200,
|
||||
provider: StreamProviderId::OpenAi,
|
||||
transport: StreamTransport::Http,
|
||||
response_headers: Vec::new(),
|
||||
},
|
||||
chunks,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn decoder(&self) -> Self::Decoder {
|
||||
FailureDecoder(self.0)
|
||||
}
|
||||
|
||||
fn normalize(&self, event: Self::WireEvent) -> Result<Vec<String>, Error> {
|
||||
if matches!(self.0, FailurePoint::Normalize) {
|
||||
return Err(Error::InvalidResponse("normalize failed".to_string()));
|
||||
}
|
||||
Ok(vec![event])
|
||||
}
|
||||
}
|
||||
|
||||
struct PendingUntilDropped(Arc<AtomicBool>);
|
||||
|
||||
impl Stream for PendingUntilDropped {
|
||||
type Item = Result<ProviderStreamChunk, Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PendingUntilDropped {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
struct PendingProvider(Arc<AtomicBool>);
|
||||
|
||||
impl StreamProvider<FakeRequest, String> for PendingProvider {
|
||||
type PreparedRequest = PreparedRequest;
|
||||
type WireEvent = String;
|
||||
type Decoder = FailureDecoder;
|
||||
|
||||
fn transform_request(&self, _request: FakeRequest) -> Result<Self::PreparedRequest, Error> {
|
||||
Ok(PreparedRequest)
|
||||
}
|
||||
|
||||
fn call(
|
||||
&'static self,
|
||||
_request: Self::PreparedRequest,
|
||||
) -> BoxFuture<'static, Result<OpenedWireStream, Error>> {
|
||||
async move {
|
||||
Ok(OpenedWireStream {
|
||||
metadata: StreamMetadata {
|
||||
status_code: 200,
|
||||
provider: StreamProviderId::OpenAi,
|
||||
transport: StreamTransport::Http,
|
||||
response_headers: Vec::new(),
|
||||
},
|
||||
chunks: Box::pin(PendingUntilDropped(self.0.clone())),
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn decoder(&self) -> Self::Decoder {
|
||||
FailureDecoder(FailurePoint::Push)
|
||||
}
|
||||
|
||||
fn normalize(&self, event: Self::WireEvent) -> Result<Vec<String>, Error> {
|
||||
Ok(vec![event])
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamDecoder for FakeDecoder {
|
||||
type WireEvent = String;
|
||||
|
||||
fn push(&mut self, chunk: ProviderStreamChunk) -> Result<Vec<Self::WireEvent>, Error> {
|
||||
self.calls.lock().expect("call log").push("decode");
|
||||
self.pending
|
||||
.push_str(std::str::from_utf8(chunk.as_bytes()).expect("test utf-8"));
|
||||
let mut parts = self
|
||||
.pending
|
||||
.split('|')
|
||||
.map(str::to_string)
|
||||
.collect::<Vec<_>>();
|
||||
self.pending = parts.pop().expect("split always returns one item");
|
||||
Ok(parts)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<Self::WireEvent>, Error> {
|
||||
if self.pending.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Ok(vec![std::mem::take(&mut self.pending)])
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamProvider<FakeRequest, String> for FakeProvider {
|
||||
type PreparedRequest = PreparedRequest;
|
||||
type WireEvent = String;
|
||||
type Decoder = FakeDecoder;
|
||||
|
||||
fn transform_request(&self, _request: FakeRequest) -> Result<Self::PreparedRequest, Error> {
|
||||
self.calls.lock().expect("call log").push("transform");
|
||||
Ok(PreparedRequest)
|
||||
}
|
||||
|
||||
fn call(
|
||||
&'static self,
|
||||
_request: Self::PreparedRequest,
|
||||
) -> BoxFuture<'static, Result<OpenedWireStream, Error>> {
|
||||
self.calls.lock().expect("call log").push("call");
|
||||
async move {
|
||||
Ok(OpenedWireStream {
|
||||
metadata: StreamMetadata {
|
||||
status_code: 200,
|
||||
provider: StreamProviderId::Anthropic,
|
||||
transport: StreamTransport::Http,
|
||||
response_headers: vec![Header {
|
||||
name: "x-test".to_string(),
|
||||
value: "ready".to_string(),
|
||||
}],
|
||||
},
|
||||
chunks: Box::pin(stream::iter([
|
||||
Ok(ProviderStreamChunk::new(b"one|tw".to_vec())),
|
||||
Ok(ProviderStreamChunk::new(b"o|three".to_vec())),
|
||||
])),
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn decoder(&self) -> Self::Decoder {
|
||||
FakeDecoder {
|
||||
calls: self.calls.clone(),
|
||||
pending: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize(&self, event: Self::WireEvent) -> Result<Vec<String>, Error> {
|
||||
self.calls.lock().expect("call log").push("normalize");
|
||||
Ok(vec![event.to_uppercase()])
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fake_provider_proves_pipeline_order_and_fragmentation() {
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let provider = Box::leak(Box::new(FakeProvider {
|
||||
calls: calls.clone(),
|
||||
}));
|
||||
let mut opened = open_provider_stream(provider, FakeRequest)
|
||||
.await
|
||||
.expect("stream opens");
|
||||
let events = opened
|
||||
.events
|
||||
.by_ref()
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.expect("events normalize");
|
||||
|
||||
assert_eq!(events, ["ONE", "TWO", "THREE"]);
|
||||
assert_eq!(opened.metadata.response_headers[0].name, "x-test");
|
||||
assert_eq!(
|
||||
*calls.lock().expect("call log"),
|
||||
[
|
||||
"transform",
|
||||
"call",
|
||||
"decode",
|
||||
"normalize",
|
||||
"decode",
|
||||
"normalize",
|
||||
"normalize",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_and_open_failures_stop_before_a_stream_is_returned() {
|
||||
for (point, expected) in [
|
||||
(FailurePoint::Transform, "invalid request: transform failed"),
|
||||
(FailurePoint::Call, "upstream network error: open failed"),
|
||||
] {
|
||||
let provider = Box::leak(Box::new(FailureProvider(point)));
|
||||
let error = match open_provider_stream(provider, FakeRequest).await {
|
||||
Ok(_) => panic!("failure should prevent the stream from opening"),
|
||||
Err(error) => error,
|
||||
};
|
||||
assert_eq!(error.to_string(), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipeline_failures_are_emitted_once_and_then_terminate() {
|
||||
for (point, expected) in [
|
||||
(FailurePoint::Chunk, "upstream network error: source failed"),
|
||||
(FailurePoint::Push, "invalid response: decoder push failed"),
|
||||
(
|
||||
FailurePoint::Finish,
|
||||
"invalid response: decoder finish failed",
|
||||
),
|
||||
(
|
||||
FailurePoint::Normalize,
|
||||
"invalid response: normalize failed",
|
||||
),
|
||||
] {
|
||||
let provider = Box::leak(Box::new(FailureProvider(point)));
|
||||
let mut opened = open_provider_stream(provider, FakeRequest)
|
||||
.await
|
||||
.expect("stream should open before its terminal failure");
|
||||
let error = opened
|
||||
.events
|
||||
.next()
|
||||
.await
|
||||
.expect("stream should emit its error")
|
||||
.expect_err("first event should be the configured failure");
|
||||
assert_eq!(error.to_string(), expected);
|
||||
assert!(opened.events.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropping_the_event_stream_drops_the_provider_chunk_stream() {
|
||||
let dropped = Arc::new(AtomicBool::new(false));
|
||||
let provider = Box::leak(Box::new(PendingProvider(dropped.clone())));
|
||||
let opened = open_provider_stream(provider, FakeRequest)
|
||||
.await
|
||||
.expect("stream should open");
|
||||
|
||||
drop(opened.events);
|
||||
|
||||
assert!(dropped.load(Ordering::SeqCst));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,8 @@
|
|||
//! Enforcement: the litellm-rust workspace has exactly three crates.
|
||||
//! Enforcement: the litellm-rust workspace has exactly four crates.
|
||||
//!
|
||||
//! `core` (pure translation), `ai-gateway` (routes + all network I/O), and
|
||||
//! `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a
|
||||
//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host),
|
||||
//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
|
||||
//! PyO3 cdylib). Adding or removing a crate must be a
|
||||
//! deliberate act: this test fails until the allowlist here is updated, forcing
|
||||
//! whoever changes the crate set to justify the new crate per the rule that a
|
||||
//! crate is a layer needing independent compilation / its own deps / a separate
|
||||
|
|
@ -16,10 +17,15 @@ use std::path::{Path, PathBuf};
|
|||
|
||||
/// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the
|
||||
/// workspace legitimately gains or loses a crate.
|
||||
const EXPECTED_MEMBERS: &[&str] = &["crates/core", "crates/ai-gateway", "crates/python-bridge"];
|
||||
const EXPECTED_MEMBERS: &[&str] = &[
|
||||
"crates/core",
|
||||
"crates/ai-gateway",
|
||||
"crates/python-interop",
|
||||
"crates/python-bridge",
|
||||
];
|
||||
|
||||
/// The crate subdirectory names that must exist under `crates/`.
|
||||
const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-bridge"];
|
||||
const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"];
|
||||
|
||||
const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact).";
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue