chainlit 2.0.0__py3-none-any.whl → 2.0.dev0__py3-none-any.whl

chainlit/oauth_providers.py

@@ -4,10 +4,9 @@ import urllib.parse
 from typing import Dict, List, Optional, Tuple
 
 import httpx
-from fastapi import HTTPException
-
 from chainlit.secret import random_secret
 from chainlit.user import User
+from fastapi import HTTPException
 
 
 class OAuthProvider:
@@ -17,31 +16,15 @@ 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
-
-    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
+        raise NotImplementedError()
 
 
 class GithubOAuthProvider(OAuthProvider):
@@ -54,11 +37,9 @@ class GithubOAuthProvider(OAuthProvider):
         self.client_secret = os.environ.get("OAUTH_GITHUB_CLIENT_SECRET")
         self.authorize_params = {
             "scope": "user:email",
+            "prompt": "consent",
         }
 
-        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,
@@ -115,11 +96,9 @@ class GoogleOAuthProvider(OAuthProvider):
             "scope": "https://www.googleapis.com/auth/userinfo.profile https://www.googleapis.com/auth/userinfo.email",
             "response_type": "code",
             "access_type": "offline",
+            "prompt": "login",
         }
 
-        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,
@@ -185,11 +164,9 @@ class AzureADOAuthProvider(OAuthProvider):
             "response_type": "code",
             "scope": "https://graph.microsoft.com/User.Read",
             "response_mode": "query",
+            "prompt": "login",
         }
 
-        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,
@@ -272,11 +249,9 @@ class AzureADHybridOAuthProvider(OAuthProvider):
             "scope": "https://graph.microsoft.com/User.Read https://graph.microsoft.com/openid",
             "response_mode": "form_post",
             "nonce": nonce,
+            "prompt": "login",
         }
 
-        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,
@@ -354,11 +329,9 @@ class OktaOAuthProvider(OAuthProvider):
             "response_type": "code",
             "scope": "openid profile email",
             "response_mode": "query",
+            "prompt": "login",
         }
 
-        if prompt := self.get_prompt():
-            self.authorize_params["prompt"] = prompt
-
     def get_authorization_server_path(self):
         if not self.authorization_server_id:
             return "/default"
@@ -428,11 +401,9 @@ class Auth0OAuthProvider(OAuthProvider):
             "response_type": "code",
             "scope": "openid profile email",
             "audience": f"{self.original_domain}/userinfo",
+            "prompt": "login",
         }
 
-        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,
@@ -488,11 +459,9 @@ class DescopeOAuthProvider(OAuthProvider):
             "response_type": "code",
             "scope": "openid profile email",
             "audience": f"{self.domain}/userinfo",
+            "prompt": "login",
         }
 
-        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,
@@ -549,11 +518,9 @@ class AWSCognitoOAuthProvider(OAuthProvider):
             "response_type": "code",
             "client_id": self.client_id,
             "scope": "openid profile email",
+            "prompt": "login",
         }
 
-        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,
@@ -620,11 +587,9 @@ class GitlabOAuthProvider(OAuthProvider):
         self.authorize_params = {
             "scope": "openid profile email",
             "response_type": "code",
+            "prompt": "login",
         }
 
-        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,
@@ -665,71 +630,6 @@ class GitlabOAuthProvider(OAuthProvider):
             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(),
@@ -740,7 +640,6 @@ providers = [
     DescopeOAuthProvider(),
     AWSCognitoOAuthProvider(),
     GitlabOAuthProvider(),
-    KeycloakOAuthProvider(),
 ]
 
 

chainlit/openai/__init__.py

@@ -1,12 +1,11 @@
 import asyncio
 from typing import Union
 
-from literalai import ChatGeneration, CompletionGeneration
-from literalai.helper import timestamp_utc
-
-from chainlit.context import local_steps
+from chainlit.context import get_context
 from chainlit.step import Step
 from chainlit.utils import check_module_version
+from literalai import ChatGeneration, CompletionGeneration
+from literalai.helper import timestamp_utc
 
 
 def instrument_openai():
@@ -20,9 +19,11 @@ def instrument_openai():
     def on_new_generation(
         generation: Union["ChatGeneration", "CompletionGeneration"], timing
     ):
-        previous_steps = local_steps.get()
+        context = get_context()
 
-        parent_id = previous_steps[-1].id if previous_steps else None
+        parent_id = None
+        if context.current_step:
+            parent_id = context.current_step.id
 
         step = Step(
             name=generation.model if generation.model else generation.provider,

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)))