litellm

1
# What is this?
2
## Tests `litellm.transcription` endpoint. Outside litellm module b/c of audio file used in testing (it's ~700kb).
3

4
import pytest
5
import asyncio, time
6
import aiohttp, traceback
7
from openai import AsyncOpenAI
8
import sys, os, dotenv
9
from typing import Optional
10
from dotenv import load_dotenv
11

12
# Get the current directory of the file being run
13
pwd = os.path.dirname(os.path.realpath(__file__))
14
print(pwd)
15

16
file_path = os.path.join(pwd, "gettysburg.wav")
17

18
audio_file = open(file_path, "rb")
19

20
load_dotenv()
21

22
sys.path.insert(
23
    0, os.path.abspath("../")
24
)  # Adds the parent directory to the system path
25
import litellm
26
from litellm import Router
27

28

29
def test_transcription():
30
    transcript = litellm.transcription(
31
        model="whisper-1",
32
        file=audio_file,
33
    )
34
    print(f"transcript: {transcript.model_dump()}")
35
    print(f"transcript: {transcript._hidden_params}")
36

37

38
# test_transcription()
39

40

41
def test_transcription_azure():
42
    litellm.set_verbose = True
43
    transcript = litellm.transcription(
44
        model="azure/azure-whisper",
45
        file=audio_file,
46
        api_base="https://my-endpoint-europe-berri-992.openai.azure.com/",
47
        api_key=os.getenv("AZURE_EUROPE_API_KEY"),
48
        api_version="2024-02-15-preview",
49
    )
50

51
    print(f"transcript: {transcript}")
52
    assert transcript.text is not None
53
    assert isinstance(transcript.text, str)
54

55

56
# test_transcription_azure()
57

58

59
@pytest.mark.asyncio
60
async def test_transcription_async_azure():
61
    transcript = await litellm.atranscription(
62
        model="azure/azure-whisper",
63
        file=audio_file,
64
        api_base="https://my-endpoint-europe-berri-992.openai.azure.com/",
65
        api_key=os.getenv("AZURE_EUROPE_API_KEY"),
66
        api_version="2024-02-15-preview",
67
    )
68

69
    assert transcript.text is not None
70
    assert isinstance(transcript.text, str)
71

72

73
# asyncio.run(test_transcription_async_azure())
74

75

76
@pytest.mark.asyncio
77
async def test_transcription_async_openai():
78
    transcript = await litellm.atranscription(
79
        model="whisper-1",
80
        file=audio_file,
81
    )
82

83
    assert transcript.text is not None
84
    assert isinstance(transcript.text, str)
85

86

87
@pytest.mark.asyncio
88
async def test_transcription_on_router():
89
    litellm.set_verbose = True
90
    print("\n Testing async transcription on router\n")
91
    try:
92
        model_list = [
93
            {
94
                "model_name": "whisper",
95
                "litellm_params": {
96
                    "model": "whisper-1",
97
                },
98
            },
99
            {
100
                "model_name": "whisper",
101
                "litellm_params": {
102
                    "model": "azure/azure-whisper",
103
                    "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com/",
104
                    "api_key": os.getenv("AZURE_EUROPE_API_KEY"),
105
                    "api_version": "2024-02-15-preview",
106
                },
107
            },
108
        ]
109

110
        router = Router(model_list=model_list)
111
        response = await router.atranscription(
112
            model="whisper",
113
            file=audio_file,
114
        )
115
        print(response)
116
    except Exception as e:
117
        traceback.print_exc()
118
        pytest.fail(f"Error occurred: {e}")
119