File size: 1,523 Bytes
cc2ce8c
f0fc5f8
6e28a81
fa9f031
 
6e28a81
fa9f031
 
 
f0fc5f8
 
6e28a81
cc2ce8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f0fc5f8
 
 
 
 
6e28a81
f0fc5f8
 
fa9f031
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
from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
import os

try:
    from dotenv import load_dotenv

    load_dotenv()
except:
    pass


def get_llm(max_tokens=1024, temperature=0.0, verbose=True, streaming=False, **kwargs):
    if has_azure_openai_config():
        return get_azure_llm(
            max_tokens=max_tokens,
            temperature=temperature,
            verbose=verbose,
            streaming=streaming,
            **kwargs,
        )
    return get_open_ai_llm(
        max_tokens=max_tokens,
        temperature=temperature,
        verbose=verbose,
        streaming=streaming,
        **kwargs,
    )


def has_azure_openai_config():
    """
    Checks if the necessary environment variables for Azure Blob Storage are set.
    Returns True if they are set, False otherwise.
    """
    return all(
        key in os.environ
        for key in [
            "AZURE_OPENAI_API_BASE_URL",
            "AZURE_OPENAI_API_VERSION",
            "AZURE_OPENAI_API_DEPLOYMENT_NAME",
            "AZURE_OPENAI_API_KEY",
        ]
    )


def get_open_ai_llm(**kwargs):
    return ChatOpenAI(**kwargs)


def get_azure_llm(**kwargs):
    llm = AzureChatOpenAI(
        openai_api_base=os.environ["AZURE_OPENAI_API_BASE_URL"],
        openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"],
        deployment_name=os.environ["AZURE_OPENAI_API_DEPLOYMENT_NAME"],
        openai_api_key=os.environ["AZURE_OPENAI_API_KEY"],
        openai_api_type="azure",
        **kwargs,
    )
    return llm