Add opentelemetry instrumentions

Author: chatrtham

src/fastapi_app/__init__.py

@@ -3,9 +3,12 @@
 import os
 
 import azure.identity.aio
+import fastapi
+from azure.monitor.opentelemetry import configure_azure_monitor
 from dotenv import load_dotenv
 from environs import Env
-from fastapi import FastAPI
+from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 
 from .globals import global_storage
 from .openai_clients import create_openai_chat_client
@@ -15,7 +18,7 @@
 
 
 @contextlib.asynccontextmanager
-async def lifespan(app: FastAPI):
+async def lifespan(app: fastapi.FastAPI):
     load_dotenv(override=True)
 
     azure_credential = None
@@ -40,6 +43,8 @@ async def lifespan(app: FastAPI):
     global_storage.openai_chat_client = openai_chat_client
     global_storage.openai_chat_model = openai_chat_model
 
+    if os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING"):
+        SQLAlchemyInstrumentor().instrument(engine=engine.sync_engine)
     yield
 
     await engine.dispose()
@@ -54,7 +59,16 @@ def create_app():
     else:
         logging.basicConfig(level=logging.WARNING)
 
-    app = FastAPI(docs_url="/docs", lifespan=lifespan)
+    # Turn off particularly noisy INFO level logs from Azure Core SDK:
+    logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(logging.WARNING)
+
+    if os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING"):
+        logger.info("Configuring Azure Monitor")
+        configure_azure_monitor(logger_name="ragapp")
+        # OpenAI SDK requests use httpx and are thus not auto-instrumented:
+        OpenAIInstrumentor().instrument()
+
+    app = fastapi.FastAPI(docs_url="/docs", lifespan=lifespan)
 
     from . import api_routes  # noqa
     from . import frontend_routes  # noqa

src/fastapi_app/dependencies.py

@@ -0,0 +1,118 @@
+import logging
+import os
+from collections.abc import AsyncGenerator
+from typing import Annotated
+
+import azure.identity
+from fastapi import Depends, Request
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+
+logger = logging.getLogger("ragapp")
+
+
+class OpenAIClient(BaseModel):
+    """
+    OpenAI client
+    """
+
+    client: AsyncOpenAI | AsyncAzureOpenAI
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class FastAPIAppContext(BaseModel):
+    """
+    Context for the FastAPI app
+    """
+
+    openai_chat_model: str
+    openai_chat_deployment: str
+
+
+async def common_parameters():
+    """
+    Get the common parameters for the FastAPI app
+    """
+    OPENAI_CHAT_HOST = os.getenv("OPENAI_CHAT_HOST")
+
+    if OPENAI_CHAT_HOST == "azure":
+        openai_chat_deployment = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT", "gpt-35-turbo")
+        openai_chat_model = os.getenv("AZURE_OPENAI_CHAT_MODEL", "gpt-35-turbo")
+    elif OPENAI_CHAT_HOST == "ollama":
+        openai_chat_deployment = "phi3:3.8b"
+        openai_chat_model = os.getenv("OLLAMA_CHAT_MODEL", "phi3:3.8b")
+    else:
+        openai_chat_deployment = "gpt-3.5-turbo"
+        openai_chat_model = os.getenv("OPENAICOM_CHAT_MODEL", "gpt-3.5-turbo")
+    return FastAPIAppContext(
+        openai_chat_model=openai_chat_model,
+        openai_chat_deployment=openai_chat_deployment,
+    )
+
+
+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"):
+            # Authenticate using a user-assigned managed identity on Azure
+            # See web.bicep for value of APP_IDENTITY_ID
+            logger.info(
+                "Using managed identity for client ID %s",
+                client_id,
+            )
+            azure_credential = azure.identity.ManagedIdentityCredential(client_id=client_id)
+        else:
+            azure_credential = azure.identity.DefaultAzureCredential()
+        return azure_credential
+    except Exception as e:
+        logger.warning("Failed to authenticate to Azure: %s", e)
+        raise e
+
+
+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_async_sessionmaker(
+    request: Request,
+) -> AsyncGenerator[async_sessionmaker[AsyncSession], None]:
+    yield request.state.sessionmaker
+
+
+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:
+        yield session
+
+
+async def get_openai_chat_client(
+    request: Request,
+) -> OpenAIClient:
+    """Get the OpenAI chat client"""
+    return OpenAIClient(client=request.state.chat_client)
+
+
+async def get_openai_embed_client(
+    request: Request,
+) -> OpenAIClient:
+    """Get the OpenAI embed client"""
+    return OpenAIClient(client=request.state.embed_client)
+
+
+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/pyproject.toml

@@ -15,7 +15,8 @@ dependencies = [
     "pgvector",
     "openai",
     "tiktoken",
-    "openai-messages-token-helper"
+    "openai-messages-token-helper",
+    "opentelemetry-instrumentation-fastapi>=0.46b0,<1.0.0",
 ]
 
 [build-system]

src/requirements.txt

@@ -3,28 +3,39 @@ aiosignal==1.3.1
 annotated-types==0.7.0
 anyio==4.4.0
 asgiref==3.8.1
+asttokens==2.4.1
 async-timeout==4.0.3
 asyncpg==0.29.0
 attrs==23.2.0
 azure-core==1.30.1
+azure-core-tracing-opentelemetry==1.0.0b11
 azure-identity==1.16.0
+azure-monitor-opentelemetry==1.6.0
 azure-monitor-opentelemetry-exporter==1.0.0b27
+backcall==0.2.0
+beautifulsoup4==4.12.3
+bleach==6.1.0
 build==1.2.1
 certifi==2024.2.2
 cffi==1.16.0
 cfgv==3.4.0
 charset-normalizer==3.3.2
 click==8.1.7
 cryptography==42.0.7
+decorator==5.1.1
+defusedxml==0.7.1
 Deprecated==1.2.14
 distlib==0.3.8
 distro==1.9.0
 dnspython==2.6.1
+docopt==0.6.2
 email_validator==2.1.1
 environs==11.0.0
+executing==2.0.1
 fastapi==0.111.0
 fastapi-cli==0.0.4
--e git+https://github.com/chatrtham/rag-postgres-openai-python.git@7f1578f2d48381331570f68cc13167c3290ee097#egg=fastapi_app&subdirectory=src
+-e git+https://github.com/chatrtham/rag-postgres-openai-python.git@20d8fe30d4373e9290dbd18c012f6610393fd5e7#egg=fastapi_app&subdirectory=src
+fastjsonschema==2.20.0
 filelock==3.14.0
 fixedint==0.1.6
 frozenlist==1.4.1
@@ -41,16 +52,28 @@ httpx==0.27.0
 identify==2.5.36
 idna==3.7
 importlib_metadata==7.1.0
+ipython==8.12.3
 isodate==0.6.1
+jedi==0.19.1
 Jinja2==3.1.4
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+jupyter_client==8.6.2
+jupyter_core==5.7.2
+jupyterlab_pygments==0.3.0
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 marshmallow==3.21.2
+matplotlib-inline==0.1.7
 mdurl==0.1.2
+mistune==3.0.2
 msal==1.28.0
 msal-extensions==1.1.0
 msrest==0.7.1
 multidict==6.0.5
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
 nodeenv==1.9.1
 numpy==1.26.4
 oauthlib==3.2.2
@@ -62,25 +85,43 @@ opentelemetry-exporter-otlp-proto-common==1.25.0
 opentelemetry-exporter-otlp-proto-grpc==1.25.0
 opentelemetry-exporter-otlp-proto-http==1.25.0
 opentelemetry-instrumentation==0.46b0
-opentelemetry-instrumentation-asgi==0.46b0
-opentelemetry-instrumentation-fastapi==0.46b0
-opentelemetry-instrumentation-requests==0.46b0
+opentelemetry-instrumentation-asgi==0.47b0
+opentelemetry-instrumentation-dbapi==0.47b0
+opentelemetry-instrumentation-django==0.47b0
+opentelemetry-instrumentation-fastapi==0.47b0
+opentelemetry-instrumentation-flask==0.47b0
+opentelemetry-instrumentation-openai==0.26.0
+opentelemetry-instrumentation-psycopg2==0.47b0
+opentelemetry-instrumentation-requests==0.47b0
 opentelemetry-instrumentation-sqlalchemy==0.46b0
+opentelemetry-instrumentation-urllib==0.47b0
+opentelemetry-instrumentation-urllib3==0.47b0
+opentelemetry-instrumentation-wsgi==0.47b0
 opentelemetry-proto==1.25.0
+opentelemetry-resource-detector-azure==0.1.5
 opentelemetry-sdk==1.25.0
 opentelemetry-semantic-conventions==0.46b0
-opentelemetry-util-http==0.46b0
+opentelemetry-semantic-conventions-ai==0.3.6
+opentelemetry-util-http==0.47b0
 orjson==3.10.3
 packaging==24.0
 pandas==2.2.2
+pandocfilters==1.5.1
+parso==0.8.4
+pexpect==4.9.0
 pgvector==0.2.5
+pickleshare==0.7.5
 pillow==10.3.0
 pip-tools==7.4.1
+pipreqs==0.5.0
 platformdirs==4.2.2
 portalocker==2.8.2
 pre-commit==3.7.1
+prompt_toolkit==3.0.47
 protobuf==4.25.3
 psutil==5.9.8
+ptyprocess==0.7.0
+pure_eval==0.2.3
 pycparser==2.22
 pydantic==2.7.2
 pydantic_core==2.18.3
@@ -92,20 +133,28 @@ python-dotenv==1.0.1
 python-multipart==0.0.9
 pytz==2024.1
 PyYAML==6.0.1
+pyzmq==26.0.3
+referencing==0.35.1
 regex==2024.5.15
 requests==2.32.3
 requests-oauthlib==2.0.0
 rich==13.7.1
+rpds-py==0.19.1
 ruff==0.4.8
 shellingham==1.5.4
 six==1.16.0
 smmap==5.0.1
 sniffio==1.3.1
+soupsieve==2.5
 SQLAlchemy==2.0.30
+stack-data==0.6.3
 starlette==0.37.2
 tenacity==8.4.1
 tiktoken==0.7.0
+tinycss2==1.3.0
+tornado==6.4.1
 tqdm==4.66.4
+traitlets==5.14.3
 typer==0.12.3
 typing_extensions==4.12.0
 tzdata==2024.1
@@ -115,7 +164,10 @@ uvicorn==0.30.0
 uvloop==0.19.0
 virtualenv==20.26.2
 watchfiles==0.22.0
+wcwidth==0.2.13
+webencodings==0.5.1
 websockets==12.0
 wrapt==1.16.0
+yarg==0.1.9
 yarl==1.9.4
 zipp==3.19.2