Repository: aipotheosis-labs/aci
File: aipolabs/server/routes/analytics.py (Python)

Code before:
from datetime import datetime, timedelta

from fastapi import APIRouter
from logfire.experimental.query_client import AsyncLogfireQueryClient

from aipolabs.common.logging_setup import get_logger
from aipolabs.common.schemas.analytics import DistributionDatapoint, TimeSeriesDatapoint
from aipolabs.server import config

router = APIRouter()
logger = get_logger(__name__)


@router.get("/app-usage-distribution", response_model=list[DistributionDatapoint])
async def get_app_usage_distribution() -> list[DistributionDatapoint]:
    query = """
SELECT
  regexp_replace(url_path, '/v1/functions/([^/]+?)(?:__.*)?/execute', '\\1') AS name,
  COUNT(*) AS value
FROM records
WHERE attributes->>'http.user_agent' LIKE '%python%'
  AND attributes->>'fastapi.route.name' = 'execute'
  AND trace_id IN (SELECT trace_id FROM records
WHERE attributes->>'api_key_id' = 'af24364c-30b5-4050-b1b0-92a0f904c50b')
GROUP BY name
ORDER BY value DESC;
    """

    async with AsyncLogfireQueryClient(read_token=config.LOGFIRE_READ_TOKEN) as client:
        json_rows = await client.query_json_rows(
            sql=query, min_timestamp=datetime.now() - timedelta(days=7)
        )
        print(json_rows)
        return [DistributionDatapoint(**row) for row in json_rows["rows"]]


@router.get("/function-usage-distribution", response_model=list[DistributionDatapoint])
async def get_function_usage_distribution() -> list[DistributionDatapoint]:
    query = """
SELECT
  regexp_replace(url_path, '/v1/functions/([A-Z0-9_]+)/execute', '\\1') AS name,
  COUNT(*) AS value
FROM records
WHERE attributes->>'http.user_agent' LIKE '%python%'
  AND attributes->>'fastapi.route.name' = 'execute'
  AND trace_id IN (SELECT trace_id FROM records
WHERE attributes->>'api_key_id' = 'af24364c-30b5-4050-b1b0-92a0f904c50b')
GROUP BY name
ORDER BY value DESC;

Review comment (security):
you probably already know this but raise it for awareness: using 'api_key_id' = 'xx' will only result the usage for that particular api_key/agent. So we probably want to allow user to select (dropdown or filter) which api_key/agent's stats to show.

We can just return usage_distribution for all apikeys/agents (we can get the project from the requesting api_key and get all api_key under that project) and do filter in frontend. With this approach, we can reduce the refactor work after we integrate with propelauth (where frontend no longer use agent's apikey to access backend)

Code after:
from datetime import datetime, timedelta
from typing import Annotated

from fastapi import APIRouter, Depends
from logfire.experimental.query_client import AsyncLogfireQueryClient

from aipolabs.common.db import crud
from aipolabs.common.logging_setup import get_logger
from aipolabs.common.schemas.analytics import DistributionDatapoint, TimeSeriesDatapoint
from aipolabs.server import config
from aipolabs.server import dependencies as deps

router = APIRouter()
logger = get_logger(__name__)


def _get_project_api_key_ids_sql_list(context: deps.RequestContext) -> str | None:
    project_api_key_ids = crud.projects.get_all_api_key_ids_for_project(
        context.db_session, context.project.id
    )

    if not project_api_key_ids:
        return None

    return ",".join(f"'{key_id}'" for key_id in project_api_key_ids)


@router.get("/app-usage-distribution", response_model=list[DistributionDatapoint])
async def get_app_usage_distribution(
    context: Annotated[deps.RequestContext, Depends(deps.get_request_context)],
) -> list[DistributionDatapoint]:
    api_key_ids_sql_list = _get_project_api_key_ids_sql_list(context)

    if not api_key_ids_sql_list:
        return []

    query = f"""
SELECT
  regexp_replace(url_path, '/v1/functions/([^/]+?)(?:__.*)?/execute', '\\1') AS name,
  COUNT(*) AS value
FROM records
WHERE attributes->>'http.user_agent' LIKE '%python%'
  AND attributes->>'fastapi.route.name' = 'execute'
  AND trace_id IN (SELECT trace_id FROM records
WHERE attributes->>'api_key_id' IN ({api_key_ids_sql_list}))
GROUP BY name
ORDER BY value DESC;
    """