kobai-sdk 0.2.9__py3-none-any.whl → 0.3.0rc1__py3-none-any.whl

kobai/ms_authenticate.py

@@ -0,0 +1,66 @@
+from azure.identity import DeviceCodeCredential
+from azure.identity import OnBehalfOfCredential
+from azure.core.exceptions import AzureError
+
+def get_scope(client_id: str = None, target_client_id: str = None, scope: str = None):
+
+    """
+    Get the default scopes 
+
+    Parameters:
+    client_id (str): Client ID or Application ID from app registration with IDM.
+    target_client_id (str): Kobai IDM client ID.
+    scope (str): Scope to be passed
+    """
+    if scope is not None:
+        return scope
+
+    if target_client_id is None:
+        target_client_id = client_id
+
+    return f"openid profile offline_access api://{target_client_id}/Kobai.Access"
+
+def device_code(tenant_id: str, client_id: str, target_client_id: str = None, scope: str = None):
+
+    """
+    Authenticate using the device code flow and get the access token
+
+    Parameters:
+    tenant_id (str): Tenant ID or Directory ID for IDM.
+    client_id (str): Client ID or Application ID from app registration with IDM.
+    target_client_id (str): Kobai IDM client ID.
+    scope (str): Scope to be passed
+    """
+    credential = DeviceCodeCredential(client_id=client_id, tenant_id=tenant_id)
+
+    try:
+        token = credential.get_token(get_scope(client_id, target_client_id, scope))
+        return token.token
+    except AzureError as e:
+        return e
+
+def onbehalf(tenant_id: str, client_id: str, client_secret: str, access_token: str, target_client_id: str = None, scope: str = None):
+
+    """
+    Authenticate using the onbehalf flow and get the access token
+
+    Parameters:
+    tenant_id (str): Tenant ID or Directory ID for IDM.
+    client_id (str): Client ID or Application ID from app registration with IDM.
+    client_secret (str): Client secret from app registration with IDM.
+    access_token (str): Access token to be exchanged.
+    target_client_id (str): Kobai IDM client ID.
+    scope (str): Scope to be passed
+    """
+    credential = OnBehalfOfCredential(
+        tenant_id=tenant_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        user_assertion=access_token
+    )
+
+    try:
+        token = credential.get_token(get_scope(client_id, target_client_id, scope))
+        return token.token
+    except AzureError as e:
+        return e

kobai/tenant_api.py

@@ -19,7 +19,10 @@ class TenantAPI:
         self.session = requests.Session()
 
         if token is not None:
-            self.session.headers.update({'Authorization': 'Bearer %s' % self.token})
+            if token.startswith('Bearer'):
+                 self.session.headers.update({'Authorization': '%s' % self.token})
+            else:
+                self.session.headers.update({'Authorization': 'Bearer %s' % self.token})
 
         self.ssl_verify = verify
         self.session.verify = verify
@@ -112,7 +115,7 @@ class TenantAPI:
 
         if op_desc is None:
             op_desc = "operation"
-
+        
         response = self.session.get(
             self.base_uri + uri,
             params=params,

kobai/tenant_client.py

@@ -3,7 +3,6 @@ import json
 import urllib
 import urllib.parse
 
-from azure.identity import DeviceCodeCredential
 from pyspark.sql import SparkSession
 
 from langchain_community.chat_models import ChatDatabricks
@@ -64,84 +63,76 @@ class TenantClient:
 # MS Entra Auth
 ########################################
 
-    def authenticate(self, client_id: str, tenant_id: str, run_ai_init: bool = True, override_username: str = None):
+    def use_browser_token(self, access_token, run_ai_init: bool = True):
 
         """
         Authenticate the TenantClient with the Kobai instance. Returns nothing, but stores bearer token in client.
-
-        Limitations:
-        Currently supports only authentication via Microsoft Entra (AzureAD) using DecideCode OAuth flow.
+        This is a fall-back method for instances not using OAuth. It is inconvenient as a Kobai Bearer Token must be retrieved from the users browser.
 
         Parameters:
-        client_id (str): Client ID or Application ID from app registration with IDM.
-        tenant_id (str): Tenant ID or Directory ID for IDM.
+        access_token (str): Bearer token for Kobai app session.
         """
+        self._init_post_auth_success(access_token, run_ai_init)
 
-        credential = DeviceCodeCredential(client_id=client_id, tenant_id=tenant_id)
-
-        access = credential.authenticate()
-
-        oauth_token = access.serialize()
-        print(oauth_token)
-        user_name = json.loads(access.serialize())["username"]
+    def use_access_token(self, access_token: str, id_token: str = None,  tenant_id: str = None, run_ai_init: bool = True):
 
-        if override_username is not None:
-            user_name = override_username
-
-        user_name_query_params={ 'userName' : user_name}
-        tenants_response = self.api_client._TenantAPI__run_get('/user-mgmt-svcs/auth/tenants?'+urllib.parse.urlencode(user_name_query_params))
-
-
-        tenant_list = json.loads(tenants_response.content.decode("utf-8"))
+        """
+        Authenticate the TenantClient with the Kobai instance. Returns nothing, but stores bearer token in client.
 
-        tenant_id = ""
-        for t in tenant_list:
-            if t["name"] == self.tenant_name:
-                tenant_id = t["id"]
+        Parameters:
+        access_token (str): Access token of the IDM server to be used to obtained the kobai access token.
+        id_token (str): ID token of the IDM server to be used to obtained the onbehalf access token.
+        tenant_id (str): Kobai tenant id.
+        """
 
         token_request_payload={
+            "tenantName" : self.tenant_name,
             "tenantId" : tenant_id,
-            "oauthToken" : oauth_token,
-            "userName" : user_name
+            "idToken" : id_token,
+            "accessToken" : access_token
         }
-        token_response = self.api_client._TenantAPI__run_post(
-            '/user-mgmt-svcs/auth/oauth/devicecode',
+
+        response = self.api_client._TenantAPI__run_post(
+            '/user-mgmt-svcs/auth/oauth/external/onbehalf/token',
             token_request_payload
         )
-        access_token = token_response.content.decode()
-        self.token = access_token
-        
-        self.__api_init_session()
-        self.__set_tenant_solutionid()
-        if run_ai_init:
-            self.init_ai_components()
-        
-        print("Authentication Successful.")
 
-    def authenticate_browser_token(self, access_token, run_ai_init: bool = True):
+        kb_access_token = response.headers.get('Authorization')
+        self.use_browser_token(kb_access_token, run_ai_init)
+
+    def get_tenants(self, id_token: str = None):
 
         """
-        Authenticate the TenantClient with the Kobai instance. Returns nothing, but stores bearer token in client.
-        This is a fall-back method for instances not using OAuth. It is inconvenient as a Kobai Bearer Token must be retrieved from the users browser.
+        Get the tenants associated with the given id token of the IDM. Returns tenants list.
 
         Parameters:
-        access_token (str): Bearer token for Kobai app session.
+        id_token (str): ID token of the IDM server to be used to obtain user tenants.
         """
 
+        if (id_token is not None) :
+            token_request_payload={
+                "idToken" : id_token
+            }
+
+            response = self.api_client._TenantAPI__run_post(
+                '/user-mgmt-svcs/auth/oauth/external/token/tenants',
+                token_request_payload
+            )
+
+            self.tenant_list = response.json()       
+        return self.tenant_list
+
+    def __api_init_session(self):
+        self.api_client = tenant_api.TenantAPI(self.token, self.uri, verify=self.ssl_verify, proxies=self.proxies )
+
+    def _init_post_auth_success(self, access_token, run_ai_init: bool = True):
         self.token = access_token
-        
         self.__api_init_session()
         self.__set_tenant_solutionid()
         if run_ai_init:
             self.init_ai_components()
-
-
         print("Authentication Successful.")
 
-    def __api_init_session(self):
-        self.api_client = tenant_api.TenantAPI(self.token, self.uri, verify=self.ssl_verify, proxies=self.proxies )
-        
-
 ########################################
 # Basic Config
 ########################################

{kobai_sdk-0.2.9.dist-info → kobai_sdk-0.3.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kobai-sdk
-Version: 0.2.9
+Version: 0.3.0rc1
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License

kobai_sdk-0.3.0rc1.dist-info/RECORD

@@ -0,0 +1,14 @@
+kobai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kobai/ai_query.py,sha256=xZh_OyakU01gIrnzaW4v_TdfzG51nPu0ntXJw1WEnvw,9424
+kobai/ai_rag.py,sha256=8B3HM4GoGVrgxJG678NN4vGaDwZRYnQiK5SCGiMIYkM,15186
+kobai/databricks_client.py,sha256=fyqqMly2Qm0r1AHWsQjkYeNsDdH0G1JSgTkF9KJ55qA,2118
+kobai/demo_tenant_client.py,sha256=wlNc-bdI2wotRXo8ppUOalv4hYdBlek_WzJNARZV-AE,9293
+kobai/ms_authenticate.py,sha256=rlmhtvAaSRBlYmvIBy5epMVa4MBGBLPaMwawu1T_xDQ,2252
+kobai/spark_client.py,sha256=opM_F-4Ut5Hq5zZjWMuLvUps9sDULvyPNZHXGL8dW1k,776
+kobai/tenant_api.py,sha256=Q5yuFd9_V4lo3LWzvYEEO3LpDRWFgQD4TlRPXDTGbiE,4368
+kobai/tenant_client.py,sha256=o2bifvmYwxL3_gpRMJaVRXrt9CN8Kfcx7p5r3jR9aGg,39415
+kobai_sdk-0.3.0rc1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+kobai_sdk-0.3.0rc1.dist-info/METADATA,sha256=9yGBoPdo2wWB_wYmAGG00dEY-4xqP7p9Grm1wUWcWfg,19263
+kobai_sdk-0.3.0rc1.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+kobai_sdk-0.3.0rc1.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
+kobai_sdk-0.3.0rc1.dist-info/RECORD,,

kobai/test.py

@@ -1,5 +0,0 @@
-import llm_config, ai_query
-
-llm_config = llm_config.LLMConfig(api_key="sV9LuoA5n0PwqggMXOYMhhZlt56FpgnMXFohimPhD7Ug3CnBLbO8JQQJ99ALACYeBjFXJ3w3AAABACOGZm8X", llm_provider="azure_openai")
-llm_config.get_azure_ad_token()
-ai_query.followup_question_1(question="abc", data={}, question_name="sample", llm_config=llm_config)

kobai_sdk-0.2.9.dist-info/RECORD

@@ -1,14 +0,0 @@
-kobai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kobai/ai_query.py,sha256=xZh_OyakU01gIrnzaW4v_TdfzG51nPu0ntXJw1WEnvw,9424
-kobai/ai_rag.py,sha256=8B3HM4GoGVrgxJG678NN4vGaDwZRYnQiK5SCGiMIYkM,15186
-kobai/databricks_client.py,sha256=fyqqMly2Qm0r1AHWsQjkYeNsDdH0G1JSgTkF9KJ55qA,2118
-kobai/demo_tenant_client.py,sha256=wlNc-bdI2wotRXo8ppUOalv4hYdBlek_WzJNARZV-AE,9293
-kobai/spark_client.py,sha256=opM_F-4Ut5Hq5zZjWMuLvUps9sDULvyPNZHXGL8dW1k,776
-kobai/tenant_api.py,sha256=9U6UbxpaAb-kpbuADXx3kbkNKaOzYy0I-GGwbpiCCOk,4212
-kobai/tenant_client.py,sha256=ocAjiX7ewb7JjCy7PiISX2in55dLVpkf8HGQv2fE6Fc,39565
-kobai/test.py,sha256=LHCKClyFNPgLA1hF62RlDzilaFw2Nd3BWcaLtOc5xdc,320
-kobai_sdk-0.2.9.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-kobai_sdk-0.2.9.dist-info/METADATA,sha256=HmKJKi_XVB2MECaQXziP8Cvfe594I09ztK_2Ne7KmHI,19260
-kobai_sdk-0.2.9.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-kobai_sdk-0.2.9.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
-kobai_sdk-0.2.9.dist-info/RECORD,,