- Notifications
You must be signed in to change notification settings - Fork 2.8k
/
Copy pathtest_passthrough_endpoints.py
66 lines (54 loc) · 2.37 KB
/
test_passthrough_endpoints.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
importpytest
importasyncio
importaiohttp, openai
fromopenaiimportOpenAI, AsyncOpenAI
fromtypingimportOptional, List, Union
importaiohttp
importasyncio
importjson
importos
importdotenv
dotenv.load_dotenv()
asyncdefcohere_rerank(session):
url="http://localhost:4000/v1/rerank"
headers= {
"Authorization": f"bearer {os.getenv('COHERE_API_KEY')}",
"Content-Type": "application/json",
"Accept": "application/json",
}
data= {
"model": "rerank-english-v3.0",
"query": "What is the capital of the United States?",
"top_n": 3,
"documents": [
"Carson City is the capital city of the American state of Nevada.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
"Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
"Capitalization or capitalisation in English grammar is the use of a capital letter at the start of a word. English usage varies from capitalization in other languages.",
"Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
],
}
asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()
print(f"Status: {status}")
print(f"Response:\n{response_text}")
print()
ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()
@pytest.mark.asyncio
@pytest.mark.skip(
reason="new test just added by @ishaan-jaff, still figuring out how to run this in ci/cd"
)
asyncdeftest_basic_passthrough():
"""
- Make request to pass through endpoint
- This SHOULD not go through LiteLLM user_api_key_auth
- This should forward headers from request to pass through endpoint
"""
asyncwithaiohttp.ClientSession() assession:
response=awaitcohere_rerank(session)
print("response from cohere rerank", response)
assertresponse["id"] isnotNone
assertresponse["results"] isnotNone