diff --git a/packages/gen/gen_ai_hub/proxy/langchain/init_models.py b/packages/gen/gen_ai_hub/proxy/langchain/init_models.py index 0f3794a8..a5be61dc 100644 --- a/packages/gen/gen_ai_hub/proxy/langchain/init_models.py +++ b/packages/gen/gen_ai_hub/proxy/langchain/init_models.py @@ -93,6 +93,7 @@ def _init_model(proxy_client: Optional[BaseProxyClient], kwargs: Dict[str, Any], init_func: Optional[Callable] = None, model_kwargs: Optional[Dict[str, Any]] = None): + proxy_client = proxy_client or get_proxy_client() model_kwargs = model_kwargs or {} if init_func: return _init_custom_model(proxy_client=proxy_client, init_func=init_func, args=args, kwargs=kwargs, diff --git a/packages/gen/tests/proxy/langchain_/test_init_models.py b/packages/gen/tests/proxy/langchain_/test_init_models.py index bf78687d..ee57664b 100644 --- a/packages/gen/tests/proxy/langchain_/test_init_models.py +++ b/packages/gen/tests/proxy/langchain_/test_init_models.py @@ -1,4 +1,5 @@ import unittest +import unittest.mock from datetime import datetime from unittest.mock import MagicMock @@ -75,6 +76,24 @@ def test_init_embedding_model(self): model = init_embedding_model(model_name, proxy_client=self.proxy_client) self.assertIsInstance(model, model_class) + def test_init_llm_without_proxy_client(self): + """init_llm should fall back to get_proxy_client() when proxy_client is not passed.""" + with unittest.mock.patch( + 'gen_ai_hub.proxy.langchain.init_models.get_proxy_client', + return_value=self.proxy_client, + ): + model = init_llm('gpt-5') + self.assertIsInstance(model, openai.ChatOpenAI) + + def test_init_embedding_model_without_proxy_client(self): + """init_embedding_model should fall back to get_proxy_client() when proxy_client is not passed.""" + with unittest.mock.patch( + 'gen_ai_hub.proxy.langchain.init_models.get_proxy_client', + return_value=self.proxy_client, + ): + model = init_embedding_model('text-embedding-3-small') + self.assertIsInstance(model, openai.OpenAIEmbeddings) + def test_init_llm(self): model_kwargs = {'top_k': 3} for model_name, model_class in self.llm.items():