FilesExpand file tree

main

/

chat_async.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

82 lines (67 loc) · 2.62 KB

main

/

chat_async.py

Copy path

Top

File metadata and controls

• Code
• Blame

82 lines (67 loc) · 2.62 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

import asyncio

import os

import azure.identity.aio

import openai

from dotenv import load_dotenv

# Setup the OpenAI client to use either Azure, OpenAI.com, or Ollama API

load_dotenv(override=True)

API_HOST = os.getenv("API_HOST", "github")

azure_credential = None # Will hold the Azure credential so we can close it properly.

if API_HOST == "azure":

azure_credential = azure.identity.aio.DefaultAzureCredential()

token_provider = azure.identity.aio.get_bearer_token_provider(

azure_credential, "https://cognitiveservices.azure.com/.default"

)

client = openai.AsyncOpenAI(

base_url=os.environ["AZURE_OPENAI_ENDPOINT"],

api_key=token_provider,

)

MODEL_NAME = os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"]

elif API_HOST == "ollama":

client = openai.AsyncOpenAI(base_url=os.environ["OLLAMA_ENDPOINT"], api_key="nokeyneeded")

MODEL_NAME = os.environ["OLLAMA_MODEL"]

elif API_HOST == "github":

client = openai.AsyncOpenAI(base_url="https://models.github.ai/inference", api_key=os.environ["GITHUB_TOKEN"])

MODEL_NAME = os.getenv("GITHUB_MODEL", "openai/gpt-4o")

else:

client = openai.AsyncOpenAI(api_key=os.environ["OPENAI_KEY"])

MODEL_NAME = os.environ["OPENAI_MODEL"]

async def generate_response(location):

print("Generating response for", location)

response = await client.chat.completions.create(

model=MODEL_NAME,

messages=[

{"role": "system", "content": "You are a helpful assistant."},

{

"role": "user",

"content": f"Name a single place I should visit on my trip to {location} and describe in one sentence",

},

],

temperature=0.7,

)

print("Got response for ", location)

return response.choices[0].message.content

async def single() -> None:

"""Run a single request example and handle cleanup."""

print(await generate_response("Tokyo"))

async def multiple() -> None:

"""Run multiple requests concurrently and handle cleanup."""

answers = await asyncio.gather(

generate_response("Tokyo"),

generate_response("Berkeley"),

generate_response("Singapore"),

)

for answer in answers:

print(answer, "\n")

async def close_clients() -> None:

"""Close the OpenAI async client and (if applicable) the Azure credential."""

await client.close()

if azure_credential is not None:

await azure_credential.close()

async def main():

try:

await single() # Change to await multiple() to run multiple requests concurrently

finally:

await close_clients()

if __name__ == "__main__":

asyncio.run(main())