| 104 | |
| 105 | class AzureLLM(LLM): |
| 106 | def __init__(self, config): |
| 107 | super().__init__(config) |
| 108 | if "openai_api_key" in config: |
| 109 | openai_api_key = config["openai_api_key"] |
| 110 | elif "OPENAI_API_KEY" in os.environ: |
| 111 | openai_api_key = os.environ["OPENAI_API_KEY"] |
| 112 | else: |
| 113 | assert False, "OpenAI API key not found" |
| 114 | |
| 115 | self.client = AzureOpenAI( |
| 116 | azure_endpoint=os.environ["AZURE_ENDPOINT_GPT35"] if "gpt35" in self.config["model"] else os.environ["AZURE_ENDPOINT_GPT4"], |
| 117 | api_key=openai_api_key, |
| 118 | api_version=os.environ["API_VERSION"], |
| 119 | azure_deployment="ai4code-research-gpt4o" |
| 120 | ) |
| 121 | |
| 122 | def __call__(self, prompt: str): |
| 123 | # prompt = truncate_tokens(prompt, encoding_name=self.config["model"], max_length=self.config["max_tokens"]) |