Coverage for aipyapp/llm/client_oauth2.py: 29%

1#! /usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4import time 

5import httpx 

6from typing import Optional, Dict 

7from .base_openai import OpenAIBaseClient 

8 

9class OAuth2Client(OpenAIBaseClient): 

10 """ 

11 OAuth2-based OpenAI LLM client that: 

12 1. First gets token using client_id/client_secret 

13 2. Then uses the token to make LLM API calls 

14 """ 

15 def __init__(self, config: Dict): 

16 super().__init__(config) 

17 self._token_url = config.get("token_url") 

18 self._client_id = config.get("client_id") 

19 self._client_secret = config.get("client_secret") 

20 self._access_token: Optional[str] = None 

21 self._token_expires = 0 

22 

23 def usable(self) -> bool: 

24 return all([ 

25 self._token_url, 

26 self._client_id, 

27 self._client_secret, 

28 self._base_url 

29 ]) 

30 

31 def _get_client(self): 

32 self._api_key = self._get_access_token() 

33 return super()._get_client() 

34 

35 def _get_access_token(self) -> str: 

36 """Get OAuth2 access token using client credentials""" 

37 current_time = time.time() 

38 

39 # Return existing token if it's still valid (with 300 seconds buffer) 

40 if self._access_token and current_time < (self._token_expires - 300): 

41 return self._access_token 

42 

43 auth_data = { 

44 'client_id': self._client_id, 

45 'client_secret': self._client_secret, 

46 'grant_type': 'client_credentials' 

47 } 

48 

49 with httpx.Client(timeout=self._timeout, verify=self._tls_verify) as client: 

50 response = client.post( 

51 self._token_url, 

52 data=auth_data 

53 ) 

54 response.raise_for_status() 

55 token_data = response.json() 

56 self._access_token = token_data['access_token'] 

57 

58 # Calculate expiration time (default to 5 mins if not provided) 

59 expires_in = token_data.get("expires_in", 300) 

60 self._token_expires = current_time + expires_in 

61 

62 return self._access_token 

63 

64 def get_completion(self, messages): 

65 response = super().get_completion(messages) 

66 self._client = None 

67 

68 return response