chainlit 1.0.401__py3-none-any.whl → 2.0.4__py3-none-any.whl

chainlit/oauth_providers.py

@@ -4,9 +4,11 @@ import urllib.parse
 from typing import Dict, List, Optional, Tuple
 
 import httpx
-from chainlit.user import User
 from fastapi import HTTPException
 
+from chainlit.secret import random_secret
+from chainlit.user import User
+
 
 class OAuthProvider:
     id: str
@@ -15,15 +17,31 @@ class OAuthProvider:
     client_secret: str
     authorize_url: str
     authorize_params: Dict[str, str]
+    default_prompt: Optional[str] = None
 
     def is_configured(self):
         return all([os.environ.get(env) for env in self.env])
 
     async def get_token(self, code: str, url: str) -> str:
-        raise NotImplementedError()
+        raise NotImplementedError
 
     async def get_user_info(self, token: str) -> Tuple[Dict[str, str], User]:
-        raise NotImplementedError()
+        raise NotImplementedError
+
+    def get_env_prefix(self) -> str:
+        """Return environment prefix, like AZURE_AD."""
+
+        return self.id.replace("-", "_").upper()
+
+    def get_prompt(self) -> Optional[str]:
+        """Return OAuth prompt param."""
+        if prompt := os.environ.get(f"OAUTH_{self.get_env_prefix()}_PROMPT"):
+            return prompt
+
+        if prompt := os.environ.get("OAUTH_PROMPT"):
+            return prompt
+
+        return self.default_prompt
 
 
 class GithubOAuthProvider(OAuthProvider):
@@ -38,6 +56,9 @@ class GithubOAuthProvider(OAuthProvider):
             "scope": "user:email",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -96,6 +117,9 @@ class GoogleOAuthProvider(OAuthProvider):
             "access_type": "offline",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -163,6 +187,9 @@ class AzureADOAuthProvider(OAuthProvider):
             "response_mode": "query",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -203,10 +230,97 @@ class AzureADOAuthProvider(OAuthProvider):
                 )
                 photo_data = await photo_response.aread()
                 base64_image = base64.b64encode(photo_data)
-                azure_user[
-                    "image"
-                ] = f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
-            except Exception as e:
+                azure_user["image"] = (
+                    f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
+                )
+            except Exception:
+                # Ignore errors getting the photo
+                pass
+
+            user = User(
+                identifier=azure_user["userPrincipalName"],
+                metadata={"image": azure_user.get("image"), "provider": "azure-ad"},
+            )
+            return (azure_user, user)
+
+
+class AzureADHybridOAuthProvider(OAuthProvider):
+    id = "azure-ad-hybrid"
+    env = [
+        "OAUTH_AZURE_AD_HYBRID_CLIENT_ID",
+        "OAUTH_AZURE_AD_HYBRID_CLIENT_SECRET",
+        "OAUTH_AZURE_AD_HYBRID_TENANT_ID",
+    ]
+    authorize_url = (
+        f"https://login.microsoftonline.com/{os.environ.get('OAUTH_AZURE_AD_HYBRID_TENANT_ID', '')}/oauth2/v2.0/authorize"
+        if os.environ.get("OAUTH_AZURE_AD_HYBRID_ENABLE_SINGLE_TENANT")
+        else "https://login.microsoftonline.com/common/oauth2/v2.0/authorize"
+    )
+    token_url = (
+        f"https://login.microsoftonline.com/{os.environ.get('OAUTH_AZURE_AD_HYBRID_TENANT_ID', '')}/oauth2/v2.0/token"
+        if os.environ.get("OAUTH_AZURE_AD_HYBRID_ENABLE_SINGLE_TENANT")
+        else "https://login.microsoftonline.com/common/oauth2/v2.0/token"
+    )
+
+    def __init__(self):
+        self.client_id = os.environ.get("OAUTH_AZURE_AD_HYBRID_CLIENT_ID")
+        self.client_secret = os.environ.get("OAUTH_AZURE_AD_HYBRID_CLIENT_SECRET")
+        nonce = random_secret(16)
+        self.authorize_params = {
+            "tenant": os.environ.get("OAUTH_AZURE_AD_HYBRID_TENANT_ID"),
+            "response_type": "code id_token",
+            "scope": "https://graph.microsoft.com/User.Read https://graph.microsoft.com/openid",
+            "response_mode": "form_post",
+            "nonce": nonce,
+        }
+
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
+    async def get_token(self, code: str, url: str):
+        payload = {
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "code": code,
+            "grant_type": "authorization_code",
+            "redirect_uri": url,
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                self.token_url,
+                data=payload,
+            )
+            response.raise_for_status()
+            json = response.json()
+
+            token = json["access_token"]
+            if not token:
+                raise HTTPException(
+                    status_code=400, detail="Failed to get the access token"
+                )
+            return token
+
+    async def get_user_info(self, token: str):
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                "https://graph.microsoft.com/v1.0/me",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            response.raise_for_status()
+
+            azure_user = response.json()
+
+            try:
+                photo_response = await client.get(
+                    "https://graph.microsoft.com/v1.0/me/photos/48x48/$value",
+                    headers={"Authorization": f"Bearer {token}"},
+                )
+                photo_data = await photo_response.aread()
+                base64_image = base64.b64encode(photo_data)
+                azure_user["image"] = (
+                    f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
+                )
+            except Exception:
                 # Ignore errors getting the photo
                 pass
 
@@ -242,6 +356,9 @@ class OktaOAuthProvider(OAuthProvider):
             "response_mode": "query",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     def get_authorization_server_path(self):
         if not self.authorization_server_id:
             return "/default"
@@ -313,6 +430,9 @@ class Auth0OAuthProvider(OAuthProvider):
             "audience": f"{self.original_domain}/userinfo",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -357,7 +477,7 @@ class DescopeOAuthProvider(OAuthProvider):
     id = "descope"
     env = ["OAUTH_DESCOPE_CLIENT_ID", "OAUTH_DESCOPE_CLIENT_SECRET"]
     # Ensure that the domain does not have a trailing slash
-    domain = f"https://api.descope.com/oauth2/v1"
+    domain = "https://api.descope.com/oauth2/v1"
 
     authorize_url = f"{domain}/authorize"
 
@@ -370,6 +490,9 @@ class DescopeOAuthProvider(OAuthProvider):
             "audience": f"{self.domain}/userinfo",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -428,6 +551,9 @@ class AWSCognitoOAuthProvider(OAuthProvider):
             "scope": "openid profile email",
         }
 
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
     async def get_token(self, code: str, url: str):
         payload = {
             "client_id": self.client_id,
@@ -475,14 +601,146 @@ class AWSCognitoOAuthProvider(OAuthProvider):
             return (cognito_user, user)
 
 
+class GitlabOAuthProvider(OAuthProvider):
+    id = "gitlab"
+    env = [
+        "OAUTH_GITLAB_CLIENT_ID",
+        "OAUTH_GITLAB_CLIENT_SECRET",
+        "OAUTH_GITLAB_DOMAIN",
+    ]
+
+    def __init__(self):
+        self.client_id = os.environ.get("OAUTH_GITLAB_CLIENT_ID")
+        self.client_secret = os.environ.get("OAUTH_GITLAB_CLIENT_SECRET")
+        # Ensure that the domain does not have a trailing slash
+        self.domain = f"https://{os.environ.get('OAUTH_GITLAB_DOMAIN', '').rstrip('/')}"
+
+        self.authorize_url = f"{self.domain}/oauth/authorize"
+
+        self.authorize_params = {
+            "scope": "openid profile email",
+            "response_type": "code",
+        }
+
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
+    async def get_token(self, code: str, url: str):
+        payload = {
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "code": code,
+            "grant_type": "authorization_code",
+            "redirect_uri": url,
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.domain}/oauth/token",
+                data=payload,
+            )
+            response.raise_for_status()
+            json_content = response.json()
+            token = json_content.get("access_token")
+            if not token:
+                raise HTTPException(
+                    status_code=400, detail="Failed to get the access token"
+                )
+            return token
+
+    async def get_user_info(self, token: str):
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"{self.domain}/oauth/userinfo",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            response.raise_for_status()
+            gitlab_user = response.json()
+            user = User(
+                identifier=gitlab_user.get("email"),
+                metadata={
+                    "image": gitlab_user.get("picture", ""),
+                    "provider": "gitlab",
+                },
+            )
+            return (gitlab_user, user)
+
+
+class KeycloakOAuthProvider(OAuthProvider):
+    env = [
+        "OAUTH_KEYCLOAK_CLIENT_ID",
+        "OAUTH_KEYCLOAK_CLIENT_SECRET",
+        "OAUTH_KEYCLOAK_REALM",
+        "OAUTH_KEYCLOAK_BASE_URL",
+    ]
+    id = os.environ.get("OAUTH_KEYCLOAK_NAME", "keycloak")
+
+    def __init__(self):
+        self.client_id = os.environ.get("OAUTH_KEYCLOAK_CLIENT_ID")
+        self.client_secret = os.environ.get("OAUTH_KEYCLOAK_CLIENT_SECRET")
+        self.realm = os.environ.get("OAUTH_KEYCLOAK_REALM")
+        self.base_url = os.environ.get("OAUTH_KEYCLOAK_BASE_URL")
+        self.authorize_url = (
+            f"{self.base_url}/realms/{self.realm}/protocol/openid-connect/auth"
+        )
+
+        self.authorize_params = {
+            "scope": "profile email openid",
+            "response_type": "code",
+        }
+
+        if prompt := self.get_prompt():
+            self.authorize_params["prompt"] = prompt
+
+    async def get_token(self, code: str, url: str):
+        payload = {
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "code": code,
+            "grant_type": "authorization_code",
+            "redirect_uri": url,
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/realms/{self.realm}/protocol/openid-connect/token",
+                data=payload,
+            )
+            response.raise_for_status()
+            json = response.json()
+            token = json.get("access_token")
+            if not token:
+                raise httpx.HTTPStatusError(
+                    "Failed to get the access token",
+                    request=response.request,
+                    response=response,
+                )
+            return token
+
+    async def get_user_info(self, token: str):
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"{self.base_url}/realms/{self.realm}/protocol/openid-connect/userinfo",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            response.raise_for_status()
+            kc_user = response.json()
+            user = User(
+                identifier=kc_user["email"],
+                metadata={"provider": "keycloak"},
+            )
+            return (kc_user, user)
+
+
 providers = [
     GithubOAuthProvider(),
     GoogleOAuthProvider(),
     AzureADOAuthProvider(),
+    AzureADHybridOAuthProvider(),
     OktaOAuthProvider(),
     Auth0OAuthProvider(),
     DescopeOAuthProvider(),
     AWSCognitoOAuthProvider(),
+    GitlabOAuthProvider(),
+    KeycloakOAuthProvider(),
 ]
 
 

chainlit/openai/__init__.py

@@ -1,12 +1,13 @@
+import asyncio
 from typing import Union
 
-from chainlit.context import get_context
-from chainlit.step import Step
-from chainlit.sync import run_sync
-from chainlit.utils import check_module_version
 from literalai import ChatGeneration, CompletionGeneration
 from literalai.helper import timestamp_utc
 
+from chainlit.context import local_steps
+from chainlit.step import Step
+from chainlit.utils import check_module_version
+
 
 def instrument_openai():
     if not check_module_version("openai", "1.0.0"):
@@ -16,16 +17,12 @@ def instrument_openai():
 
     from literalai.instrumentation.openai import instrument_openai
 
-    async def on_new_generation(
+    def on_new_generation(
         generation: Union["ChatGeneration", "CompletionGeneration"], timing
     ):
-        context = get_context()
+        previous_steps = local_steps.get()
 
-        parent_id = None
-        if context.current_step:
-            parent_id = context.current_step.id
-        elif context.session.root_message:
-            parent_id = context.session.root_message.id
+        parent_id = previous_steps[-1].id if previous_steps else None
 
         step = Step(
             name=generation.model if generation.model else generation.provider,
@@ -52,11 +49,6 @@ def instrument_openai():
             step.input = generation.prompt
             step.output = generation.completion
 
-        await step.send()
-
-    def on_new_generation_sync(
-        generation: Union["ChatGeneration", "CompletionGeneration"], timing
-    ):
-        run_sync(on_new_generation(generation, timing))
+        asyncio.create_task(step.send())
 
-    instrument_openai(None, on_new_generation_sync)
+    instrument_openai(None, on_new_generation)

chainlit/secret.py

@@ -6,4 +6,4 @@ chars = string.ascii_letters + string.digits + "$%*,-./:=>?@^_~"
 
 
 def random_secret(length: int = 64):
-    return "".join((secrets.choice(chars) for i in range(length)))
+    return "".join(secrets.choice(chars) for i in range(length))