diff --git a/README.md b/README.md index c3ea724..b75596e 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ After you have Python and PostgreSQL installed, follow these steps: 5. `pip install -r requirements.txt` (install the dependencies) 6. `cd pilot` 7. `mv .env.example .env` (create the .env file) -8. Add your OpenAI API key and the PostgreSQL database info to the `.env` file +8. Add your environment (OpenAI/Azure), your API key and the PostgreSQL database info to the `.env` file 9. `python db_init.py` (initialize the database) 10. `python main.py` (start GPT Pilot) diff --git a/pilot/.env.example b/pilot/.env.example index d53210e..e7b661c 100644 --- a/pilot/.env.example +++ b/pilot/.env.example @@ -1,4 +1,11 @@ +#OPENAI or AZURE +ENDPOINT=OPENAI OPENAI_API_KEY= +AZURE_API_KEY= +AZURE_ENDPOINT= +#In case of Azure endpoint, change this to your deployed model name +MODEL_NAME=gpt-4 +MAX_TOKENS=8192 DB_NAME=gpt-pilot DB_HOST=localhost DB_PORT=5432 diff --git a/pilot/const/llm.py b/pilot/const/llm.py index ca42267..8ae8ab8 100644 --- a/pilot/const/llm.py +++ b/pilot/const/llm.py @@ -1,4 +1,5 @@ +import os +MAX_GPT_MODEL_TOKENS = int(os.getenv('MAX_TOKENS')) MIN_TOKENS_FOR_GPT_RESPONSE = 600 -MAX_GPT_MODEL_TOKENS = 8192 MAX_QUESTIONS = 5 -END_RESPONSE = "EVERYTHING_CLEAR" +END_RESPONSE = "EVERYTHING_CLEAR" \ No newline at end of file diff --git a/pilot/utils/llm_connection.py b/pilot/utils/llm_connection.py index 58f4b6a..369af7c 100644 --- a/pilot/utils/llm_connection.py +++ b/pilot/utils/llm_connection.py @@ -46,8 +46,11 @@ def get_tokens_in_messages(messages: List[str]) -> int: tokenized_messages = [tokenizer.encode(message['content']) for message in messages] return sum(len(tokens) for tokens in tokenized_messages) +#get endpoint and model name from .ENV file +model = os.getenv('MODEL_NAME') +endpoint = os.getenv('ENDPOINT') -def num_tokens_from_functions(functions, model="gpt-4"): +def num_tokens_from_functions(functions, model=model): """Return the number of tokens used by a list of functions.""" encoding = tiktoken.get_encoding("cl100k_base") @@ -171,13 +174,22 @@ def stream_gpt_completion(data, req_type): # spinner = spinner_start(colored("Waiting for OpenAI API response...", 'yellow')) # print(colored("Stream response from OpenAI:", 'yellow')) - api_key = os.getenv("OPENAI_API_KEY") logger.info(f'Request data: {data}') + # Check if the ENDPOINT is AZURE + if endpoint == 'AZURE': + # If yes, get the AZURE_ENDPOINT from .ENV file + endpoint_url = os.getenv('AZURE_ENDPOINT') + '/openai/deployments/' + model + '/chat/completions?api-version=2023-05-15' + headers = {'Content-Type': 'application/json', 'api-key': os.getenv('AZURE_API_KEY')} + else: + # If not, send the request to the OpenAI endpoint + headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer ' + os.getenv("OPENAI_API_KEY")} + endpoint_url = 'https://api.openai.com/v1/chat/completions' + response = requests.post( - 'https://api.openai.com/v1/chat/completions', - headers={'Content-Type': 'application/json', 'Authorization': 'Bearer ' + api_key}, + endpoint_url, + headers=headers, json=data, stream=True )