promptflow

1
import os
2

3
from dotenv import load_dotenv
4
from openai.version import VERSION as OPENAI_VERSION
5

6
from promptflow.tracing import trace
7

8

9
def get_client():
10
    if OPENAI_VERSION.startswith("0."):
11
        raise Exception(
12
            "Please upgrade your OpenAI package to version >= 1.0.0 or using the command: pip install --upgrade openai."
13
        )
14
    api_key = os.environ.get("OPENAI_API_KEY", None)
15
    if api_key:
16
        from openai import OpenAI
17

18
        return OpenAI()
19
    else:
20
        from openai import AzureOpenAI
21

22
        return AzureOpenAI(
23
            api_version=os.environ.get("OPENAI_API_VERSION", "2023-07-01-preview")
24
        )
25

26

27
@trace
28
def my_llm_tool(
29
    prompt: str,
30
    # for AOAI, deployment name is customized by user, not model name.
31
    deployment_name: str,
32
    max_tokens: int = 120,
33
    temperature: float = 1.0,
34
    top_p: float = 1.0,
35
    n: int = 1,
36
) -> str:
37
    if "OPENAI_API_KEY" not in os.environ and "AZURE_OPENAI_API_KEY" not in os.environ:
38
        # load environment variables from .env file
39
        load_dotenv()
40

41
    if "OPENAI_API_KEY" not in os.environ and "AZURE_OPENAI_API_KEY" not in os.environ:
42
        raise Exception(
43
            "Please specify environment variables: OPENAI_API_KEY or AZURE_OPENAI_API_KEY"
44
        )
45
    messages = [{"content": prompt, "role": "system"}]
46
    response = get_client().chat.completions.create(
47
        messages=messages,
48
        model=deployment_name,
49
        max_tokens=int(max_tokens),
50
        temperature=float(temperature),
51
        top_p=float(top_p),
52
        n=int(n),
53
    )
54

55
    # get first element because prompt is single.
56
    return response.choices[0].message.content
57

58

59
if __name__ == "__main__":
60
    result = my_llm_tool(
61
        prompt="Write a simple Hello, world! program that displays the greeting message.",
62
        deployment_name="gpt-35-turbo",
63
    )
64
    print(result)
65