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

1#! /usr/bin/env python

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

3

4 import time

5 import httpx

6 from typing import Optional, Dict

7 from .base_openai import OpenAIBaseClient

8

9 class 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