my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import json
#2	import os
#3	from typing import Dict, List, Optional, Union
#4
#5	from azure.identity import DefaultAzureCredential, get_bearer_token_provider
#6	from openai import AzureOpenAI
#7
#8	from mem0.configs.llms.azure import AzureOpenAIConfig
#9	from mem0.configs.llms.base import BaseLlmConfig
#10	from mem0.llms.base import LLMBase
#11	from mem0.memory.utils import extract_json
#12
#13	SCOPE = "https://cognitiveservices.azure.com/.default"
#14
#15
#16	class AzureOpenAILLM(LLMBase):
#17	def __init__(self, config: Optional[Union[BaseLlmConfig, AzureOpenAIConfig, Dict]] = None):
#18	# Convert to AzureOpenAIConfig if needed
#19	if config is None:
#20	config = AzureOpenAIConfig()
#21	elif isinstance(config, dict):
#22	config = AzureOpenAIConfig(**config)
#23	elif isinstance(config, BaseLlmConfig) and not isinstance(config, AzureOpenAIConfig):
#24	# Convert BaseLlmConfig to AzureOpenAIConfig
#25	config = AzureOpenAIConfig(
#26	model=config.model,
#27	temperature=config.temperature,
#28	api_key=config.api_key,
#29	max_tokens=config.max_tokens,
#30	top_p=config.top_p,
#31	top_k=config.top_k,
#32	enable_vision=config.enable_vision,
#33	vision_details=config.vision_details,
#34	http_client_proxies=config.http_client,
#35	)
#36
#37	super().__init__(config)
#38
#39	# Model name should match the custom deployment name chosen for it.
#40	if not self.config.model:
#41	self.config.model = "gpt-4.1-nano-2025-04-14"
#42
#43	api_key = self.config.azure_kwargs.api_key or os.getenv("LLM_AZURE_OPENAI_API_KEY")
#44	azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("LLM_AZURE_DEPLOYMENT")
#45	azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("LLM_AZURE_ENDPOINT")
#46	api_version = self.config.azure_kwargs.api_version or os.getenv("LLM_AZURE_API_VERSION")
#47	default_headers = self.config.azure_kwargs.default_headers
#48
#49	# If the API key is not provided or is a placeholder, use DefaultAzureCredential.
#50	if api_key is None or api_key == "" or api_key == "your-api-key":
#51	self.credential = DefaultAzureCredential()
#52	azure_ad_token_provider = get_bearer_token_provider(
#53	self.credential,
#54	SCOPE,
#55	)
#56	api_key = None
#57	else:
#58	azure_ad_token_provider = None
#59
#60	self.client = AzureOpenAI(
#61	azure_deployment=azure_deployment,
#62	azure_endpoint=azure_endpoint,
#63	azure_ad_token_provider=azure_ad_token_provider,
#64	api_version=api_version,
#65	api_key=api_key,
#66	http_client=self.config.http_client,
#67	default_headers=default_headers,
#68	)
#69
#70	def _parse_response(self, response, tools):
#71	"""
#72	Process the response based on whether tools are used or not.
#73
#74	Args:
#75	response: The raw response from API.
#76	tools: The list of tools provided in the request.
#77
#78	Returns:
#79	str or dict: The processed response.
#80	"""
#81	if tools:
#82	processed_response = {
#83	"content": response.choices[0].message.content,
#84	"tool_calls": [],
#85	}
#86
#87	if response.choices[0].message.tool_calls:
#88	for tool_call in response.choices[0].message.tool_calls:
#89	processed_response["tool_calls"].append(
#90	{
#91	"name": tool_call.function.name,
#92	"arguments": json.loads(extract_json(tool_call.function.arguments)),
#93	}
#94	)
#95
#96	return processed_response
#97	else:
#98	return response.choices[0].message.content
#99
#100	def generate_response(
#101	self,
#102	messages: List[Dict[str, str]],
#103	response_format=None,
#104	tools: Optional[List[Dict]] = None,
#105	tool_choice: str = "auto",
#106	**kwargs,
#107	):
#108	"""
#109	Generate a response based on the given messages using Azure OpenAI.
#110
#111	Args:
#112	messages (list): List of message dicts containing 'role' and 'content'.
#113	response_format (str or object, optional): Format of the response. Defaults to "text".
#114	tools (list, optional): List of tools that the model can call. Defaults to None.
#115	tool_choice (str, optional): Tool choice method. Defaults to "auto".
#116	**kwargs: Additional Azure OpenAI-specific parameters.
#117
#118	Returns:
#119	str: The generated response.
#120	"""
#121
#122	user_prompt = messages[-1]["content"]
#123
#124	user_prompt = user_prompt.replace("assistant", "ai")
#125
#126	messages[-1]["content"] = user_prompt
#127
#128	params = self._get_supported_params(messages=messages, **kwargs)
#129
#130	# Add model and messages
#131	params.update({
#132	"model": self.config.model,
#133	"messages": messages,
#134	})
#135
#136	if tools:
#137	params["tools"] = tools
#138	params["tool_choice"] = tool_choice
#139
#140	response = self.client.chat.completions.create(**params)
#141	return self._parse_response(response, tools)
#142

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public