auth0-ai-langchain 0.2.0__py3-none-any.whl → 1.0.0b2__py3-none-any.whl

auth0_ai_langchain/FGARetriever.py

@@ -32,7 +32,7 @@ class FGARetriever(BaseRetriever):
         Args:
             retriever (BaseRetriever): The retriever used to fetch documents.
             build_query (Callable[[Document], ClientBatchCheckItem]): Function to convert documents into FGA queries.
-            fga_configuration (Optional[ClientConfiguration]): Configuration for the OpenFGA client. If not provided, defaults to environment variables.
+            fga_configuration (ClientConfiguration, optional): Configuration for the OpenFGA client. If not provided, defaults to environment variables.
         """
         super().__init__()
         self._retriever = retriever
@@ -95,7 +95,7 @@ class FGARetriever(BaseRetriever):
 
         Args:
             query (str): The query for retrieving documents.
-            run_manager (Optional[object]): Optional manager for tracking runs.
+            run_manager (object, optional): Optional manager for tracking runs.
 
         Returns:
             List[Document]: Filtered and relevant documents.
@@ -148,7 +148,7 @@ class FGARetriever(BaseRetriever):
 
         Args:
             query (str): The query for retrieving documents.
-            run_manager (Optional[object]): Optional manager for tracking runs.
+            run_manager (object, optional): Optional manager for tracking runs.
 
         Returns:
             List[Document]: Filtered and relevant documents.

auth0_ai_langchain/auth0_ai.py

@@ -1,43 +1,113 @@
 from typing import Callable, Optional
-from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
-from auth0_ai.credentials import Credential
-from auth0_ai.authorizers.types import AuthorizerParams
-from auth0_ai.authorizers.federated_connection_authorizer import FederatedConnectionAuthorizerParams 
-from .federated_connections.federated_connection_authorizer import FederatedConnectionAuthorizer
-from .ciba.ciba_graph.ciba_graph import CIBAGraph
-from .ciba.ciba_graph.types import CIBAGraphOptions
-
-def get_access_token(config: RunnableConfig) -> Credential:
-    """
-    Fetch the access token obtained during the CIBA flow.
+from auth0_ai.authorizers.ciba import CIBAAuthorizerParams
+from auth0_ai.authorizers.federated_connection_authorizer import FederatedConnectionAuthorizerParams
+from auth0_ai.authorizers.types import Auth0ClientParams
+from auth0_ai_langchain.ciba.ciba_authorizer import CIBAAuthorizer
+from auth0_ai_langchain.federated_connections.federated_connection_authorizer import FederatedConnectionAuthorizer
+
 
-    Attributes:
-        config(RunnableConfig): LangGraph runnable configuration instance.
+class Auth0AI:
+    """Provides decorators to secure LangChain tools using Auth0 authorization flows.
     """
-    return config.get("configurable", {}).get("_credentials", {}).get("access_token")
 
-class Auth0AI():
-    def __init__(self, config: Optional[AuthorizerParams] = None):
-        self._graph: Optional[CIBAGraph] = None
-        self.config = config
+    def __init__(self, auth0: Optional[Auth0ClientParams] = None):
+        """Initializes the Auth0AI instance.
 
-    def with_async_user_confirmation(self, **options: CIBAGraphOptions) -> CIBAGraph:
+        Args:
+            auth0 (Optional[Auth0ClientParams]): Parameters for the Auth0 client.
+                If not provided, values will be automatically read from environment
+                variables: `AUTH0_DOMAIN`, `AUTH0_CLIENT_ID`, and `AUTH0_CLIENT_SECRET`.
         """
-        Initializes and registers a state graph for conditional trade operations using CIBA.
+        self.auth0 = auth0
 
-        Attributes:
-            options (Optional[CIBAGraphOptions]): The base CIBA options.
-        """
-        self._graph = CIBAGraph(CIBAGraphOptions(**options), self.config)
-        return self._graph
-    
-    def with_federated_connection(self, **options: FederatedConnectionAuthorizerParams) -> Callable[[BaseTool], BaseTool]:
+    def with_async_user_confirmation(self, **params: CIBAAuthorizerParams) -> Callable[[BaseTool], BaseTool]:
+        """Protects a tool with the CIBA (Client-Initiated Backchannel Authentication) flow.
+
+        Requires user confirmation via a second device (e.g., phone)
+        before allowing the tool to execute.
+
+        Args:
+            **params: Parameters defined in `CIBAAuthorizerParams`.
+
+        Returns:
+            Callable[[BaseTool], BaseTool]: A decorator to wrap a LangChain tool.
+
+        Example:
+            ```python
+            import os
+            from auth0_ai_langchain.auth0_ai import Auth0AI
+            from auth0_ai_langchain.ciba import get_ciba_credentials
+            from langchain_core.runnables import ensure_config
+            from langchain_core.tools import StructuredTool
+
+            auth0_ai = Auth0AI()
+
+            with_async_user_confirmation = auth0_ai.with_async_user_confirmation(
+                scopes=["stock:trade"],
+                audience=os.getenv("AUDIENCE"),
+                binding_message=lambda ticker, qty: f"Authorize the purchase of {qty} {ticker}",
+                user_id=lambda *_, **__: ensure_config().get("configurable", {}).get("user_id")
+            )
+
+            def tool_function(ticker: str, qty: int) -> str:
+                credentials = get_ciba_credentials()
+                headers = {
+                    "Authorization": f"{credentials['token_type']} {credentials['access_token']}",
+                    # ...
+                }
+                # Call API
+
+            trade_tool = with_async_user_confirmation(
+                StructuredTool(
+                    name="trade_tool",
+                    description="Use this function to trade a stock",
+                    func=tool_function,
+                )
+            )
+            ```
         """
-        Protects a tool execution with the Federated Connection authorizer.
+        authorizer = CIBAAuthorizer(CIBAAuthorizerParams(**params), self.auth0)
+        return authorizer.authorizer()
+
+    def with_federated_connection(self, **params: FederatedConnectionAuthorizerParams) -> Callable[[BaseTool], BaseTool]:
+        """Enables a tool to obtain an access token from a federated identity provider (e.g., Google, Azure AD).
+
+        The token can then be used within the tool to call third-party APIs on behalf of the user.
+
+        Args:
+            **params: Parameters defined in `FederatedConnectionAuthorizerParams`.
+
+        Returns:
+            Callable[[BaseTool], BaseTool]: A decorator to wrap a LangChain tool.
+
+        Example:
+            ```python
+            from auth0_ai_langchain.auth0_ai import Auth0AI
+            from auth0_ai_langchain.federated_connections import get_credentials_for_connection
+            from langchain_core.tools import StructuredTool
+            from datetime import datetime
+
+            auth0_ai = Auth0AI()
+
+            with_google_calendar_access = auth0_ai.with_federated_connection(
+                connection="google-oauth2",
+                scopes=["https://www.googleapis.com/auth/calendar.freebusy"]
+            )
+
+            def tool_function(date: datetime):
+                credentials = get_credentials_for_connection()
+                # Call Google API using credentials["access_token"]
 
-        Attributes:
-            options (FederatedConnectionAuthorizerParams): The Federated Connections authorizer options.
+            check_calendar_tool = with_google_calendar_access(
+                StructuredTool(
+                    name="check_user_calendar",
+                    description="Use this function to check if the user is available on a certain date and time",
+                    func=tool_function,
+                )
+            )
+            ```
         """
-        authorizer = FederatedConnectionAuthorizer(FederatedConnectionAuthorizerParams(**options), self.config)
+        authorizer = FederatedConnectionAuthorizer(
+            FederatedConnectionAuthorizerParams(**params), self.auth0)
         return authorizer.authorizer()

auth0_ai_langchain/ciba/__init__.py

@@ -0,0 +1,3 @@
+from auth0_ai.authorizers.ciba.ciba_authorizer_base import get_ciba_credentials as get_ciba_credentials
+from auth0_ai_langchain.ciba.ciba_authorizer import CIBAAuthorizer as CIBAAuthorizer
+from auth0_ai_langchain.ciba.graph_resumer import GraphResumer as GraphResumer

auth0_ai_langchain/ciba/ciba_authorizer.py

@@ -0,0 +1,17 @@
+from abc import ABC
+from typing import Union
+from auth0_ai.authorizers.ciba import CIBAAuthorizerBase
+from auth0_ai.interrupts.ciba_interrupts import AuthorizationPendingInterrupt, AuthorizationPollingInterrupt
+from auth0_ai_langchain.utils.interrupt import to_graph_interrupt
+from auth0_ai_langchain.utils.tool_wrapper import tool_wrapper
+from langchain_core.tools import BaseTool
+
+class CIBAAuthorizer(CIBAAuthorizerBase, ABC):
+    def _handle_authorization_interrupts(self, err: Union[AuthorizationPendingInterrupt, AuthorizationPollingInterrupt]) -> None:
+        raise to_graph_interrupt(err)
+    
+    def authorizer(self):
+        def wrap_tool(tool: BaseTool) -> BaseTool:
+            return tool_wrapper(tool, self.protect)
+        
+        return wrap_tool

auth0_ai_langchain/ciba/graph_resumer.py

@@ -0,0 +1,154 @@
+import asyncio
+from threading import Event
+from typing import Callable, Optional, Dict, Any, List, TypedDict
+from auth0_ai.authorizers.ciba import CIBAAuthorizationRequest
+from auth0_ai.interrupts.ciba_interrupts import CIBAInterrupt, AuthorizationPendingInterrupt, AuthorizationPollingInterrupt
+from auth0_ai_langchain.utils.interrupt import get_auth0_interrupts
+from langgraph_sdk.client import LangGraphClient
+from langgraph_sdk.schema import Thread, Interrupt
+
+class WatchedThread(TypedDict):
+    thread_id: str
+    assistant_id: str
+    interruption_id: str
+    auth_request: CIBAAuthorizationRequest
+    config: Dict[str, Any]
+    last_run: float
+
+class GraphResumerFilters(TypedDict):
+    graph_id: str
+
+class GraphResumer:
+    def __init__(self, lang_graph: LangGraphClient, filters: Optional[GraphResumerFilters] = None):
+        self.lang_graph = lang_graph
+        self.filters = filters or {}
+        self.map: Dict[str, WatchedThread] = {}
+        self._stop_event = Event()
+        self._loop_task: Optional[asyncio.Task] = None
+
+        # Event callbacks
+        self._resume_callbacks: List[Callable[[WatchedThread], None]] = []
+        self._error_callbacks: List[Callable[[Exception], None]] = []
+
+    # Public API to register event callbacks
+    def on_resume(self, callback: Callable[[WatchedThread], None]) -> "GraphResumer":
+        self._resume_callbacks.append(callback)
+        return self
+
+    def on_error(self, callback: Callable[[Exception], None]) -> "GraphResumer":
+        self._error_callbacks.append(callback)
+        return self
+
+    def _emit_resume(self, thread: WatchedThread) -> None:
+        for callback in self._resume_callbacks:
+            callback(thread)
+
+    def _emit_error(self, error: Exception) -> None:
+        for callback in self._error_callbacks:
+            callback(error)
+
+    async def _get_all_interrupted_threads(self) -> List[Thread]:
+        interrupted_threads: List[Thread] = []
+        offset = 0
+
+        while True:
+            page = await self.lang_graph.threads.search(
+                status="interrupted",
+                limit=100,
+                offset=offset,
+                metadata={"graph_id": self.filters["graph_id"]} if "graph_id" in self.filters else None
+            )
+
+            if not page:
+                break
+
+            for t in page:
+                interrupt = self._get_first_interrupt(t)
+                if interrupt and CIBAInterrupt.is_interrupt(interrupt["value"]) and CIBAInterrupt.has_request_data(interrupt["value"]):
+                    interrupted_threads.append(t)
+
+            offset += len(page)
+            if len(page) < 100:
+                break
+
+        return interrupted_threads
+
+    def _get_first_interrupt(self, thread: Thread) -> Optional[Interrupt]:
+        interrupts = thread["interrupts"]
+        if interrupts:
+            values = list(interrupts.values())
+            if values and values[0]:
+                return values[0][0]
+        return None
+
+    def _get_hash_map_id(self, thread: Thread) -> str:
+        return f"{thread['thread_id']}:{next(iter(thread['interrupts']))}"
+
+    async def _resume_thread(self, t: WatchedThread):
+        self._emit_resume(t)
+
+        await self.lang_graph.runs.wait(t["thread_id"], t["assistant_id"], config=t["config"])
+
+        t["last_run"] = asyncio.get_event_loop().time() * 1000
+
+    async def loop(self):
+        all_threads = await self._get_all_interrupted_threads()
+
+        # Remove old interrupted threads
+        active_keys = {self._get_hash_map_id(t) for t in all_threads}
+
+        for key in list(self.map.keys()):
+            if key not in active_keys:
+                del self.map[key]
+
+        # Add new interrupted threads
+        for thread in all_threads:
+            interrupt = next(
+                (i for i in get_auth0_interrupts(thread)
+                 if AuthorizationPendingInterrupt.is_interrupt(i["value"])
+                 or AuthorizationPollingInterrupt.is_interrupt(i["value"])),
+                None
+            )
+
+            if not interrupt or not interrupt["value"].get("_request"):
+                continue
+
+            key = self._get_hash_map_id(thread)
+            if key not in self.map:
+                self.map[key] = {
+                    "thread_id": thread["thread_id"],
+                    "assistant_id": thread["metadata"].get("graph_id"),
+                    "config": getattr(thread, "config", {}),
+                    "interruption_id": next(iter(thread["interrupts"])),
+                    "auth_request": interrupt["value"]["_request"],
+                }
+
+        threads_to_resume = [
+            t for t in self.map.values()
+            if "last_run" not in t or (t["last_run"] + t["auth_request"]["interval"] * 1000 < asyncio.get_event_loop().time() * 1000)
+        ]
+
+        await asyncio.gather(*[
+            self._resume_thread(t) for t in threads_to_resume
+        ])
+
+    def start(self):
+        if self._loop_task and not self._loop_task.done():
+            return
+
+        self._stop_event.clear()
+
+        async def _run_loop():
+            while not self._stop_event.is_set():
+                try:
+                    await self.loop()
+                except Exception as e:
+                    self._emit_error(e)
+                await asyncio.sleep(5)
+
+        self._loop_task = asyncio.create_task(_run_loop())
+
+    def stop(self):
+        self._stop_event.set()
+        if self._loop_task:
+            self._loop_task.cancel()

auth0_ai_langchain/federated_connections/__init__.py

@@ -1,4 +1,10 @@
-from auth0_ai.interrupts.federated_connection_interrupt import FederatedConnectionError as FederatedConnectionError
-from auth0_ai.interrupts.federated_connection_interrupt import FederatedConnectionInterrupt as FederatedConnectionInterrupt
-from auth0_ai.authorizers.federated_connection_authorizer import get_access_token_for_connection as get_access_token_for_connection
+from auth0_ai.interrupts.federated_connection_interrupt import (
+    FederatedConnectionError as FederatedConnectionError,
+    FederatedConnectionInterrupt as FederatedConnectionInterrupt
+)
+
+from auth0_ai.authorizers.federated_connection_authorizer import (
+    get_credentials_for_connection as get_credentials_for_connection,
+    get_access_token_for_connection as get_access_token_for_connection
+)
 from .federated_connection_authorizer import FederatedConnectionAuthorizer as FederatedConnectionAuthorizer

auth0_ai_langchain/federated_connections/federated_connection_authorizer.py

@@ -1,52 +1,33 @@
 import copy
 from abc import ABC
 from auth0_ai.authorizers.federated_connection_authorizer import FederatedConnectionAuthorizerBase, FederatedConnectionAuthorizerParams
-from auth0_ai.authorizers.types import AuthorizerParams
+from auth0_ai.authorizers.types import Auth0ClientParams
 from auth0_ai.interrupts.federated_connection_interrupt import FederatedConnectionInterrupt
-from langchain_core.tools import BaseTool, tool
+from auth0_ai_langchain.utils.interrupt import to_graph_interrupt
+from auth0_ai_langchain.utils.tool_wrapper import tool_wrapper
+from langchain_core.tools import BaseTool
 from langchain_core.runnables import ensure_config
-from ..utils.interrupt import to_graph_interrupt
 
-async def get_refresh_token(*_args, **_kwargs) -> str | None:
+async def default_get_refresh_token(*_, **__) -> str | None:
     return ensure_config().get("configurable", {}).get("_credentials", {}).get("refresh_token")
 
 class FederatedConnectionAuthorizer(FederatedConnectionAuthorizerBase, ABC):
     def __init__(
         self, 
-        options: FederatedConnectionAuthorizerParams,
-        config: AuthorizerParams = None,
+        params: FederatedConnectionAuthorizerParams,
+        auth0: Auth0ClientParams = None,
     ):
-        if options.refresh_token.value is None:
-            options = copy.copy(options)
-            options.refresh_token.value = get_refresh_token
+        if params.refresh_token.value is None:
+            params = copy.copy(params)
+            params.refresh_token.value = default_get_refresh_token
 
-        super().__init__(options, config)
+        super().__init__(params, auth0)
     
     def _handle_authorization_interrupts(self, err: FederatedConnectionInterrupt) -> None:
         raise to_graph_interrupt(err)
     
     def authorizer(self):
-        def wrapped_tool(t: BaseTool) -> BaseTool:
-            async def execute_fn(*_args, **kwargs):
-                return await t.ainvoke(input=kwargs)
-
-            tool_fn = self.protect(
-                lambda *_args, **_kwargs: {
-                    "thread_id": ensure_config().get("configurable", {}).get("thread_id"),
-                    "checkpoint_ns": ensure_config().get("configurable", {}).get("checkpoint_ns"),
-                    "run_id": ensure_config().get("configurable", {}).get("run_id"),
-                    "tool_call_id": ensure_config().get("configurable", {}).get("tool_call_id"), # TODO: review this
-                },
-                execute_fn
-            )
-            tool_fn.__name__ = t.name
-            
-            return tool(
-                tool_fn,
-                description=t.description,
-                return_direct=t.return_direct,
-                args_schema=t.args_schema,
-                response_format=t.response_format,
-            )
+        def wrap_tool(tool: BaseTool) -> BaseTool:
+            return tool_wrapper(tool, self.protect)
         
-        return wrapped_tool
+        return wrap_tool

auth0_ai_langchain/fga/__init__.py

@@ -0,0 +1,4 @@
+from auth0_ai.authorizers.fga_authorizer import (
+    FGAAuthorizer as FGAAuthorizer,
+    FGAAuthorizerOptions as FGAAuthorizerOptions
+)

auth0_ai_langchain/utils/interrupt.py

@@ -1,6 +1,9 @@
+from typing import List
 from auth0_ai.interrupts.auth0_interrupt import Auth0Interrupt
 from langgraph.errors import GraphInterrupt
 from langgraph.types import Interrupt
+from langgraph_sdk.schema import Thread
+
 
 def to_graph_interrupt(interrupt: Auth0Interrupt) -> GraphInterrupt:
     return GraphInterrupt([
@@ -8,6 +11,20 @@ def to_graph_interrupt(interrupt: Auth0Interrupt) -> GraphInterrupt:
             value=interrupt.to_json(),
             when="during",
             resumable=True,
-            ns=[f"auth0AI:{interrupt.__class__.__name__}:{interrupt.code}"]
+            ns=[f"auth0AI:{interrupt.name}:{interrupt.code}"]
         )
     ])
+
+
+def get_auth0_interrupts(thread: Thread) -> List[Interrupt]:
+    result = []
+
+    if "interrupts" not in thread:
+        return result
+
+    for interrupt_list in thread["interrupts"].values():
+        for interrupt in interrupt_list:
+            if Auth0Interrupt.is_interrupt(interrupt["value"]):
+                result.append(interrupt)
+
+    return result

auth0_ai_langchain/utils/tool_wrapper.py

@@ -0,0 +1,34 @@
+from typing import Callable
+from typing_extensions import Annotated
+from pydantic import create_model
+from langchain_core.tools import BaseTool, tool as create_tool, InjectedToolCallId
+from langchain_core.runnables import RunnableConfig
+
+def tool_wrapper(tool: BaseTool, protect_fn: Callable) -> BaseTool:
+
+    # Workaround: extend existing args_schema to be able to get the tool_call_id value
+    args_schema = create_model(
+        tool.args_schema.__name__ + "Extended",
+        __base__=tool.args_schema,
+        **{"tool_call_id": (Annotated[str, InjectedToolCallId])}
+    )
+
+    @create_tool(
+        tool.name,
+        description=tool.description,
+        args_schema=args_schema
+    )
+    async def wrapped_tool(config: RunnableConfig, tool_call_id: Annotated[str, InjectedToolCallId], **input):
+        async def execute_fn(*_, **__):
+            return await tool.ainvoke(input, config)
+
+        return await protect_fn(
+            lambda *_, **__: {
+                "thread_id": config.get("configurable", {}).get("thread_id"),
+                "tool_call_id": tool_call_id,
+                "tool_name": tool.name,
+            },
+            execute_fn,
+        )(**input)
+
+    return wrapped_tool