use app state to store global vars

there is only one engine, sessionmaker, azure_credentials, context, chat_client, and embed_client created during the lifespan of the fastapi app

Author: john0isaac

src/fastapi_app/__init__.py

@@ -1,12 +1,46 @@
 import logging
 import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TypedDict
 
 from dotenv import load_dotenv
 from fastapi import FastAPI
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from fastapi_app.dependencies import (
+    FastAPIAppContext,
+    common_parameters,
+    create_async_sessionmaker,
+    get_azure_credentials,
+)
+from fastapi_app.openai_clients import create_openai_chat_client, create_openai_embed_client
+from fastapi_app.postgres_engine import create_postgres_engine_from_env
 
 logger = logging.getLogger("ragapp")
 
 
+class State(TypedDict):
+    sessionmaker: async_sessionmaker[AsyncSession]
+    context: FastAPIAppContext
+    chat_client: AsyncOpenAI | AsyncAzureOpenAI
+    embed_client: AsyncOpenAI | AsyncAzureOpenAI
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator[State]:
+    context = await common_parameters()
+    azure_credential = await get_azure_credentials()
+    engine = await create_postgres_engine_from_env(azure_credential)
+    sessionmaker = await create_async_sessionmaker(engine)
+    chat_client = await create_openai_chat_client(azure_credential)
+    embed_client = await create_openai_embed_client(azure_credential)
+
+    yield {"sessionmaker": sessionmaker, "context": context, "chat_client": chat_client, "embed_client": embed_client}
+    await engine.dispose()
+
+
 def create_app(testing: bool = False):
     if os.getenv("RUNNING_IN_PRODUCTION"):
         logging.basicConfig(level=logging.WARNING)
@@ -15,7 +49,7 @@ def create_app(testing: bool = False):
             load_dotenv(override=True)
         logging.basicConfig(level=logging.INFO)
 
-    app = FastAPI(docs_url="/docs")
+    app = FastAPI(docs_url="/docs", lifespan=lifespan)
 
     from fastapi_app.routes import api_routes, frontend_routes
 

src/fastapi_app/dependencies.py

@@ -1,16 +1,14 @@
 import logging
 import os
+from collections.abc import AsyncGenerator
 from typing import Annotated
 
 import azure.identity
-from fastapi import Depends
+from fastapi import Depends, Request
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 from pydantic import BaseModel
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
 
-from fastapi_app.openai_clients import create_openai_chat_client, create_openai_embed_client
-from fastapi_app.postgres_engine import create_postgres_engine_from_env
-
 logger = logging.getLogger("ragapp")
 
 
@@ -67,7 +65,7 @@ async def common_parameters():
     )
 
 
-def get_azure_credentials() -> azure.identity.DefaultAzureCredential | azure.identity.ManagedIdentityCredential:
+async def get_azure_credentials() -> azure.identity.DefaultAzureCredential | azure.identity.ManagedIdentityCredential:
     azure_credential: azure.identity.DefaultAzureCredential | azure.identity.ManagedIdentityCredential
     try:
         if client_id := os.getenv("APP_IDENTITY_ID"):
@@ -86,35 +84,55 @@ def get_azure_credentials() -> azure.identity.DefaultAzureCredential | azure.ide
         raise e
 
 
-azure_credentials = get_azure_credentials()
+async def create_async_sessionmaker(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
+    """Get the agent database"""
+    return async_sessionmaker(
+        engine,
+        expire_on_commit=False,
+        autoflush=False,
+    )
 
 
-async def get_engine():
-    """Get the agent database engine"""
-    engine = await create_postgres_engine_from_env(azure_credentials)
-    return engine
+async def get_async_sessionmaker(
+    request: Request,
+) -> AsyncGenerator[async_sessionmaker[AsyncSession], None]:
+    yield request.state.sessionmaker
 
 
-async def get_async_session(engine: Annotated[AsyncEngine, Depends(get_engine)]):
-    """Get the agent database"""
-    async_session_maker = async_sessionmaker(engine, expire_on_commit=False)
-    async with async_session_maker() as async_session:
-        yield async_session
+async def get_context(
+    request: Request,
+) -> FastAPIAppContext:
+    return request.state.context
+
+
+async def get_async_db_session(
+    sessionmaker: Annotated[async_sessionmaker[AsyncSession], Depends(get_async_sessionmaker)],
+) -> AsyncGenerator[AsyncSession, None]:
+    async with sessionmaker() as session:
+        try:
+            yield session
+        except:
+            await session.rollback()
+            raise
+        else:
+            await session.commit()
 
 
-async def get_openai_chat_client():
+async def get_openai_chat_client(
+    request: Request,
+) -> OpenAIClient:
     """Get the OpenAI chat client"""
-    chat_client = await create_openai_chat_client(azure_credentials)
-    return OpenAIClient(client=chat_client)
+    return OpenAIClient(client=request.state.chat_client)
 
 
-async def get_openai_embed_client():
+async def get_openai_embed_client(
+    request: Request,
+) -> OpenAIClient:
     """Get the OpenAI embed client"""
-    embed_client = await create_openai_embed_client(azure_credentials)
-    return OpenAIClient(client=embed_client)
+    return OpenAIClient(client=request.state.embed_client)
 
 
-CommonDeps = Annotated[FastAPIAppContext, Depends(common_parameters)]
-DBSession = Annotated[AsyncSession, Depends(get_async_session)]
+CommonDeps = Annotated[FastAPIAppContext, Depends(get_context)]
+DBSession = Annotated[AsyncSession, Depends(get_async_db_session)]
 ChatClient = Annotated[OpenAIClient, Depends(get_openai_chat_client)]
 EmbeddingsClient = Annotated[OpenAIClient, Depends(get_openai_embed_client)]

src/fastapi_app/update_embeddings.py

@@ -4,14 +4,17 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
 
-from fastapi_app.dependencies import common_parameters, get_engine, get_openai_embed_client
+from fastapi_app.dependencies import common_parameters, get_azure_credentials
 from fastapi_app.embeddings import compute_text_embedding
+from fastapi_app.openai_clients import create_openai_embed_client
+from fastapi_app.postgres_engine import create_postgres_engine_from_env
 from fastapi_app.postgres_models import Item
 
 
 async def update_embeddings():
-    engine = await get_engine()
-    openai_embed = await get_openai_embed_client()
+    azure_credential = await get_azure_credentials()
+    engine = await create_postgres_engine_from_env(azure_credential)
+    openai_embed_client = await create_openai_embed_client(azure_credential)
     common_params = await common_parameters()
 
     async with async_sessionmaker(engine, expire_on_commit=False)() as session:
@@ -21,7 +24,7 @@ async def update_embeddings():
             for item in items:
                 item.embedding = await compute_text_embedding(
                     item.to_str_for_embedding(),
-                    openai_client=openai_embed.client,
+                    openai_client=openai_embed_client,
                     embed_model=common_params.openai_embed_model,
                     embedding_dimensions=common_params.openai_embed_dimensions,
                 )