chore(examples): minor improvements to async-user-confirmation and dependencies updates (#36)

Author: cristiandouce

.vscode/settings.json

@@ -1,18 +1,16 @@
 {
   "poetryMonorepo.appendExtraPaths": true,
   "python.analysis.extraPaths": [
+    "examples/async-user-confirmation/langchain-examples",
+    "examples/async-user-confirmation/llama-index-examples",
+    "examples/async-user-confirmation/sample-api",
+    "examples/authorization-for-tools/langchain-examples",
+    "examples/authorization-for-tools/llama-index-examples",
     "examples/calling-apis/langchain-examples",
     "examples/calling-apis/llama-index-examples",
-    "examples/async-user-confirmation/langchain-examples/src",
-    "examples/async-user-confirmation/llama-index-examples/src",
-    "examples/async-user-confirmation/sample-api",
-    "examples/authorization-for-rag/langchain-examples/langchain_rag",
-    "examples/authorization-for-rag/llama-index-examples/llama_index_rag",
-    "examples/authorization-for-tools/langchain-examples/src",
-    "examples/authorization-for-tools/llama-index-examples/src",
-    "packages/auth0-ai/auth0_ai/auth0_ai",
     "packages/auth0-ai-langchain/auth0_ai_langchain/auth0_ai_langchain",
-    "packages/auth0-ai-llamaindex/auth0_ai_llamaindex/auth0_ai_llamaindex"
+    "packages/auth0-ai-llamaindex/auth0_ai_llamaindex/auth0_ai_llamaindex",
+    "packages/auth0-ai/auth0_ai/auth0_ai"
   ],
   "python.testing.pytestArgs": [
     "packages/auth0-ai-langchain/tests",

examples/async-user-confirmation/langchain-examples/poetry.lock

@@ -20,13 +20,13 @@ httpx = "^0.28.1"
 python-dotenv = "^1.0.1"
 pydantic = "^2.10.6"
 nanoid = "^2.0.0"
-langchain-openai = "^0.3.16"
 langchain = "^0.3.25"
 langchain-core = "^0.3.59"
+langchain-openai = "^0.3.16"
 langgraph = "^0.4.3"
 langgraph-cli = "^0.2.8"
-langgraph-api = "^0.2.17"
-langgraph-runtime-inmem = "^0.0.9"
+langgraph-api = "^0.2.19"
+langgraph-runtime-inmem = "^0.0.10"
 auth0-ai-langchain = { path = "../../../packages/auth0-ai-langchain", develop = true }
 
 [build-system]

examples/async-user-confirmation/langchain-examples/pyproject.toml

@@ -3,9 +3,9 @@
 from langgraph.prebuilt import ToolNode
 from langchain_core.messages import BaseMessage, AIMessage
 from typing import TypedDict, Annotated, Sequence
-from tools.trade import trade_tool
-from tools.conditional_trade import conditional_trade_tool
 from langchain_openai import ChatOpenAI
+from src.agents.tools.trade import trade_tool
+from src.agents.tools.conditional_trade import conditional_trade_tool
 
 class State(TypedDict):
     messages: Annotated[Sequence[BaseMessage], add_messages]

examples/async-user-confirmation/langchain-examples/src/agents/agent.py

@@ -10,7 +10,7 @@
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
 from src.agents.clients.scheduler import SchedulerClient
-from tools.trade import trade_tool
+from src.agents.tools.trade import trade_tool
 
 class ConditionalTrade(TypedDict):
     ticker: str

examples/async-user-confirmation/langchain-examples/src/agents/conditional_trade.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from langchain_core.tools import StructuredTool
 from langchain_core.runnables.config import RunnableConfig
-from clients.scheduler import SchedulerClient
+from src.agents.clients.scheduler import SchedulerClient
 
 class MetricEnum(str, Enum):
     PE = "P/E"

examples/async-user-confirmation/langchain-examples/src/agents/tools/conditional_trade.py

@@ -15,11 +15,11 @@ python = "^3.11"
 python-dotenv = "^1.0.1"
 flask = { extras = ["async"], version = "^3.1.0" }
 requests = "^2.32.3"
-authlib = "^1.5.1"
 hypercorn = "^0.17.3"
 llama-index = "^0.12.33"
 llama-index-agent-openai = "^0.4.6"
 auth0-ai-llamaindex = { path = "../../../packages/auth0-ai-llamaindex", develop = true }
+auth0-server-python = "^1.0.0b4"
 
 [build-system]
 requires = ["poetry-core"]

examples/async-user-confirmation/llama-index-examples/poetry.lock

@@ -1,14 +1,14 @@
 import asyncio
+from asgiref.wsgi import WsgiToAsgi
 from hypercorn.config import Config
 from hypercorn.asyncio import serve
 from src.app.app import app
 
+
 def main():
     config = Config()
     config.bind = ["0.0.0.0:3000"]
     config.worker_class = "asyncio"
     config.use_reloader = True
 
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    loop.run_until_complete(serve(app, config))
+    asyncio.run(serve(WsgiToAsgi(app), config))

examples/async-user-confirmation/llama-index-examples/pyproject.toml

@@ -0,0 +1,29 @@
+from datetime import datetime
+from llama_index.agent.openai import OpenAIAgent
+
+from .memory import get_memory
+from .tools.trade import trade_tool
+
+
+system_prompt = f"""You are a specialized stock trading assistant designed to
+guide users through the process of buying stocks step by step.
+
+**Important Constraints**:
+- You cannot discuss, buy, or sell any stocks outside this limited list, whether real or fictional.
+- You and the user can discuss the prices of these stocks, adjust stock amounts, and place buy orders through the UI.
+
+**Additional Guidelines**:
+- Today’s date for reference: {datetime.now().isoformat()}
+- You may perform calculations as needed and engage in general discussion with the user.
+"""
+
+
+async def get_agent(user_id: str, chat_id: str):
+    chat_memory = await get_memory(user_id, chat_id)
+    return OpenAIAgent.from_tools(
+        model="gpt-4o",
+        memory=chat_memory,
+        system_prompt=system_prompt,
+        tools=[trade_tool],
+        verbose=True
+    )

examples/async-user-confirmation/llama-index-examples/scripts/start.py

@@ -0,0 +1,12 @@
+from llama_index.core.storage.chat_store import SimpleChatStore
+from llama_index.core.memory import ChatMemoryBuffer
+
+chat_store = SimpleChatStore()
+
+
+async def get_memory(user_id: str, chat_id: str):
+    return ChatMemoryBuffer.from_defaults(
+        token_limit=3000,
+        chat_store=chat_store,
+        chat_store_key=f"{user_id}_{chat_id}",
+    )

examples/async-user-confirmation/llama-index-examples/src/agents/agent.py

@@ -1,9 +1,12 @@
 import os
-
+from dotenv import load_dotenv
 import httpx
-from auth0_ai_llamaindex.ciba import get_ciba_credentials
 from llama_index.core.tools import FunctionTool
 
+from auth0_ai_llamaindex.ciba import get_ciba_credentials
+from ...auth0.auth0_ai import with_async_user_confirmation
+load_dotenv()
+
 
 def trade_tool_function(ticker: str, qty: int) -> str:
     credentials = get_ciba_credentials()
@@ -29,8 +32,8 @@ def trade_tool_function(ticker: str, qty: int) -> str:
         return f"HTTP request failed: {str(e)}"
 
 
-trade_tool = FunctionTool.from_defaults(
+trade_tool = with_async_user_confirmation(FunctionTool.from_defaults(
     name="trade_tool",
     description="Use this function to trade a stock",
     fn=trade_tool_function,
-)
+))

examples/async-user-confirmation/llama-index-examples/src/agents/memory.py

@@ -1,117 +1,67 @@
 import os
 import uuid
-from datetime import datetime
-from urllib.parse import quote_plus, urlencode
-
-from auth0_ai_llamaindex.auth0_ai import Auth0AI, set_ai_context
-from authlib.integrations.flask_client import OAuth
 from dotenv import load_dotenv
 from flask import Flask, jsonify, redirect, render_template, request, session, url_for
-from llama_index.agent.openai import OpenAIAgent
 
-from ..tools.trade import trade_tool
+from auth0_ai_llamaindex.auth0_ai import set_ai_context
 
-load_dotenv()
+from ..agents.agent import get_agent
+from ..agents.memory import get_memory
+from ..auth0.routes import login_bp
 
-agents = {}
-system_prompt = f"""You are a specialized stock trading assistant designed to
-guide users through the process of buying stocks step by step.
 
-**Important Constraints**:
-- You cannot discuss, buy, or sell any stocks outside this limited list, whether real or fictional.
-- You and the user can discuss the prices of these stocks, adjust stock amounts, and place buy orders through the UI.
+load_dotenv()
 
-**Additional Guidelines**:
-- Today’s date for reference: {datetime.now().isoformat()}
-- You may perform calculations as needed and engage in general discussion with the user.
-"""
+app = Flask(__name__)
+app.secret_key = os.getenv("APP_SECRET_KEY", "SOME_RANDOM_SECRET_KEY")
 
-auth0_ai = Auth0AI()
-with_async_user_confirmation = auth0_ai.with_async_user_confirmation(
-    scope="stock:trade",
-    audience=os.getenv("AUDIENCE"),
-    binding_message=lambda ticker, qty: f"Authorize the purchase of {qty} {ticker}",
-    user_id=lambda *_, **__: session["user"]["userinfo"]["sub"],
-    # When this flag is set to `"block"`, the execution of the tool awaits until the user approves or rejects the request.
-    # Given the asynchronous nature of the CIBA flow, this mode is only useful during development.
-    on_authorization_request="block",
-)
+app.register_blueprint(login_bp)
 
-tools = [with_async_user_confirmation(trade_tool)]
 
+@app.route("/")
+async def home():
+    if "user" not in session:
+        return redirect(url_for("auth0.login", _external=True))
 
-def get_agent():
-    user_id = session["user"]["userinfo"]["sub"]
-    if user_id not in agents:
-        agents[user_id] = OpenAIAgent.from_tools(
-            tools=tools, model="gpt-4o", system_prompt=system_prompt, verbose=True)
-    return agents[user_id]
+    session["thread_id"] = str(uuid.uuid4())
+    return redirect(url_for("chat", thread_id=session["thread_id"], _external=True))
 
 
-app = Flask(__name__)
-app.secret_key = os.getenv("APP_SECRET_KEY", "YOUR_SECRET_KEY")
+@app.route("/chat/<thread_id>")
+async def chat(thread_id: str):
+    if "user" not in session:
+        return "please login", 401
 
-oauth = OAuth(app)
-oauth.register(
-    "auth0",
-    client_id=os.getenv("AUTH0_CLIENT_ID"),
-    client_secret=os.getenv("AUTH0_CLIENT_SECRET"),
-    client_kwargs={
-        "scope": "openid profile",
-    },
-    server_metadata_url=f'https://{os.getenv("AUTH0_DOMAIN")}/.well-known/openid-configuration'
-)
+    if ("thread_id" not in session) or (session["thread_id"] != thread_id):
+        return "Invalid or mismatched chat session. Please start a new chat.", 400
 
+    user_id = session["user"]["sub"]
 
-@app.route("/")
-def home():
-    if "user" not in session:
-        return redirect("/login")
-    
-    thread_id = str(uuid.uuid4())
-    set_ai_context(thread_id)
+    memory = await get_memory(user_id, thread_id)
+
+    messages = [
+        {
+            "role": m.role,
+            "content": m.content
+        } for m in memory.get_all()
+    ]
 
-    return render_template("index.html", user=session.get('user'))
+    return render_template("index.html", user=session["user"], messages=messages, interrupt=None)
 
 
-@app.route("/chat", methods=["POST"])
-async def chat():
+@app.route("/api/chat", methods=["POST"])
+async def api_chat():
     if "user" not in session:
         return jsonify({"error": "unauthorized"}), 401
 
+    user_id = session["user"]["sub"]
+    thread_id = session["thread_id"]
+    set_ai_context(thread_id)
+
     try:
         message = request.json.get("message")
-        response = await get_agent().achat(message)
+        agent = await get_agent(user_id, thread_id)
+        response = await agent.achat(message)
         return jsonify({"response": str(response)})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
-
-
-@app.route("/login")
-def login():
-    return oauth.auth0.authorize_redirect(
-        redirect_uri=url_for("login_callback", _external=True)
-    )
-
-
-@app.route("/login/callback", methods=["GET", "POST"])
-def login_callback():
-    token = oauth.auth0.authorize_access_token()
-    session["user"] = token
-    return redirect("/")
-
-
-@app.route("/logout")
-def logout():
-    session.clear()
-    return redirect(
-        "https://" + os.getenv("AUTH0_DOMAIN")
-        + "/v2/logout?"
-        + urlencode(
-            {
-                "returnTo": url_for("home", _external=True),
-                "client_id": os.getenv("AUTH0_CLIENT_ID"),
-            },
-            quote_via=quote_plus,
-        )
-    )