| 103 | |
| 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"]) |
| 124 | response = self.client.chat.completions.create( |
| 125 | temperature=0, |
| 126 | model=self.config["model"], |
| 127 | messages=[ |
| 128 | {"role": "system", "content": self.system_prompt}, |
| 129 | {"role": "user", "content": prompt}, |
| 130 | ] |
| 131 | ) |
| 132 | return response.choices[0].message.content |
| 133 | |
| 134 | class VLLM(LLM): |
| 135 | def __init__(self, config): |