[BUG] OpenAIRerankerClient does not support AzureOpenAILLMClient #1006
Description
Bug Description
The OpenAIRerankerClient supports OpenAIClient but does not correctly support its sibling class AzureOpenAILLMClient.
In the init method, only OpenAIClient is explicitly unwrapped, while AzureOpenAILLMClient is passed through as-is. As a result, self.client is not the expected AsyncAzureOpenAI instance, and subsequent API calls (e.g., chat.completions) fail.
Steps to Reproduce
Use Azure OpenAI with an AzureOpenAILLMClient and call the search() method
Expected Behavior
No error should occur, and a valid response should be returned from the model.
Actual Behavior
The following error occurs when calling search():
AttributeError: 'AzureOpenAILLMClient' object has no attribute 'chat'Environment
- Graphiti Version: 0.22.0
- Python Version: 3.10.12
- Operating System: Ubuntu 22.04
- Database Backend: Neo4j 5.26.2
- LLM Provider & Model: Azure OpenAI gpt-4.1
Installation Method
- pip install
- uv add
- Development installation (git clone)
Error Messages/Traceback
AttributeError: 'AzureOpenAILLMClient' object has no attribute 'chat'Configuration
Additional Context
Possible Solution
The issue can be resolved by either of the following approaches:
- Use the shared base class for detection:
class OpenAIRerankerClient(CrossEncoderClient):
def __init__(
self,
config: LLMConfig | None = None,
client: AsyncOpenAI | AsyncAzureOpenAI | BaseOpenAIClient | None = None,
):
"""
Initialize the OpenAIRerankerClient with the provided configuration and client.
This reranker uses the OpenAI API to run a simple boolean classifier prompt concurrently
for each passage. Log-probabilities are used to rank the passages.
Args:
config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
client (AsyncOpenAI | AsyncAzureOpenAI | BaseOpenAIClient | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
"""
if config is None:
config = LLMConfig()
self.config = config
if client is None:
self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
elif isinstance(client, BaseOpenAIClient):
self.client = client.client
else:
self.client = client- Explicitly support both wrapper classes:
elif isinstance(client, (OpenAIClient, AzureOpenAILLMClient)):
self.client = client.client