Coverage for src/local_deep_research/embeddings/providers/implementations/owned_openai_embeddings.py: 89%

1"""OpenAI embedding clients with explicit, retryable resource cleanup.""" 

2 

3import asyncio 

4import threading 

5from typing import Any 

6 

7from langchain_openai import OpenAIEmbeddings 

8from pydantic import ConfigDict, PrivateAttr 

9 

10from ....security.secure_logging import logger 

11 

12 

13async def _aclose_http_clients( 

14 http_client: Any, http_async_client: Any 

15) -> bool: 

16 closed = True 

17 try: 

18 http_client.close() 

19 except Exception: 

20 closed = False 

21 logger.warning("Failed to close OpenAI embeddings sync httpx client") 

22 try: 

23 await http_async_client.aclose() 

24 except Exception: 

25 closed = False 

26 logger.warning("Failed to close OpenAI embeddings async httpx client") 

27 return closed 

28 

29 

30def close_openai_embedding_clients( 

31 http_client: Any, http_async_client: Any 

32) -> bool: 

33 try: 

34 asyncio.get_running_loop() 

35 except RuntimeError: 

36 try: 

37 return asyncio.run( 

38 _aclose_http_clients(http_client, http_async_client) 

39 ) 

40 except Exception: 

41 logger.warning("Failed to close OpenAI embeddings httpx clients") 

42 return False 

43 

44 closed = False 

45 

46 def _close_in_thread() -> None: 

47 nonlocal closed 

48 try: 

49 closed = asyncio.run( 

50 _aclose_http_clients(http_client, http_async_client) 

51 ) 

52 except Exception: 

53 logger.warning("Failed to close OpenAI embeddings httpx clients") 

54 

55 close_thread = threading.Thread( 

56 target=_close_in_thread, 

57 daemon=True, 

58 name="ldr-async-openai-embedding-close", 

59 ) 

60 close_thread.start() 

61 close_thread.join(timeout=5) 

62 if close_thread.is_alive(): 

63 logger.warning("OpenAI embeddings async httpx close exceeded 5s") 

64 return False 

65 return closed 

66 

67 

68class OwnedOpenAIEmbeddings(OpenAIEmbeddings): 

69 model_config = ConfigDict(arbitrary_types_allowed=True) 

70 http_client: Any 

71 http_async_client: Any 

72 _closed: bool = PrivateAttr(default=False) 

73 

74 def close(self) -> None: 

75 """Close from synchronous code; use ``aclose`` on the owning loop. 

76 

77 The fallback for a running loop waits up to five seconds. A client 

78 whose transport cannot close on another loop remains retryable. 

79 """ 

80 if self._closed: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true

81 return 

82 if close_openai_embedding_clients( 

83 self.http_client, self.http_async_client 

84 ): 

85 self._closed = True 

86 

87 async def aclose(self) -> None: 

88 if self._closed: 

89 return 

90 if await _aclose_http_clients(self.http_client, self.http_async_client): 

91 self._closed = True