kobai-sdk 0.2.8rc15__tar.gz → 0.3.0rc1__tar.gz

{kobai_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: kobai-sdk
-Version: 0.2.8rc15
+Version: 0.3.0rc1
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License
@@ -223,12 +223,14 @@ Requires-Dist: langchain-core
 Requires-Dist: langchain-community
 Requires-Dist: langchain_openai
 Requires-Dist: databricks_langchain
+Requires-Dist: sentence-transformers
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: bumpver; extra == "dev"
 Requires-Dist: isort; extra == "dev"
 Requires-Dist: pip-tools; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
 
 # Kobai SDK for Python (Alpha)
 

kobai_sdk-0.3.0rc1/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_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/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_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/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.8rc15 → kobai_sdk-0.3.0rc1}/kobai_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: kobai-sdk
-Version: 0.2.8rc15
+Version: 0.3.0rc1
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License
@@ -223,12 +223,14 @@ Requires-Dist: langchain-core
 Requires-Dist: langchain-community
 Requires-Dist: langchain_openai
 Requires-Dist: databricks_langchain
+Requires-Dist: sentence-transformers
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: bumpver; extra == "dev"
 Requires-Dist: isort; extra == "dev"
 Requires-Dist: pip-tools; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
 
 # Kobai SDK for Python (Alpha)
 

{kobai_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/kobai_sdk.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ kobai/ai_query.py
 kobai/ai_rag.py
 kobai/databricks_client.py
 kobai/demo_tenant_client.py
+kobai/ms_authenticate.py
 kobai/spark_client.py
 kobai/tenant_api.py
 kobai/tenant_client.py

{kobai_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/kobai_sdk.egg-info/requires.txt

@@ -7,6 +7,7 @@ langchain-core
 langchain-community
 langchain_openai
 databricks_langchain
+sentence-transformers
 
 [dev]
 black

{kobai_sdk-0.2.8rc15 → kobai_sdk-0.3.0rc1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "kobai-sdk"
-version = "0.2.8rc15"
+version = "0.3.0rc1"
 description = "A package that enables interaction with a Kobai tenant."
 readme = "README.md"
 authors = [{ name = "Ryan Oattes", email = "ryan@kobai.io" }]
@@ -26,7 +26,8 @@ dependencies = [
     "langchain-core",
     "langchain-community",
     "langchain_openai",
-    "databricks_langchain"
+    "databricks_langchain",
+    "sentence-transformers"
     ]
 requires-python = ">=3.11"