From ffeae6bc3011e0f613b03e84656e2888e7aa21a1 Mon Sep 17 00:00:00 2001 From: Muhammad Yaseen Date: Fri, 21 Nov 2025 12:10:18 +0100 Subject: [PATCH 1/4] feat: added support for azure open ai api --- .env.example | 8 ++++++ graphgen/models/llm/api/openai_client.py | 32 ++++++++++++++++++++---- graphgen/operators/init/init_llm.py | 7 +++--- 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/.env.example b/.env.example index 94064796..5101b269 100644 --- a/.env.example +++ b/.env.example @@ -14,6 +14,14 @@ TRAINEE_MODEL=gpt-4o-mini TRAINEE_BASE_URL= TRAINEE_API_KEY= +# azure_openai_api +# SYNTHESIZER_BACKEND=azure_openai_api +# The following is the same as your "Deployment name" in Azure +# SYNTHESIZER_MODEL= +# SYNTHESIZER_BASE_URL=https://.openai.azure.com/openai/deployments//chat/completions +# SYNTHESIZER_API_KEY= +# SYNTHESIZER_API_VERSION= + # # ollama_api # SYNTHESIZER_BACKEND=ollama_api # SYNTHESIZER_MODEL=gemma3 diff --git a/graphgen/models/llm/api/openai_client.py b/graphgen/models/llm/api/openai_client.py index 448d6625..3c97c170 100644 --- a/graphgen/models/llm/api/openai_client.py +++ b/graphgen/models/llm/api/openai_client.py @@ -2,7 +2,8 @@ from typing import Any, Dict, List, Optional import openai -from openai import APIConnectionError, APITimeoutError, AsyncOpenAI, RateLimitError +from openai import APIConnectionError, APITimeoutError, AsyncOpenAI, AsyncAzureOpenAI, RateLimitError +from pyparsing import Literal from tenacity import ( retry, retry_if_exception_type, @@ -35,17 +36,20 @@ def __init__( model: str = "gpt-4o-mini", api_key: Optional[str] = None, base_url: Optional[str] = None, + api_version: Optional[str] = None, json_mode: bool = False, seed: Optional[int] = None, topk_per_token: int = 5, # number of topk tokens to generate for each token request_limit: bool = False, rpm: Optional[RPM] = None, tpm: Optional[TPM] = None, + backend: str = "openai_api", **kwargs: Any, ): super().__init__(**kwargs) self.model = model self.api_key = api_key + self.api_version = api_version # required for Azure OpenAI self.base_url = base_url self.json_mode = json_mode self.seed = seed @@ -56,13 +60,31 @@ def __init__( self.rpm = rpm or RPM() self.tpm = tpm or TPM() + assert backend in ["openai_api", "azure_openai_api"], f"Unsupported backend {backend}. Use 'openai_api' or 'azure_openai_api'." + self.backend = backend + self.__post_init__() def __post_init__(self): - assert self.api_key is not None, "Please provide api key to access openai api." - self.client = AsyncOpenAI( - api_key=self.api_key or "dummy", base_url=self.base_url - ) + + api_name = self.backend.replace("_", " ") + assert self.api_key is not None, f"Please provide api key to access {api_name}." + + if self.backend == "openai_api": + self.client = AsyncOpenAI( + api_key=self.api_key or "dummy", base_url=self.base_url + ) + elif self.backend == "azure_openai_api": + assert self.api_version is not None, f"Please provide api_version for {api_name}." + assert self.base_url is not None, f"Please provide base_url for {api_name}." + self.client = AsyncAzureOpenAI( + api_key=self.api_key, + azure_endpoint=self.base_url, + api_version=self.api_version, + azure_deployment=self.model, + ) + else: + raise ValueError(f"Unsupported backend {self.backend}. Use 'openai_api' or 'azure_openai_api'.") def _pre_generate(self, text: str, history: List[str]) -> Dict: kwargs = { diff --git a/graphgen/operators/init/init_llm.py b/graphgen/operators/init/init_llm.py index 79a3618b..0aabf4f6 100644 --- a/graphgen/operators/init/init_llm.py +++ b/graphgen/operators/init/init_llm.py @@ -27,10 +27,11 @@ def create_llm_wrapper(backend: str, config: Dict[str, Any]) -> BaseLLMWrapper: from graphgen.models.llm.api.http_client import HTTPClient return HTTPClient(**config) - if backend == "openai_api": + if backend == "openai_api" or backend == "azure_openai_api": from graphgen.models.llm.api.openai_client import OpenAIClient - - return OpenAIClient(**config) + + # pass in concrete backend to the OpenAIClient so that internally we can distinguish between OpenAI and Azure OpenAI + return OpenAIClient(**config, backend=backend) if backend == "ollama_api": from graphgen.models.llm.api.ollama_client import OllamaClient From 26ce90c7b164a9de555ceb609e15245567522407 Mon Sep 17 00:00:00 2001 From: Muhammad Yaseen Date: Fri, 21 Nov 2025 12:18:07 +0100 Subject: [PATCH 2/4] misc: removed unused import --- graphgen/models/llm/api/openai_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/graphgen/models/llm/api/openai_client.py b/graphgen/models/llm/api/openai_client.py index 3c97c170..a8c9a6e6 100644 --- a/graphgen/models/llm/api/openai_client.py +++ b/graphgen/models/llm/api/openai_client.py @@ -3,7 +3,6 @@ import openai from openai import APIConnectionError, APITimeoutError, AsyncOpenAI, AsyncAzureOpenAI, RateLimitError -from pyparsing import Literal from tenacity import ( retry, retry_if_exception_type, From 821f6441d04b90df70761b7ed9fc3ed066089c4f Mon Sep 17 00:00:00 2001 From: Muhammad Yaseen Date: Fri, 21 Nov 2025 13:34:48 +0000 Subject: [PATCH 3/4] misc: fixing pylint issues --- graphgen/models/llm/api/openai_client.py | 5 +++-- graphgen/operators/init/init_llm.py | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/graphgen/models/llm/api/openai_client.py b/graphgen/models/llm/api/openai_client.py index a8c9a6e6..a4d3dc12 100644 --- a/graphgen/models/llm/api/openai_client.py +++ b/graphgen/models/llm/api/openai_client.py @@ -59,7 +59,9 @@ def __init__( self.rpm = rpm or RPM() self.tpm = tpm or TPM() - assert backend in ["openai_api", "azure_openai_api"], f"Unsupported backend {backend}. Use 'openai_api' or 'azure_openai_api'." + assert ( + backend in ("openai_api", "azure_openai_api") + ), f"Unsupported backend '{backend}'. Use 'openai_api' or 'azure_openai_api'." self.backend = backend self.__post_init__() @@ -68,7 +70,6 @@ def __post_init__(self): api_name = self.backend.replace("_", " ") assert self.api_key is not None, f"Please provide api key to access {api_name}." - if self.backend == "openai_api": self.client = AsyncOpenAI( api_key=self.api_key or "dummy", base_url=self.base_url diff --git a/graphgen/operators/init/init_llm.py b/graphgen/operators/init/init_llm.py index 0aabf4f6..1ce44d60 100644 --- a/graphgen/operators/init/init_llm.py +++ b/graphgen/operators/init/init_llm.py @@ -27,10 +27,10 @@ def create_llm_wrapper(backend: str, config: Dict[str, Any]) -> BaseLLMWrapper: from graphgen.models.llm.api.http_client import HTTPClient return HTTPClient(**config) - if backend == "openai_api" or backend == "azure_openai_api": + if backend in ("openai_api", "azure_openai_api"): from graphgen.models.llm.api.openai_client import OpenAIClient - - # pass in concrete backend to the OpenAIClient so that internally we can distinguish between OpenAI and Azure OpenAI + # pass in concrete backend to the OpenAIClient so that internally we can distinguish + # between OpenAI and Azure OpenAI return OpenAIClient(**config, backend=backend) if backend == "ollama_api": from graphgen.models.llm.api.ollama_client import OllamaClient From dca04cd767263d93b301bf8e00896cf9854009bd Mon Sep 17 00:00:00 2001 From: Muhammad Yaseen Date: Fri, 21 Nov 2025 14:39:29 +0000 Subject: [PATCH 4/4] misc: pylint fixes --- graphgen/models/llm/api/openai_client.py | 4 ++-- graphgen/operators/init/init_llm.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/graphgen/models/llm/api/openai_client.py b/graphgen/models/llm/api/openai_client.py index a4d3dc12..532b981c 100644 --- a/graphgen/models/llm/api/openai_client.py +++ b/graphgen/models/llm/api/openai_client.py @@ -68,9 +68,9 @@ def __init__( def __post_init__(self): - api_name = self.backend.replace("_", " ") + api_name = self.backend.replace("_", " ") assert self.api_key is not None, f"Please provide api key to access {api_name}." - if self.backend == "openai_api": + if self.backend == "openai_api": self.client = AsyncOpenAI( api_key=self.api_key or "dummy", base_url=self.base_url ) diff --git a/graphgen/operators/init/init_llm.py b/graphgen/operators/init/init_llm.py index 1ce44d60..e294d2c3 100644 --- a/graphgen/operators/init/init_llm.py +++ b/graphgen/operators/init/init_llm.py @@ -29,7 +29,7 @@ def create_llm_wrapper(backend: str, config: Dict[str, Any]) -> BaseLLMWrapper: return HTTPClient(**config) if backend in ("openai_api", "azure_openai_api"): from graphgen.models.llm.api.openai_client import OpenAIClient - # pass in concrete backend to the OpenAIClient so that internally we can distinguish + # pass in concrete backend to the OpenAIClient so that internally we can distinguish # between OpenAI and Azure OpenAI return OpenAIClient(**config, backend=backend) if backend == "ollama_api":