feat: improved logging

aidial_analytics_realtime/analytics.py

@@ -1,8 +1,7 @@
+import logging
 from datetime import datetime
 from decimal import Decimal
 from enum import Enum
-from logging import Logger
-from typing import Awaitable, Callable
 from uuid import uuid4
 
 from influxdb_client import Point
@@ -14,11 +13,13 @@
     get_chat_completion_response_contents,
     get_embeddings_request_contents,
 )
+from aidial_analytics_realtime.influx_writer import InfluxWriterAsync
 from aidial_analytics_realtime.rates import RatesCalculator
 from aidial_analytics_realtime.topic_model import TopicModel
 from aidial_analytics_realtime.utils.concurrency import (
     run_in_cpu_tasks_executor,
 )
+from aidial_analytics_realtime.utils.timer import Timer
 
 identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)
 
@@ -29,19 +30,15 @@ class RequestType(Enum):
 
 
 async def detect_lang(
-    logger: Logger, request: dict, response: dict, request_type: RequestType
+    request: dict, response: dict, request_type: RequestType
 ) -> str:
     match request_type:
         case RequestType.CHAT_COMPLETION:
-            request_contents = get_chat_completion_request_contents(
-                logger, request
-            )
-            response_content = get_chat_completion_response_contents(
-                logger, response
-            )
+            request_contents = get_chat_completion_request_contents(request)
+            response_content = get_chat_completion_response_contents(response)
             text = "\n\n".join(request_contents[-1:] + response_content)
         case RequestType.EMBEDDING:
-            text = "\n\n".join(get_embeddings_request_contents(logger, request))
+            text = "\n\n".join(get_embeddings_request_contents(request))
         case _:
             assert_never(request_type)
 
@@ -54,12 +51,18 @@ async def detect_lang_by_text(text: str) -> str | None:
     if not text:
         return None
 
+    logger = logging.getLogger("app.langid")
+
     try:
-        lang, prob = await run_in_cpu_tasks_executor(identifier.classify, text)
+        with Timer(logger.debug):
+            lang, prob = await run_in_cpu_tasks_executor(
+                identifier.classify, text
+            )
+
         if prob > 0.998:
             return lang
-    except Exception:
-        pass
+    except Exception as e:
+        logger.error(f"error: {str(e)}")
 
     return None
 
@@ -73,7 +76,6 @@ def build_execution_path(path: list | None):
 
 
 async def make_point(
-    logger: Logger,
     deployment: str,
     model: str,
     project_id: str,
@@ -100,25 +102,21 @@ async def make_point(
         match request_type:
             case RequestType.CHAT_COMPLETION:
                 response_contents = get_chat_completion_response_contents(
-                    logger, response
-                )
-                request_contents = get_chat_completion_request_contents(
-                    logger, request
+                    response
                 )
+                request_contents = get_chat_completion_request_contents(request)
 
                 request_content = "\n".join(request_contents)
                 response_content = "\n".join(response_contents)
 
                 if chat_id:
                     topic = to_string(
                         await topic_model.get_topic_by_text(
-                            "\n\n".join(request_contents + response_contents)
+                            "\n\n".join(request_contents + response_contents),
                         )
                     )
             case RequestType.EMBEDDING:
-                request_contents = get_embeddings_request_contents(
-                    logger, request
-                )
+                request_contents = get_embeddings_request_contents(request)
 
                 request_content = "\n".join(request_contents)
                 if chat_id:
@@ -167,7 +165,7 @@ async def make_point(
             (
                 "undefined"
                 if not chat_id or request is None or response is None
-                else await detect_lang(logger, request, response, request_type)
+                else await detect_lang(request, response, request_type)
             ),
         )
         .tag("upstream", to_string(upstream_url))
@@ -267,8 +265,7 @@ async def parse_usage_per_model(response: dict | None):
 
 
 async def on_message(
-    logger: Logger,
-    influx_writer: Callable[[Point], Awaitable[None]],
+    influx_writer: InfluxWriterAsync,
     deployment: str,
     model: str,
     project_id: str,
@@ -287,14 +284,11 @@ async def on_message(
     trace: dict | None,
     execution_path: list | None,
 ):
-    logger.info(f"Chat completion response length {len(response or [])}")
-
     usage_per_model = await parse_usage_per_model(response)
     response_usage = None if response is None else response.get("usage")
 
     if token_usage is not None:
         point = await make_point(
-            logger,
             deployment,
             model,
             project_id,
@@ -316,7 +310,6 @@ async def on_message(
         await influx_writer(point)
     elif len(usage_per_model) == 0:
         point = await make_point(
-            logger,
             deployment,
             model,
             project_id,
@@ -338,7 +331,6 @@ async def on_message(
         await influx_writer(point)
     else:
         point = await make_point(
-            logger,
             deployment,
             model,
             project_id,
@@ -361,7 +353,6 @@ async def on_message(
 
         for usage in usage_per_model:
             point = await make_point(
-                logger,
                 deployment,
                 usage["model"],
                 project_id,

aidial_analytics_realtime/app.py

@@ -1,9 +1,13 @@
+import asyncio
 import contextlib
 import json
 import logging
 import re
+import sys
 from datetime import datetime
 
+import aiohttp
+import starlette.requests
 import uvicorn
 from fastapi import Depends, FastAPI, Request
 from fastapi.responses import JSONResponse
@@ -22,7 +26,10 @@
 from aidial_analytics_realtime.topic_model import TopicModel
 from aidial_analytics_realtime.universal_api_utils import merge
 from aidial_analytics_realtime.utils.concurrency import cpu_task_executor
-from aidial_analytics_realtime.utils.log_config import configure_loggers, logger
+from aidial_analytics_realtime.utils.logging import add_logger_prefix
+from aidial_analytics_realtime.utils.logging import app_logger as logger
+from aidial_analytics_realtime.utils.logging import configure_loggers
+from aidial_analytics_realtime.utils.timer import Timer
 
 RATE_PATTERN = r"/v1/(.+?)/rate"
 CHAT_COMPLETION_PATTERN = r"/openai/deployments/(.+?)/chat/completions"
@@ -61,7 +68,6 @@ async def on_rate_message(
     response: dict,
     influx_writer: InfluxWriterAsync,
 ):
-    logger.info(f"Rate message length {len(request) + len(response)}")
     request_body = json.loads(request["body"])
     point = make_rate_point(
         deployment,
@@ -133,7 +139,6 @@ async def on_chat_completion_message(
             response_body = json.loads(response["body"])
 
     await on_message(
-        logger,
         influx_writer,
         deployment,
         model or deployment,
@@ -187,7 +192,6 @@ async def on_embedding_message(
     )
 
     await on_message(
-        logger,
         influx_writer,
         deployment,
         deployment,
@@ -298,28 +302,76 @@ async def on_log_messages(
     topic_model: TopicModel = Depends(),
     rates_calculator: RatesCalculator = Depends(),
 ):
+
     data = await request.json()
 
-    statuses = []
-    for idx, item in enumerate(data):
-        try:
-            await on_log_message(
-                json.loads(item["message"]),
-                influx_writer,
-                topic_model,
-                rates_calculator,
-            )
-        except Exception as e:
-            logging.exception(f"Error processing message #{idx}")
-            statuses.append({"status": "error", "error": str(e)})
-        else:
-            statuses.append({"status": "success"})
+    n = len(data)
+    logger.info(f"number of messages: {n}")
+
+    statuses: list[dict] = []
+
+    async with Timer(logger.debug):
+
+        async def _task(i: int, message_str: str) -> dict:
+            add_logger_prefix(f"[{i}/{n}]")
+
+            async with Timer(logger.debug):
+                return await process_message(
+                    json.loads(message_str),
+                    influx_writer,
+                    topic_model,
+                    rates_calculator,
+                )
+
+        statuses = await asyncio.gather(
+            *[
+                _task(i, message["message"])
+                for i, message in enumerate(data, start=1)
+            ]
+        )
+
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(f"response: {json.dumps(statuses)}")
 
     # Returning 200 code even if processing of some messages has failed,
     # since the log broker that sends the messages may decide to retry the failed requests.
     return JSONResponse(content=statuses, status_code=200)
 
 
+async def process_message(
+    message: dict,
+    influx_writer: InfluxWriterAsync,
+    topic_model: TopicModel,
+    rates_calculator: RatesCalculator,
+) -> dict:
+    def _error(reason: str | None = None) -> dict:
+        error = str(sys.exc_info()[1])
+        ret = {"status": "error"}
+        if error:
+            ret["error"] = error
+        if reason:
+            ret["reason"] = reason
+            logger.error(reason)
+        else:
+            logger.exception("caught exception")
+        return ret
+
+    try:
+        await on_log_message(
+            message, influx_writer, topic_model, rates_calculator
+        )
+        logger.info("success")
+        return {"status": "success"}
+    except starlette.requests.ClientDisconnect:
+        return _error("client disconnect")
+    except aiohttp.ClientConnectionError:
+        return _error("connection error")
+    except asyncio.TimeoutError:
+        return _error("timeout")
+    except Exception:
+        return _error()
+
+
 @app.get("/health")
 def health():
     return {"status": "ok"}

aidial_analytics_realtime/dial.py

@@ -1,40 +1,31 @@
-from logging import Logger
 from typing import Iterator, List
 
+from aidial_analytics_realtime.utils.logging import app_logger as logger
 
-def get_chat_completion_request_contents(
-    logger: Logger, request: dict
-) -> List[str]:
-    return list(_chat_completion_request_contents(logger, request))
 
+def get_chat_completion_request_contents(request: dict) -> List[str]:
+    return list(_chat_completion_request_contents(request))
 
-def get_chat_completion_response_contents(
-    logger: Logger, response: dict
-) -> List[str]:
-    return list(_chat_completion_response_contents(logger, response))
 
+def get_chat_completion_response_contents(response: dict) -> List[str]:
+    return list(_chat_completion_response_contents(response))
 
-def get_embeddings_request_contents(logger: Logger, request: dict) -> List[str]:
-    return list(_embeddings_request_contents(logger, request))
 
+def get_embeddings_request_contents(request: dict) -> List[str]:
+    return list(_embeddings_request_contents(request))
 
-def _chat_completion_request_contents(
-    logger: Logger, request: dict
-) -> Iterator[str]:
+
+def _chat_completion_request_contents(request: dict) -> Iterator[str]:
     for message in request["messages"]:
-        yield from _chat_completion_message_contents(logger, message)
+        yield from _chat_completion_message_contents(message)
 
 
-def _chat_completion_response_contents(
-    logger: Logger, response: dict
-) -> Iterator[str]:
+def _chat_completion_response_contents(response: dict) -> Iterator[str]:
     message = response["choices"][0]["message"]
-    yield from _chat_completion_message_contents(logger, message)
+    yield from _chat_completion_message_contents(message)
 
 
-def _embeddings_request_contents(
-    logger: Logger, request: dict
-) -> Iterator[str]:
+def _embeddings_request_contents(request: dict) -> Iterator[str]:
     inp = request.get("input")
 
     if isinstance(inp, str):
@@ -47,9 +38,7 @@ def _embeddings_request_contents(
         logger.warning(f"Unexpected type of embeddings input: {type(inp)}")
 
 
-def _chat_completion_message_contents(
-    logger: Logger, message: dict
-) -> Iterator[str]:
+def _chat_completion_message_contents(message: dict) -> Iterator[str]:
     content = message.get("content")
     if content is None:
         return

aidial_analytics_realtime/influx_writer.py

@@ -1,11 +1,16 @@
+import logging
 import os
 from typing import Awaitable, Callable, Tuple
 
 from influxdb_client import Point
 from influxdb_client.client.influxdb_client_async import InfluxDBClientAsync
 
+from aidial_analytics_realtime.utils.timer import Timer
+
 InfluxWriterAsync = Callable[[Point], Awaitable[None]]
 
+_logger = logging.getLogger("app.influx")
+
 
 def create_influx_writer() -> Tuple[InfluxDBClientAsync, InfluxWriterAsync]:
     influx_url = os.environ["INFLUX_URL"]
@@ -19,6 +24,7 @@ def create_influx_writer() -> Tuple[InfluxDBClientAsync, InfluxWriterAsync]:
     influx_write_api = client.write_api()
 
     async def influx_writer_impl(record: Point):
-        await influx_write_api.write(bucket=influx_bucket, record=record)
+        with Timer(_logger.debug):
+            await influx_write_api.write(bucket=influx_bucket, record=record)
 
     return client, influx_writer_impl