diff --git a/service/README.md b/service/README.md index 78ad1db..d74665f 100644 --- a/service/README.md +++ b/service/README.md @@ -45,6 +45,10 @@ The file path can be overridden globally via `AI_DOCUMENT_PLUGIN_CONFIG_PATH`. S Prompt templates and LLM parameters for each step. In `config.yaml`, `files.prompts_path` must be a relative path and is resolved relative to the active config file. +## Logging + +The service logs only to standard output. Configure the verbosity under `logging.level`. + ## Make commands The project `Makefile` provides a few shortcuts for common development tasks: diff --git a/service/config.template.yaml b/service/config.template.yaml index 7705409..a1db2b6 100644 --- a/service/config.template.yaml +++ b/service/config.template.yaml @@ -29,4 +29,4 @@ files: # This specifies how many DMPs can be processed at the same time. # Each DMP process can create further parallel lines of execution depending on the tenants setting. # This limit is server-wide, that means it has the same counter for all tenants. -max_parallel_executions: 2 \ No newline at end of file +max_parallel_executions: 2 diff --git a/service/pyproject.toml b/service/pyproject.toml index baf37bf..fccfeac 100644 --- a/service/pyproject.toml +++ b/service/pyproject.toml @@ -34,7 +34,6 @@ dependencies = [ "sqlalchemy[asyncio]", "tabulate", "tiktoken", - "tqdm", ] [project.scripts] diff --git a/service/requirements.txt b/service/requirements.txt index 74171d6..b7c71b6 100644 --- a/service/requirements.txt +++ b/service/requirements.txt @@ -157,7 +157,6 @@ tiktoken==0.13.0 # via ai-document-plugin-service (pyproject.toml) tqdm==4.70.0 # via - # ai-document-plugin-service (pyproject.toml) # haystack-ai # openai truststore==0.10.4 diff --git a/service/src/ai_document_plugin_service/ai/assignment/assignment_component.py b/service/src/ai_document_plugin_service/ai/assignment/assignment_component.py index 84d9e70..2888b5d 100644 --- a/service/src/ai_document_plugin_service/ai/assignment/assignment_component.py +++ b/service/src/ai_document_plugin_service/ai/assignment/assignment_component.py @@ -1,6 +1,7 @@ import asyncio import itertools import logging +import time from collections.abc import Callable from typing import Any, TypedDict @@ -73,10 +74,19 @@ async def run_async( km: dict[str, Any], on_progress: Callable[[str], None] | None = None, ) -> AssignmentComponentResult: + started = time.perf_counter() logger.debug('Step 1: Assigning questions to sections...') sections = build_section_records(template_data) question_chunks, question_id_to_path = build_question_chunks(data) + logger.info( + 'Starting question-to-section assignment', + extra={ + 'section_count': len(sections), + 'question_chunk_count': len(question_chunks), + 'question_path_count': len(question_id_to_path), + }, + ) stats = AssignmentStats() section_formatter = SectionFormatter(sections) @@ -119,6 +129,15 @@ async def match_chunk(question_chunk: str) -> dict[str, list[str]]: result_mapping, km, ) + logger.info( + 'Completed question-to-section assignment', + extra={ + 'assignment_count': len(assignments), + 'mapped_question_count': len(result_mapping), + 'llm_call_count': stats.total_calls, + }, + ) + stats.set_duration_ms(round((time.perf_counter() - started) * 1000, 3)) return { 'assignments': assignments, diff --git a/service/src/ai_document_plugin_service/ai/assignment/llm.py b/service/src/ai_document_plugin_service/ai/assignment/llm.py index c48eb57..e5ef493 100644 --- a/service/src/ai_document_plugin_service/ai/assignment/llm.py +++ b/service/src/ai_document_plugin_service/ai/assignment/llm.py @@ -1,6 +1,5 @@ import json import logging -import pathlib from abc import ABC, abstractmethod from json import JSONDecodeError from typing import TYPE_CHECKING @@ -78,6 +77,7 @@ async def match_questions_to_sections( async def call_and_parse() -> dict[str, list[str]]: response = await self.client.completion( + stats=stats, messages=messages, temperature=self.config.assignment.temperature, max_tokens=self.config.assignment.max_tokens, @@ -85,9 +85,9 @@ async def call_and_parse() -> dict[str, list[str]]: ) choice = response.choices[0] if choice.finish_reason != 'stop': - logger.debug( - 'Model did not stop generating naturally: %s', - choice, + logger.error( + 'Model did not stop generating naturally', + extra={'finish_reason': choice.finish_reason}, ) msg = 'Model did not stop generating naturally.' raise ModelDidNotStopError(msg) @@ -97,6 +97,7 @@ async def call_and_parse() -> dict[str, list[str]]: return self._parse_json_question_to_sections(content) except JSONDecodeError as e: msg = 'Unable to parse: ' + content + logger.exception('Unable to parse LLM assignment response as JSON', exc_info=e) raise UnableToParseResponseError(msg) from e return await call_with_retry(call_and_parse) @@ -113,6 +114,7 @@ def _parse_json_question_to_sections(content: str) -> dict[str, list[str]]: if not isinstance(data, dict): msg = 'LLM response is not a JSON object.' + logger.error('LLM assignment response is not a JSON object') raise TypeError(msg) result: dict[str, list[str]] = {} for id_str, section_list in data.items(): @@ -123,31 +125,3 @@ def _parse_json_question_to_sections(content: str) -> dict[str, list[str]]: else: result[str(id_str)] = [] return result - - -class LoggingNoopLayerMatcher(LayerMatcher): - """Logs assignment inputs to logger and a JSONL file; returns no mappings.""" - - def __init__(self, log_path: str | pathlib.Path) -> None: - self._log_path = pathlib.Path(log_path) - - async def match_questions_to_sections( - self, - sections_xml: str, - question_chunk_xml: str, - stats: AssignmentStats, - ) -> dict[str, list[str]]: - record = { - 'sections_xml': sections_xml, - 'question_chunk_xml': question_chunk_xml, - 'stats': stats.to_dict(), - } - self._log_path.parent.mkdir(parents=True, exist_ok=True) - with self._log_path.open('a', encoding='utf-8') as handle: - handle.write(json.dumps(record, ensure_ascii=False) + '\n') - logger.info( - 'LoggingNoopLayerMatcher appended inputs to %s stats=%s', - self._log_path, - stats, - ) - return {} diff --git a/service/src/ai_document_plugin_service/ai/assignment/sections/llm.py b/service/src/ai_document_plugin_service/ai/assignment/sections/llm.py index 167fdef..696f06d 100644 --- a/service/src/ai_document_plugin_service/ai/assignment/sections/llm.py +++ b/service/src/ai_document_plugin_service/ai/assignment/sections/llm.py @@ -1,8 +1,8 @@ +import json +import logging import typing from abc import ABC, abstractmethod -from tqdm import tqdm - from ai_document_plugin_service.ai.assignment.types import LeafSection from ai_document_plugin_service.ai.common.config import Config from ai_document_plugin_service.ai.common.llm_client import ( @@ -12,6 +12,8 @@ ) from ai_document_plugin_service.ai.common.types import AssignmentStats +logger = logging.getLogger(__name__) + class SectionIdGenerator(ABC): @abstractmethod @@ -47,8 +49,9 @@ async def generate_leaf_section_ids( user_tpl = self.config.section_id.user_message result: dict[str, str] = {} used_ids: set[str] = set() + total_sections = len(leaf_sections) - for leaf in tqdm(leaf_sections): + for index, leaf in enumerate(leaf_sections, start=1): existing_str = ', '.join(sorted(used_ids)) if used_ids else '(none yet)' content_block = leaf.text.strip() user_msg = ( @@ -77,6 +80,17 @@ async def generate_leaf_section_ids( sid = f'{sid}_{len(used_ids)}' used_ids.add(sid) result[leaf.id] = sid + logger.info( + json.dumps( + { + 'event': 'generating_section_identifiers_progress', + 'completed_sections': index, + 'total_sections': total_sections, + }, + ensure_ascii=False, + sort_keys=True, + ), + ) return result @@ -98,6 +112,18 @@ async def generate_leaf_section_ids( # ty: ignore[invalid-method-override] _: AssignmentStats, ) -> dict[str, str]: res = {} - for i, leaf in tqdm(enumerate(leaf_sections)): - res[leaf.id] = f'{leaf.id}_{i}' + total_sections = len(leaf_sections) + for index, leaf in enumerate(leaf_sections, start=1): + res[leaf.id] = f'{leaf.id}_{index - 1}' + logger.info( + json.dumps( + { + 'event': 'generating_section_identifiers_progress', + 'completed_sections': index, + 'total_sections': total_sections, + }, + ensure_ascii=False, + sort_keys=True, + ), + ) return res diff --git a/service/src/ai_document_plugin_service/ai/common/config.py b/service/src/ai_document_plugin_service/ai/common/config.py index 4d2c6d7..8b558ac 100644 --- a/service/src/ai_document_plugin_service/ai/common/config.py +++ b/service/src/ai_document_plugin_service/ai/common/config.py @@ -79,7 +79,7 @@ def _normalize_path(path: str) -> str: return str(pathlib.Path(_expand_env_vars(path).strip()).expanduser()) -def _get(config: dict[str, Any], *path: str) -> Any: # noqa: ANN401 +def _get(config: dict[str, Any], *path: str, allow_empty_string: bool = False) -> Any: # noqa: ANN401 current = config for key in path: if not isinstance(current, dict) or key not in current: @@ -92,7 +92,7 @@ def _get(config: dict[str, Any], *path: str) -> Any: # noqa: ANN401 if current is None: msg = f"Missing required config value: '{'.'.join(path)}'" raise ValueError(msg) - if isinstance(current, str) and not current.strip(): + if isinstance(current, str) and not current.strip() and not allow_empty_string: msg = f"Missing required config value: '{'.'.join(path)}'" raise ValueError(msg) return current @@ -204,7 +204,6 @@ def resolve_config_path(config_path: str | None = None) -> str: def load_config(config_path: str | None = None) -> Config: resolved_config_path = _resolve_existing_path(resolve_config_path(config_path)) config_dir = pathlib.Path(resolved_config_path).parent - with pathlib.Path(resolved_config_path).open(encoding='utf-8') as handle: config = yaml.safe_load(handle) diff --git a/service/src/ai_document_plugin_service/ai/common/dynamic_semaphore.py b/service/src/ai_document_plugin_service/ai/common/dynamic_semaphore.py index 3e94652..8dd41d1 100644 --- a/service/src/ai_document_plugin_service/ai/common/dynamic_semaphore.py +++ b/service/src/ai_document_plugin_service/ai/common/dynamic_semaphore.py @@ -57,6 +57,14 @@ def set_limit(self, new_limit: int) -> None: self.limit = new_limit self._wake_waiters() + @property + def active_count(self) -> int: + return self._active_count + + @property + def queued_count(self) -> int: + return len(self._waiters) + async def acquire(self) -> None: loop = self._bind_loop() if self._active_count < self.limit: diff --git a/service/src/ai_document_plugin_service/ai/common/execution_logging.py b/service/src/ai_document_plugin_service/ai/common/execution_logging.py new file mode 100644 index 0000000..e806ed6 --- /dev/null +++ b/service/src/ai_document_plugin_service/ai/common/execution_logging.py @@ -0,0 +1,76 @@ +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterator + +_run_log_context: ContextVar['RunLogContext | None'] = ContextVar('run_log_context', default=None) +_LLM_EVENT_LOGGER = logging.getLogger('ai_document_plugin_service.execution.llm') +_PIPELINE_EVENT_LOGGER = logging.getLogger('ai_document_plugin_service.execution.pipeline') + + +@dataclass(frozen=True) +class RunLogContext: + run_id: str + questionnaire_uuid: str | None = None + template_uuid: str | None = None + template_title: str | None = None + knowledge_model_uuid: str | None = None + user_uuid: str | None = None + tenant_uuid: str | None = None + + +@contextmanager +def run_log_context(context: RunLogContext) -> 'Iterator[None]': + token = _run_log_context.set(context) + try: + yield + finally: + _run_log_context.reset(token) + + +def get_run_log_context() -> RunLogContext | None: + return _run_log_context.get() + + +def utc_now_iso() -> str: + return datetime.now(tz=UTC).isoformat() + + +def log_llm_event(record: dict[str, Any]) -> None: + _emit_event(_LLM_EVENT_LOGGER, record) + + +def log_timing_event(event: str, **fields: Any) -> None: # noqa: ANN401 + _emit_event(_PIPELINE_EVENT_LOGGER, {'event': event, **fields}) + + +def log_semaphore_event(event: str, **fields: Any) -> None: # noqa: ANN401 + _emit_event(_LLM_EVENT_LOGGER, {'event': event, **fields}) + + +def _with_context(record: dict[str, Any]) -> dict[str, Any]: + payload: dict[str, Any] = { + **record, + } + context = get_run_log_context() + if context is None: + return payload + + payload.update({field: value for field, value in context.__dict__.items() if value is not None}) + return payload + + +def _emit_event(logger: logging.Logger, record: dict[str, Any]) -> None: + payload = _with_context(record) + logger.log(_event_log_level(payload), 'Execution event', extra=payload) + + +def _event_log_level(record: dict[str, Any]) -> int: + if record.get('status') in {'error', 'failed'} or record.get('state') == 'failed' or 'error_type' in record: + return logging.ERROR + return logging.INFO diff --git a/service/src/ai_document_plugin_service/ai/common/json_log_writer.py b/service/src/ai_document_plugin_service/ai/common/json_log_writer.py new file mode 100644 index 0000000..bc9e679 --- /dev/null +++ b/service/src/ai_document_plugin_service/ai/common/json_log_writer.py @@ -0,0 +1,33 @@ +import json +import threading +from collections.abc import Mapping +from pathlib import Path + +from pydantic import BaseModel + +type JsonScalar = str | int | float | bool | None +type JsonValue = JsonScalar | list[JsonValue] | dict[str, JsonValue] + +_WRITE_LOCK = threading.Lock() + + +def append_jsonl(path: Path, payload: dict[str, JsonValue]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + serialized = json.dumps(payload, ensure_ascii=False) + with _WRITE_LOCK, path.open('a', encoding='utf-8') as handle: + handle.write(serialized + '\n') + + +def make_json_safe(value: object) -> JsonValue: + result: JsonValue + if value is None or isinstance(value, str | int | float | bool): + result = value + elif isinstance(value, BaseModel): + result = make_json_safe(value.model_dump(mode='json')) + elif isinstance(value, Mapping): + result = {str(key): make_json_safe(item) for key, item in value.items()} + elif isinstance(value, list | tuple | set): + result = [make_json_safe(item) for item in value] + else: + result = repr(value) + return result diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index 111e580..f02db67 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -9,6 +9,9 @@ from openai.types.chat import ChatCompletion from ai_document_plugin_service.ai.common.dynamic_semaphore import DynamicSemaphore +from ai_document_plugin_service.ai.common.execution_logging import ( + log_llm_event, +) if TYPE_CHECKING: from ai_document_plugin_service.ai.common import AssignmentStats @@ -71,6 +74,7 @@ def extract_usage_tokens(response: object) -> tuple[int, int]: if input_tokens is not None and output_tokens is not None: return input_tokens, output_tokens msg = 'No token info provided in the API response' + logger.error('LLM response is missing token usage information') raise MissingTokenUsageError(msg) @@ -129,10 +133,13 @@ def update_config(self, model: str, api_key: str, api_url: str, parallel_workers self.max_workers = max(1, parallel_workers or 1) self.semaphore.set_limit(self.max_workers) self.client = AsyncOpenAI(api_key=api_key, base_url=api_url, max_retries=0) - logger.debug( - '[llm] tenant=%s: Updated LLM client config, setting semaphore limit to %s', - self.tenant_uuid, - self.max_workers, + logger.info( + 'LLM client configuration updated', + extra={ + 'tenant_uuid': str(self.tenant_uuid), + 'llm_model': self.model, + 'llm_max_workers': self.max_workers, + }, ) def get_max_workers(self) -> int: @@ -150,37 +157,152 @@ def get_model_name(self) -> str: async def completion( self, *args: Any, # noqa: ANN401 + stats: 'AssignmentStats | None' = None, **kwargs: Any, # noqa: ANN401 ) -> ChatCompletion: if self.model is None: + logger.error('LLM completion failed: model is not configured', + extra={'tenant_uuid': str(self.tenant_uuid)}) raise InvalidLLMConfigError('model', self.tenant_uuid) # noqa: EM101 if self.max_workers is None: + logger.error('LLM completion failed: max_workers is not configured', + extra={'tenant_uuid': str(self.tenant_uuid)}) raise InvalidLLMConfigError('max_workers', self.tenant_uuid) # noqa: EM101 if self.api_url is None: + logger.error('LLM completion failed: api_url is not configured', + extra={'tenant_uuid': str(self.tenant_uuid)}) raise InvalidLLMConfigError('api_url', self.tenant_uuid) # noqa: EM101 if self.api_key is None: + logger.error('LLM completion failed: api_key is not configured', + extra={'tenant_uuid': str(self.tenant_uuid)}) raise InvalidLLMConfigError('api_key', self.tenant_uuid) # noqa: EM101 if self.client is None: msg = f'LLM internal client is null but api_key and api_url is set for tenant {self.tenant_uuid}.' + logger.error('LLM completion failed: internal AsyncOpenAI client is missing', + extra={'tenant_uuid': str(self.tenant_uuid)}) raise RuntimeError(msg) req_id = uuid.uuid4().hex[:8] wait_start = time.perf_counter() - logger.debug('[llm] tenant=%s req=%s model=%s queueing', self.tenant_uuid, req_id, self.model) + log_llm_event( + { + 'state': 'waiting_for_semaphore', + 'req_id': req_id, + 'tenant_uuid': str(self.tenant_uuid), + 'model': self.model, + 'limit': self.semaphore.limit, + 'active_count': self.semaphore.active_count, + 'queued_count': self.semaphore.queued_count + 1, + }, + ) async with self.semaphore: wait_s = time.perf_counter() - wait_start - logger.debug( - '[llm] tenant=%s req=%s acquired semaphore after %.3fs (limit=%s)', - self.tenant_uuid, - req_id, - wait_s, - self.semaphore.limit, + log_llm_event( + { + 'state': 'waiting_for_llm_response', + 'req_id': req_id, + 'tenant_uuid': str(self.tenant_uuid), + 'model': self.model, + 'limit': self.semaphore.limit, + 'active_count': self.semaphore.active_count, + 'queued_count': self.semaphore.queued_count, + 'semaphore_wait_ms': _duration_ms(wait_s), + 'message_count': _count_messages(kwargs.get('messages')), + }, ) call_start = time.perf_counter() - result = await self.client.chat.completions.create(*args, model=self.model, **kwargs) - logger.debug( - '[llm] tenant=%s req=%s completed in %.3fs (releasing semaphore)', - self.tenant_uuid, - req_id, - time.perf_counter() - call_start, + try: + result = await self.client.chat.completions.create(*args, model=self.model, **kwargs) + except Exception as error: + duration_s = time.perf_counter() - call_start + _add_timing(stats, wait_s, duration_s) + self._log_llm_completion( + req_id=req_id, + status='error', + wait_s=wait_s, + duration_s=duration_s, + request_kwargs=kwargs, + error=error, + ) + raise + + duration_s = time.perf_counter() - call_start + _add_timing(stats, wait_s, duration_s) + self._log_llm_completion( + req_id=req_id, + status='success', + wait_s=wait_s, + duration_s=duration_s, + request_kwargs=kwargs, + response=result, ) return result + + def _log_llm_completion( + self, + *, + req_id: str, + status: str, + wait_s: float, + duration_s: float, + request_kwargs: dict[str, Any], + response: ChatCompletion | None = None, + error: Exception | None = None, + ) -> None: + payload = { + 'state': 'completed' if status == 'success' else 'failed', + 'status': status, + 'req_id': req_id, + 'tenant_uuid': str(self.tenant_uuid), + 'model': self.model, + 'semaphore_wait_ms': _duration_ms(wait_s), + 'llm_response_ms': _duration_ms(duration_s), + 'total_llm_ms': _duration_ms(wait_s + duration_s), + 'message_count': _count_messages(request_kwargs.get('messages')), + 'temperature': request_kwargs.get('temperature'), + 'max_tokens': request_kwargs.get('max_tokens'), + 'reasoning_effort': request_kwargs.get('reasoning_effort'), + } + if response is not None: + usage = _extract_usage(response) + payload.update( + finish_reason=response.choices[0].finish_reason if response.choices else None, + prompt_tokens=usage['prompt_tokens'], + completion_tokens=usage['completion_tokens'], + total_tokens=usage['total_tokens'], + ) + if error is not None: + payload.update( + { + 'error_type': type(error).__name__, + 'error_message': str(error), + }, + ) + log_llm_event(payload) + + +def _count_messages(messages: object) -> int | None: + if isinstance(messages, list): + return len(messages) + return None + + +def _duration_ms(duration_s: float) -> float: + return round(duration_s * 1000, 3) + + +def _add_timing(stats: 'AssignmentStats | None', wait_s: float, duration_s: float) -> None: + if stats is None: + return + stats.add_llm_timing( + wait_ms=_duration_ms(wait_s), + response_ms=_duration_ms(duration_s), + ) + + +def _extract_usage(response: ChatCompletion) -> dict[str, int | None]: + usage = getattr(response, 'usage', None) + return { + 'prompt_tokens': getattr(usage, 'prompt_tokens', None), + 'completion_tokens': getattr(usage, 'completion_tokens', None), + 'total_tokens': getattr(usage, 'total_tokens', None), + } diff --git a/service/src/ai_document_plugin_service/ai/common/logging_payloads.py b/service/src/ai_document_plugin_service/ai/common/logging_payloads.py new file mode 100644 index 0000000..f6757a8 --- /dev/null +++ b/service/src/ai_document_plugin_service/ai/common/logging_payloads.py @@ -0,0 +1,86 @@ +import json +from collections.abc import Mapping, Sequence + +from ai_document_plugin_service.ai.common.json_log_writer import make_json_safe + +SENSITIVE_FIELD_NAMES = { + 'api_key', + 'apikey', + 'authorization', + 'cookie', + 'llm_api_key', + 'password', + 'refresh_token', + 'secret', + 'set-cookie', + 'token', +} +MAX_LOG_TEXT_LENGTH = 1000 + + +def sanitize_for_logging(value: object) -> object: + safe_value = make_json_safe(value) + return _sanitize_value(safe_value) + + +def summarize_payload(value: object, *, max_chars: int = MAX_LOG_TEXT_LENGTH) -> str: + sanitized = sanitize_for_logging(value) + serialized = json.dumps(sanitized, ensure_ascii=False, sort_keys=True) + return truncate_text(serialized, max_chars=max_chars) + + +def summarize_headers(headers: Mapping[str, str]) -> dict[str, object]: + return _sanitize_mapping(headers) + + +def summarize_http_body( + body: bytes, + *, + content_type: str | None = None, + max_chars: int = MAX_LOG_TEXT_LENGTH, +) -> str | None: + if not body: + return None + + normalized_content_type = (content_type or '').lower() + if 'application/json' in normalized_content_type: + try: + parsed = json.loads(body.decode('utf-8')) + except (UnicodeDecodeError, json.JSONDecodeError): + return truncate_text(body.decode('utf-8', errors='replace'), max_chars=max_chars) + return summarize_payload(parsed, max_chars=max_chars) + + if 'application/x-www-form-urlencoded' in normalized_content_type: + return truncate_text(body.decode('utf-8', errors='replace'), max_chars=max_chars) + + if normalized_content_type.startswith('multipart/'): + return f'' + + return truncate_text(body.decode('utf-8', errors='replace'), max_chars=max_chars) + + +def truncate_text(value: str, *, max_chars: int = MAX_LOG_TEXT_LENGTH) -> str: + if len(value) <= max_chars: + return value + return f'{value[:max_chars]}... ' + + +def _sanitize_value(value: object) -> object: + if isinstance(value, Mapping): + return _sanitize_mapping(value) + if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): + return [_sanitize_value(item) for item in value] + if isinstance(value, str): + return truncate_text(value) + return value + + +def _sanitize_mapping[Key](value: Mapping[Key, object]) -> dict[str, object]: + sanitized: dict[str, object] = {} + for raw_key, raw_value in value.items(): + key = str(raw_key) + if key.lower() in SENSITIVE_FIELD_NAMES: + sanitized[key] = '' + continue + sanitized[key] = _sanitize_value(raw_value) + return sanitized diff --git a/service/src/ai_document_plugin_service/ai/common/logging_utils.py b/service/src/ai_document_plugin_service/ai/common/logging_utils.py index fb6ccf5..9cce055 100644 --- a/service/src/ai_document_plugin_service/ai/common/logging_utils.py +++ b/service/src/ai_document_plugin_service/ai/common/logging_utils.py @@ -1,26 +1,72 @@ import logging +from ai_document_plugin_service.ai.common.trace_context import get_trace_id + +LOG_FORMAT = '%(asctime)s | %(levelname)8s | %(name)s: [T:%(traceId)s] %(message)s' +_ORIGINAL_LOG_RECORD_FACTORY = logging.getLogRecordFactory() +_STANDARD_LOG_RECORD_FIELDS = frozenset(logging.makeLogRecord({}).__dict__) | {'asctime', 'message', 'traceId'} + + +class ExtraFormatter(logging.Formatter): + """Appends application fields supplied through ``extra`` to text logs.""" + + def format(self, record: logging.LogRecord) -> str: + formatted = super().format(record) + extra = { + key: value + for key, value in record.__dict__.items() + if key not in _STANDARD_LOG_RECORD_FIELDS + } + return f'{formatted} | extra={extra!r}' if extra else formatted + def configure_logging(level: int | str = logging.DEBUG) -> None: + normalized_level = _normalize_level(level) + _configure_root_stdout_logging(normalized_level) + _configure_library_log_levels() + + +def _normalize_level(level: int | str) -> int: if isinstance(level, str): normalized_level = logging.getLevelName(level.upper()) if not isinstance(normalized_level, int): msg = f'Unsupported log level: {level}' raise TypeError(msg) - level = normalized_level + return normalized_level + return level + +def _configure_root_stdout_logging(level: int) -> None: root_logger = logging.getLogger() - if root_logger.handlers: - root_logger.setLevel(level) - logging.getLogger('httpx').setLevel(logging.WARNING) - logging.getLogger('httpcore').setLevel(logging.WARNING) - logging.getLogger('openai').setLevel(logging.WARNING) - return + root_logger.setLevel(level) + _install_trace_log_record_factory() - logging.basicConfig( - level=level, - format='%(asctime)s %(levelname)s [%(name)s] %(message)s', - ) + if not root_logger.handlers: + logging.basicConfig( + level=level, + format=LOG_FORMAT, + ) + + for handler in root_logger.handlers: + handler.setLevel(level) + handler.setFormatter(ExtraFormatter(LOG_FORMAT)) + + +def _configure_library_log_levels() -> None: + logging.getLogger('haystack').setLevel(logging.WARNING) logging.getLogger('httpx').setLevel(logging.WARNING) logging.getLogger('httpcore').setLevel(logging.WARNING) logging.getLogger('openai').setLevel(logging.WARNING) + + +def _install_trace_log_record_factory() -> None: + current_factory = logging.getLogRecordFactory() + if current_factory is _trace_log_record_factory: + return + logging.setLogRecordFactory(_trace_log_record_factory) + + +def _trace_log_record_factory(*args: object, **kwargs: object) -> logging.LogRecord: + record = _ORIGINAL_LOG_RECORD_FACTORY(*args, **kwargs) + record.traceId = get_trace_id() + return record diff --git a/service/src/ai_document_plugin_service/ai/common/pipeline_metrics.py b/service/src/ai_document_plugin_service/ai/common/pipeline_metrics.py index 6d582a4..45dc351 100644 --- a/service/src/ai_document_plugin_service/ai/common/pipeline_metrics.py +++ b/service/src/ai_document_plugin_service/ai/common/pipeline_metrics.py @@ -63,6 +63,14 @@ def total_output_tokens(self) -> int: def total_cost(self) -> float: return sum(self._price(step.stats)[2] for step in self.steps) + @property + def total_llm_wait_ms(self) -> float: + return round(sum(step.stats.total_llm_wait_ms for step in self.steps), 3) + + @property + def total_llm_response_ms(self) -> float: + return round(sum(step.stats.total_llm_response_ms for step in self.steps), 3) + def _price(self, stats: AssignmentStats) -> tuple[float, float, float]: input_cost = stats.total_input_tokens * self.cost_per_mil_input / 1_000_000 output_cost = stats.total_output_tokens * self.cost_per_mil_output / 1_000_000 @@ -98,6 +106,8 @@ def _build_summary_section(self, elapsed_seconds: float) -> JsonValue: 'cost_per_mil_input': self.cost_per_mil_input, 'cost_per_mil_output': self.cost_per_mil_output, 'elapsed_seconds': elapsed_seconds, + 'total_llm_wait_ms': self.total_llm_wait_ms, + 'total_llm_response_ms': self.total_llm_response_ms, }, } diff --git a/service/src/ai_document_plugin_service/ai/common/trace_context.py b/service/src/ai_document_plugin_service/ai/common/trace_context.py new file mode 100644 index 0000000..e9f4419 --- /dev/null +++ b/service/src/ai_document_plugin_service/ai/common/trace_context.py @@ -0,0 +1,21 @@ +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + +_TRACE_ID_CONTEXT: ContextVar[str] = ContextVar('trace_id', default='-') + + +def get_trace_id() -> str: + return _TRACE_ID_CONTEXT.get() + + +@contextmanager +def trace_context(trace_id: str) -> 'Iterator[None]': + token = _TRACE_ID_CONTEXT.set(trace_id) + try: + yield + finally: + _TRACE_ID_CONTEXT.reset(token) diff --git a/service/src/ai_document_plugin_service/ai/common/types.py b/service/src/ai_document_plugin_service/ai/common/types.py index 7a7c738..292fef3 100644 --- a/service/src/ai_document_plugin_service/ai/common/types.py +++ b/service/src/ai_document_plugin_service/ai/common/types.py @@ -6,10 +6,20 @@ class AssignmentStats: total_calls: int = 0 total_input_tokens: int = 0 total_output_tokens: int = 0 + total_llm_wait_ms: float = 0.0 + total_llm_response_ms: float = 0.0 + total_duration_ms: float = 0.0 def add_usage(self, input_tokens: int, output_tokens: int) -> None: self.add_totals(1, input_tokens, output_tokens) + def add_llm_timing(self, wait_ms: float, response_ms: float) -> None: + self.total_llm_wait_ms += wait_ms + self.total_llm_response_ms += response_ms + + def set_duration_ms(self, duration_ms: float) -> None: + self.total_duration_ms = duration_ms + def add_totals( self, calls: int, @@ -20,9 +30,12 @@ def add_totals( self.total_input_tokens += input_tokens self.total_output_tokens += output_tokens - def to_dict(self) -> dict[str, int]: + def to_dict(self) -> dict[str, int | float]: return { 'total_calls': self.total_calls, 'total_input_tokens': self.total_input_tokens, 'total_output_tokens': self.total_output_tokens, + 'total_llm_wait_ms': self.total_llm_wait_ms, + 'total_llm_response_ms': self.total_llm_response_ms, + 'total_duration_ms': self.total_duration_ms, } diff --git a/service/src/ai_document_plugin_service/ai/generation/dmp_generator_component.py b/service/src/ai_document_plugin_service/ai/generation/dmp_generator_component.py index a0bbd2e..029d26c 100644 --- a/service/src/ai_document_plugin_service/ai/generation/dmp_generator_component.py +++ b/service/src/ai_document_plugin_service/ai/generation/dmp_generator_component.py @@ -1,14 +1,15 @@ import asyncio +import json import logging import math import re +import time from collections.abc import Callable, Coroutine, Iterable from dataclasses import dataclass, field from typing import Any, TypedDict import pandas as pd from haystack import component -from tqdm import tqdm from ai_document_plugin_service.ai.assignment.types import SerializedSectionAssignment from ai_document_plugin_service.ai.common.progress import progress_percent @@ -53,6 +54,7 @@ async def run_async( db_assignments: list[SerializedSectionAssignment] | None = None, on_progress: Callable[[str], None] | None = None, ) -> DmpGeneratorComponentResult: + started = time.perf_counter() """Generate full DMP markdown from nested assignments tree. Returns (markdown, debug_markdown, stats). Use markdown for the polished DMP; @@ -61,6 +63,13 @@ async def run_async( logger.debug('Step 2: Generating DMP markdown...') assignments = db_assignments or new_assignments or [] replies = self._filter_reachable_replies(replies, km) + logger.info( + 'Starting DMP generation', + extra={ + 'assignment_count': len(assignments), + 'reply_count': len(replies), + }, + ) stats = AssignmentStats() max_workers = self.dmp_generator_llm.get_max_workers() @@ -80,26 +89,47 @@ async def run_async( self._collect_leaf_sections(scheduled, leaf_sections) total_sections = len(leaf_sections) + logger.info( + 'Prepared leaf sections for DMP generation', + extra={'leaf_section_count': total_sections, 'max_workers': max_workers}, + ) section_semaphore = asyncio.Semaphore(max_workers) tasks = [ asyncio.create_task(self._execute_leaf_section(section, section_semaphore)) for section in leaf_sections ] - progress_bar = tqdm( - total=total_sections, - desc=f'Generating sections ({max_workers} workers)', - ) for i, task in enumerate(asyncio.as_completed(tasks), start=1): await task - progress_bar.update(1) + logger.info( + json.dumps( + { + 'event': 'generating_sections_progress', + 'completed_sections': i, + 'total_sections': total_sections, + 'progress_percent': progress_percent(i, total_sections), + 'max_workers': max_workers, + }, + ensure_ascii=False, + sort_keys=True, + ), + ) if on_progress is not None: on_progress( f'Writing DMP sections ({progress_percent(i, total_sections)}%)', ) - progress_bar.close() parts = [self._render_scheduled_section(scheduled) for scheduled in scheduled_sections] markdown = '\n\n'.join([s for s, _ in parts]) debug_markdown = '\n\n'.join([d for _, d in parts]) + logger.info( + 'Completed DMP generation', + extra={ + 'leaf_section_count': total_sections, + 'markdown_length': len(markdown), + 'debug_markdown_length': len(debug_markdown), + 'llm_call_count': stats.total_calls, + }, + ) + stats.set_duration_ms(round((time.perf_counter() - started) * 1000, 3)) return { 'markdown': markdown, 'debug_markdown': debug_markdown, @@ -688,6 +718,7 @@ async def _execute_leaf_section( async with semaphore: if section.leaf_coro is None: msg = 'Leaf section has no generation coroutine' + logger.error('Leaf section execution failed: generation coroutine is missing') raise RuntimeError(msg) section.result = await section.leaf_coro @@ -778,6 +809,7 @@ async def _generate_leaf_section( assignments = node['assignments'] if assignments is None: msg = f"Leaf section '{title}' is missing assignments" + logger.error('Leaf section generation failed: assignments are missing', extra={'section_title': title}) raise ValueError(msg) matches, _ = self.match_replies_selection(assignments, replies, km) rows = self._flatten_matched_questions(matches, title) diff --git a/service/src/ai_document_plugin_service/ai/generation/llm.py b/service/src/ai_document_plugin_service/ai/generation/llm.py index 1879c44..e385672 100644 --- a/service/src/ai_document_plugin_service/ai/generation/llm.py +++ b/service/src/ai_document_plugin_service/ai/generation/llm.py @@ -64,6 +64,7 @@ async def section_from_qa( messages = self._section_from_qa_messages(prompt) response = await call_with_retry( lambda: self.client.completion( + stats=stats, messages=messages, temperature=self.config.dmp_generation.temperature, max_tokens=self.config.dmp_generation.max_tokens, diff --git a/service/src/ai_document_plugin_service/ai/generation/parse_answers.py b/service/src/ai_document_plugin_service/ai/generation/parse_answers.py index 2eb946d..029458e 100644 --- a/service/src/ai_document_plugin_service/ai/generation/parse_answers.py +++ b/service/src/ai_document_plugin_service/ai/generation/parse_answers.py @@ -186,4 +186,5 @@ def parse_answer( # noqa: PLR0911 return '' msg = 'Unknown answer type' + logger.error('Unknown answer type encountered while parsing reply', extra={'answer_type': answer_type}) raise RuntimeError(msg, answer_type) diff --git a/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py b/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py index 57f9f3a..14f02d4 100644 --- a/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py +++ b/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py @@ -1,7 +1,12 @@ +import logging from uuid import UUID import httpx +from ai_document_plugin_service.ai.common.logging_payloads import summarize_payload + +logger = logging.getLogger(__name__) + class DSWClient: def __init__(self, token: str, api_url: str) -> None: @@ -10,12 +15,40 @@ def __init__(self, token: str, api_url: str) -> None: async def get_questionnaire_detail(self, questionnaire_uuid: str | UUID) -> dict: url = f'{self.api_url}/projects/{questionnaire_uuid}/questionnaire' + logger.info( + 'Fetching questionnaire detail from DSW', + extra={'url.full': url, 'questionnaire_uuid': str(questionnaire_uuid)}, + ) headers: dict[str, str] = {} if self.token: headers['Authorization'] = f'Bearer {self.token}' - async with httpx.AsyncClient() as client: - response = await client.get(url, headers=headers) - response.raise_for_status() - return response.json() + try: + async with httpx.AsyncClient() as client: + response = await client.get(url, headers=headers) + response.raise_for_status() + except httpx.HTTPError: + logger.exception( + 'Failed to fetch questionnaire detail from DSW', + extra={'url.full': url, 'questionnaire_uuid': str(questionnaire_uuid)}, + ) + raise + + payload = response.json() + logger.info( + 'Fetched questionnaire detail successfully', + extra={ + 'url.full': url, + 'questionnaire_uuid': str(questionnaire_uuid), + 'http.response.status_code': response.status_code, + }, + ) + logger.debug( + 'DSW questionnaire detail payload', + extra={ + 'questionnaire_uuid': str(questionnaire_uuid), + 'dsw.response.body': summarize_payload(payload), + }, + ) + return payload diff --git a/service/src/ai_document_plugin_service/ai/knowledgemodel/parser_component.py b/service/src/ai_document_plugin_service/ai/knowledgemodel/parser_component.py index 9bf93f6..045d907 100644 --- a/service/src/ai_document_plugin_service/ai/knowledgemodel/parser_component.py +++ b/service/src/ai_document_plugin_service/ai/knowledgemodel/parser_component.py @@ -39,10 +39,21 @@ def __init__(self) -> None: def run(self, data: dict, trigger: bool) -> dict[str, list[QuestionData]]: # noqa: FBT001, ARG002 self.km = data['knowledgeModel'] replies = data['replies'] + logger.info( + 'Parsing knowledge model replies into question tree', + extra={ + 'chapter_count': len(self.km.get('chapterUuids', [])), + 'reply_count': len(replies), + }, + ) top_level_questions: list[QuestionData] = [] for chapter_dict in self._iterate_km_chapters(): parsed_chapter = self.parse_chapter(chapter_dict['uuid'], replies) top_level_questions.append(parsed_chapter) + logger.info( + 'Finished parsing knowledge model replies', + extra={'chapter_count': len(top_level_questions)}, + ) return {'data': top_level_questions} def _iterate_km_chapters(self) -> Iterator[dict]: diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py index fcfa622..58b2a34 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py @@ -1,3 +1,4 @@ +import logging from typing import Any from uuid import UUID @@ -5,6 +6,8 @@ from ai_document_plugin_service.ai.persistence.database import Database, JsonValue +logger = logging.getLogger(__name__) + @component class AssignmentLoaderComponent: @@ -16,7 +19,19 @@ def __init__(self, database: Database) -> None: found=bool, ) async def run_async(self, knowledge_model_uuid: UUID, template_uuid: UUID) -> dict[str, Any]: + logger.debug( + 'Loading stored assignments for pipeline', + extra={'knowledge_model_uuid': knowledge_model_uuid, 'template_uuid': str(template_uuid)}, + ) assignments = await self.database.get_assignments(knowledge_model_uuid, template_uuid) + logger.info( + 'Assignment load completed', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'found': assignments is not None, + }, + ) return { 'assignments': assignments, diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index 675d417..f08d25c 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) JsonValue = Mapping[str, object] | Sequence[object] -StatsJson = dict[str, dict[str, int]] +StatsJson = dict[str, dict[str, int | float]] class AssignmentSaverComponentResult(TypedDict): @@ -55,6 +55,18 @@ async def run_async( stats: AssignmentStats | None = None, ) -> AssignmentSaverComponentResult: """Save assignments to storage, optionally including token usage stats.""" + logger.info( + 'Persisting generated assignments', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'knowledge_model_version': knowledge_model_version, + 'template_uuid': str(template_uuid), + 'template_title': template_title, + 'tenant_uuid': str(tenant_uuid), + 'assignment_count': len(assignments), + 'has_stats': stats is not None, + }, + ) serializable = [assignment.to_dict() for assignment in assignments] stats_payload = _serialize_stats(stats) @@ -150,7 +162,13 @@ async def save( json.dumps(stats, indent=2, ensure_ascii=False), encoding='utf-8', ) - logger.debug('Saved assignments to %s', output_path) + logger.info( + 'Saved assignments to filesystem', + extra={ + 'output_path': str(output_path), + 'stats_output_path': str(output_path_stats) if stats is not None else None, + }, + ) @staticmethod def _build_filename( @@ -187,6 +205,14 @@ async def save( tenant_uuid: UUID, created_at: datetime | None = None, ) -> None: + logger.debug( + 'Saving assignments through database-backed saver', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'tenant_uuid': str(tenant_uuid), + }, + ) await self.database.save_template( uuid=template_uuid, title=template_title, @@ -202,6 +228,14 @@ async def save( created_at=created_at, template_uuid=template_uuid, ) + logger.info( + 'Database-backed assignment save completed', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'tenant_uuid': str(tenant_uuid), + }, + ) def _serialize_stats(stats: AssignmentStats | None) -> StatsJson | None: @@ -212,6 +246,9 @@ def _serialize_stats(stats: AssignmentStats | None) -> StatsJson | None: 'total_calls': stats.total_calls, 'total_input_tokens': stats.total_input_tokens, 'total_output_tokens': stats.total_output_tokens, + 'total_llm_wait_ms': stats.total_llm_wait_ms, + 'total_llm_response_ms': stats.total_llm_response_ms, + 'total_duration_ms': stats.total_duration_ms, }, } diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index d36f173..f827307 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -298,10 +298,20 @@ def __init__( self.result_table = schema.result_table self.generation_table = schema.generation_table self._database_verified = False + logger.info( + 'Initialized Postgres database client', + extra={ + 'db.host': config.host, + 'db.port': config.port, + 'db.name': config.name, + 'db.schema': self.schema_name, + }, + ) async def dispose(self) -> None: """Close the engine's connection pool. Call on shutdown / when done with this instance.""" await self.engine.dispose() + logger.info('Disposed Postgres database engine', extra={'db.schema': self.schema_name}) @asynccontextmanager async def transaction(self) -> AsyncIterator[None]: @@ -341,6 +351,7 @@ def _list_existing_tables(self, connection: Connection) -> set[str]: async def _ensure_schema(self) -> None: if self._database_verified: return + logger.debug('Verifying persistence schema availability', extra={'db.schema': self.schema_name}) async with self.engine.connect() as connection: existing_tables = await connection.run_sync(self._list_existing_tables) @@ -354,9 +365,14 @@ async def _ensure_schema(self) -> None: 'Startup migrations should create or update these tables. Check the application startup logs and ' 'database configuration.' ) + logger.error( + 'Database schema is not ready', + extra={'db.schema': self.schema_name, 'missing_tables': missing_tables}, + ) raise RuntimeError(msg) self._database_verified = True + logger.info('Verified persistence schema', extra={'db.schema': self.schema_name}) async def save_assignments( self, @@ -393,6 +409,15 @@ async def save_assignments( async with self._connect() as connection: await connection.execute(upsert_statement) + logger.info( + 'Saved assignments', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'assignment_count': len(assignments) if isinstance(assignments, Sequence) else None, + 'db.schema': self.schema_name, + }, + ) logger.debug( 'Saved assignments for KM package id=%s to %s.assignments', knowledge_model_uuid, @@ -420,12 +445,22 @@ async def create_template( async with self._connect() as connection: await connection.execute(statement) except IntegrityError as exc: + msg = f'Template with title "{title}" already exists.' + logger.warning( + 'Template insert failed because title already exists', + extra={'template_title': title, 'tenant_uuid': str(tenant_uuid), 'db.schema': self.schema_name}, + ) + raise ValueError(msg) from exc raise TemplateTitleConflictError(title) from exc - logger.debug( - 'Created template uuid=%s in %s.template', - template_uuid, - self.schema_name, + logger.info( + 'Created template in database', + extra={ + 'template_uuid': str(template_uuid), + 'template_title': title, + 'tenant_uuid': str(tenant_uuid), + 'db.schema': self.schema_name, + }, ) return template_uuid @@ -504,10 +539,14 @@ async def save_template( async with self._connect() as connection: await connection.execute(upsert_statement) - logger.debug( - 'Saved template uuid=%s to %s.template', - uuid, - self.schema_name, + logger.info( + 'Saved template definition', + extra={ + 'template_uuid': str(uuid), + 'template_title': title, + 'tenant_uuid': str(tenant_uuid), + 'db.schema': self.schema_name, + }, ) async def get_assignments( @@ -527,10 +566,13 @@ async def get_assignments( row = result.fetchone() if row is None: - logger.debug( - 'No assignments found for KM package id=%s in %s.assignments', - knowledge_model_uuid, - self.schema_name, + logger.info( + 'Assignments not found', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'db.schema': self.schema_name, + }, ) return None @@ -565,6 +607,14 @@ async def list_templates(self, tenant_uuid: UUID, user_uuid: UUID) -> list[Templ result = await connection.execute(statement) rows = result.fetchall() + logger.info( + 'Listed templates from database', + extra={ + 'tenant_uuid': str(tenant_uuid), + 'template_count': len(rows), + 'db.schema': self.schema_name, + }, + ) return [TemplateRecord.from_row(row) for row in rows] async def get_template( @@ -634,10 +684,17 @@ async def save_result( async with self._connect() as connection: await connection.execute(upsert_statement) - logger.debug( - 'Saved result for KM package id=%s to %s.result', - knowledge_model_uuid, - self.schema_name, + logger.info( + 'Saved pipeline result', + extra={ + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': str(template_uuid), + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'prepolished_markdown_length': len(prepolished_markdown), + 'markdown_length': len(markdown), + 'db.schema': self.schema_name, + }, ) async def save_stats( @@ -667,12 +724,27 @@ async def save_stats( if result.rowcount == 0: msg = 'Cannot save stats because result row does not exist yet. Save dmp and dmp_pre_polished first.' + logger.error( + 'Stats update failed because result row does not exist', + extra={ + 'knowledge_model_uuid': str(knowledge_model_uuid), + 'template_uuid': str(template_uuid), + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'db.schema': self.schema_name, + }, + ) raise ValueError(msg) - logger.debug( - 'Saved stats for KM package id=%s to %s.result', - knowledge_model_uuid, - self.schema_name, + logger.info( + 'Saved pipeline stats', + extra={ + 'knowledge_model_uuid': str(knowledge_model_uuid), + 'template_uuid': str(template_uuid), + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'db.schema': self.schema_name, + }, ) async def update_result( @@ -705,12 +777,28 @@ async def update_result( if result.rowcount == 0: msg = 'Cannot save result because result row does not exist yet. Create the row first before updating dmp.' + logger.error( + 'Result update failed because result row does not exist', + extra={ + 'knowledge_model_uuid': str(knowledge_model_uuid), + 'template_uuid': str(template_uuid), + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'db.schema': self.schema_name, + }, + ) raise ValueError(msg) - logger.debug( - 'Updated result for KM package id=%s in %s.result', - knowledge_model_uuid, - self.schema_name, + logger.info( + 'Updated stored pipeline result markdown', + extra={ + 'knowledge_model_uuid': str(knowledge_model_uuid), + 'template_uuid': str(template_uuid), + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'markdown_length': len(markdown), + 'db.schema': self.schema_name, + }, ) async def create_generation( diff --git a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py index 872b6c7..117daae 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py @@ -1,3 +1,4 @@ +import logging import typing from typing import TypedDict from uuid import UUID @@ -6,6 +7,8 @@ from ai_document_plugin_service.ai.persistence.database import Database +logger = logging.getLogger(__name__) + class FileSaverComponentResult(TypedDict): markdown: str @@ -26,6 +29,17 @@ async def run_async( debug_markdown: str, markdown: str, ) -> FileSaverComponentResult: + logger.info( + 'Persisting generated markdown result', + extra={ + 'template_uuid': str(template_uuid), + 'knowledge_model_uuid': knowledge_model_uuid, + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + 'debug_markdown_length': len(debug_markdown), + 'markdown_length': len(markdown), + }, + ) await self.database.save_result( template_uuid, knowledge_model_uuid, @@ -34,6 +48,15 @@ async def run_async( debug_markdown, markdown, ) + logger.info( + 'Generated markdown result persisted', + extra={ + 'template_uuid': str(template_uuid), + 'knowledge_model_uuid': knowledge_model_uuid, + 'user_uuid': str(user_uuid), + 'tenant_uuid': str(tenant_uuid), + }, + ) return { 'markdown': markdown, diff --git a/service/src/ai_document_plugin_service/ai/polishing/dmp_polisher_component.py b/service/src/ai_document_plugin_service/ai/polishing/dmp_polisher_component.py index 4600d80..aa419a0 100644 --- a/service/src/ai_document_plugin_service/ai/polishing/dmp_polisher_component.py +++ b/service/src/ai_document_plugin_service/ai/polishing/dmp_polisher_component.py @@ -5,6 +5,7 @@ """ import logging +import time from collections.abc import Callable from typing import TypedDict @@ -33,6 +34,7 @@ async def run_async( template_data: dict | None = None, on_progress: Callable[[str], None] | None = None, ) -> DmpPolisherComponentResult: + started = time.perf_counter() """Polish the DMP by moving content to relevant sections and improving structure. Args: @@ -46,6 +48,13 @@ async def run_async( """ stats = AssignmentStats() + logger.info( + 'Starting DMP polishing', + extra={ + 'input_markdown_length': len(markdown), + 'has_template_data': template_data is not None, + }, + ) if on_progress is not None: on_progress('Polishing document') structure_str = DmpPolisherComponent._build_template_structure_string(template_data) @@ -54,6 +63,15 @@ async def run_async( structure_str=structure_str, stats=stats, ) + logger.info( + 'Completed DMP polishing', + extra={ + 'input_markdown_length': len(markdown), + 'output_markdown_length': len(polished), + 'llm_call_count': stats.total_calls, + }, + ) + stats.set_duration_ms(round((time.perf_counter() - started) * 1000, 3)) return { 'markdown': polished, 'stats': stats, diff --git a/service/src/ai_document_plugin_service/ai/polishing/llm.py b/service/src/ai_document_plugin_service/ai/polishing/llm.py index 2d798c8..2d1fe2c 100644 --- a/service/src/ai_document_plugin_service/ai/polishing/llm.py +++ b/service/src/ai_document_plugin_service/ai/polishing/llm.py @@ -48,6 +48,7 @@ async def polish_dmp( response = await call_with_retry( lambda: self.client.completion( + stats=stats, messages=messages, temperature=self.config.dmp_polishing.temperature, max_tokens=self.config.dmp_polishing.max_tokens, diff --git a/service/src/ai_document_plugin_service/ai/run_pipeline.py b/service/src/ai_document_plugin_service/ai/run_pipeline.py index 0df8efe..6b2ed08 100644 --- a/service/src/ai_document_plugin_service/ai/run_pipeline.py +++ b/service/src/ai_document_plugin_service/ai/run_pipeline.py @@ -17,6 +17,7 @@ get_component_markdown, get_component_stats, ) +from ai_document_plugin_service.ai.common.execution_logging import log_timing_event from ai_document_plugin_service.ai.generation.dmp_generator_component import DmpGeneratorComponent from ai_document_plugin_service.ai.generation.llm import SectionGenerationLLM from ai_document_plugin_service.ai.knowledgemodel.parser_component import ParserComponent @@ -122,7 +123,17 @@ async def run_pipeline( on_progress: ProgressCallback | None = None, ) -> tuple[UUID, str]: t1 = time.time() - km_data = await dsw_client.get_questionnaire_detail(questionnaire_uuid=questionnaire_uuid) + pipeline_total_started = time.perf_counter() + questionnaire_fetch_started = time.perf_counter() + try: + km_data = await dsw_client.get_questionnaire_detail(questionnaire_uuid=questionnaire_uuid) + except Exception: + logger.exception('Failed to load questionnaire detail', extra={'questionnaire_uuid': str(questionnaire_uuid)}) + raise + log_timing_event( + 'questionnaire_detail_loaded', + duration_ms=round((time.perf_counter() - questionnaire_fetch_started) * 1000, 3), + ) replies = km_data['replies'] km = km_data['knowledgeModel'] @@ -133,65 +144,116 @@ async def run_pipeline( if on_progress is not None: on_progress('Preparing document template') - result = await pipeline.run_async( - data={ - 'loader_component': { - 'knowledge_model_uuid': knowledge_model_uuid, - 'template_uuid': template_uuid, - }, - 'parser_component': {'data': km_data}, - 'assignment_component': { - 'template_data': template_data, - 'km': km, - 'on_progress': on_progress, - }, - 'assignment_saver_component': { - 'knowledge_model_uuid': knowledge_model_uuid, - 'knowledge_model_name': knowledge_model_name, - 'knowledge_model_version': knowledge_model_version, - 'template_uuid': template_uuid, - 'template_title': template_title, - 'template_data': template_data, - 'tenant_uuid': tenant_uuid, - }, - 'dmp_generator_component': { - 'replies': replies, - 'km': km, - 'on_progress': on_progress, - }, - 'dmp_polisher_component': { - 'template_data': template_data, - 'on_progress': on_progress, + pipeline_started = time.perf_counter() + try: + result = await pipeline.run_async( + data={ + 'loader_component': { + 'knowledge_model_uuid': knowledge_model_uuid, + 'template_uuid': template_uuid, + }, + 'parser_component': {'data': km_data}, + 'assignment_component': { + 'template_data': template_data, + 'km': km, + 'on_progress': on_progress, + }, + 'assignment_saver_component': { + 'knowledge_model_uuid': knowledge_model_uuid, + 'knowledge_model_name': knowledge_model_name, + 'knowledge_model_version': knowledge_model_version, + 'template_uuid': template_uuid, + 'template_title': template_title, + 'template_data': template_data, + 'tenant_uuid': tenant_uuid, + }, + 'dmp_generator_component': { + 'replies': replies, + 'km': km, + 'on_progress': on_progress, + }, + 'dmp_polisher_component': { + 'template_data': template_data, + 'on_progress': on_progress, + }, + 'saver_component': { + 'template_uuid': template_uuid, + 'knowledge_model_uuid': knowledge_model_uuid, + 'user_uuid': user_uuid, + 'tenant_uuid': tenant_uuid, + }, }, - 'saver_component': { - 'template_uuid': template_uuid, - 'knowledge_model_uuid': knowledge_model_uuid, - 'user_uuid': user_uuid, - 'tenant_uuid': tenant_uuid, + include_outputs_from={ + 'assignment_saver_component', + 'dmp_generator_component', + 'dmp_polisher_component', + 'saver_component', }, - }, - include_outputs_from={ - 'assignment_saver_component', - 'dmp_generator_component', - 'dmp_polisher_component', - 'saver_component', - }, + ) + except Exception: + logger.exception( + 'Pipeline component execution failed', + extra={'questionnaire_uuid': str(questionnaire_uuid), 'template_uuid': str(template_uuid)}, + ) + raise + log_timing_event( + 'pipeline_components_finished', + duration_ms=round((time.perf_counter() - pipeline_started) * 1000, 3), ) result_markdown = get_component_markdown(result, 'saver_component') if result_markdown is None: msg = 'Missing markdown output from saver_component' + logger.error(msg, extra={'template_uuid': str(template_uuid)}) raise RuntimeError(msg) - await write_metrics( - database, - template_uuid, - knowledge_model_uuid, - user_uuid, - tenant_uuid, - result, - model_name, - t1, + assignment_stats = get_component_stats(result, 'assignment_saver_component') + generation_stats = get_component_stats(result, 'dmp_generator_component') + polishing_stats = get_component_stats(result, 'dmp_polisher_component') + + metrics_started = time.perf_counter() + try: + await write_metrics( + database, + template_uuid, + knowledge_model_uuid, + user_uuid, + tenant_uuid, + result, + model_name, + t1, + ) + except Exception: + logger.exception( + 'Failed to persist pipeline metrics', + extra={'template_uuid': str(template_uuid), 'knowledge_model_uuid': str(knowledge_model_uuid)}, + ) + raise + log_timing_event( + 'pipeline_metrics_saved', + duration_ms=round((time.perf_counter() - metrics_started) * 1000, 3), + ) + log_timing_event( + 'pipeline_summary', + generation_ms=generation_stats.total_duration_ms if generation_stats is not None else None, + polishing_ms=polishing_stats.total_duration_ms if polishing_stats is not None else None, + total_pipeline_ms=round((time.perf_counter() - pipeline_total_started) * 1000, 3), + total_llm_wait_ms=round( + sum( + stats.total_llm_wait_ms + for stats in (assignment_stats, generation_stats, polishing_stats) + if stats is not None + ), + 3, + ), + total_llm_response_ms=round( + sum( + stats.total_llm_response_ms + for stats in (assignment_stats, generation_stats, polishing_stats) + if stats is not None + ), + 3, + ), ) return knowledge_model_uuid, result_markdown diff --git a/service/src/ai_document_plugin_service/api/auth.py b/service/src/ai_document_plugin_service/api/auth.py index 48647c5..9966059 100644 --- a/service/src/ai_document_plugin_service/api/auth.py +++ b/service/src/ai_document_plugin_service/api/auth.py @@ -1,3 +1,4 @@ +import logging from dataclasses import dataclass from typing import Annotated from uuid import UUID @@ -13,6 +14,7 @@ DSW_USER_VALIDATION_SUCCESS_STATUS = 200 DSW_ADMIN_ROLE = 'admin' DSW_ADMIN_PERMISSION = 'SettingsManageRolePermission' +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -51,6 +53,7 @@ def _parse_bearer_token(authorization: str | None) -> str | None: def _fetch_dsw_user(api_url: str, token: str) -> dict[str, object] | None: """Validate the token against DSW and return the current user, or None if rejected.""" url = f'{normalize_project_url(api_url)}/users/current' + logger.debug('Validating DSW user against current-user endpoint', extra={'url.full': url}) try: response = httpx.get( url, @@ -58,17 +61,24 @@ def _fetch_dsw_user(api_url: str, token: str) -> dict[str, object] | None: timeout=DSW_USER_VALIDATION_TIMEOUT_SECONDS, ) except httpx.HTTPError: + logger.exception('DSW user validation request failed', extra={'url.full': url}) return None if response.status_code != DSW_USER_VALIDATION_SUCCESS_STATUS: + logger.error( + 'DSW user validation rejected the token', + extra={'url.full': url, 'http.response.status_code': response.status_code}, + ) return None try: user = response.json() except ValueError: + logger.exception('DSW user validation returned invalid JSON', extra={'url.full': url}) return None if not isinstance(user, dict): + logger.error('DSW user validation returned an invalid payload', extra={'url.full': url}) return None return user @@ -89,6 +99,7 @@ def _is_admin(user: dict[str, object], api_url: str) -> bool: f'did not include a valid "role" string or a "role" object with a "permissions" list. ' f'The tenant may be running an incompatible DSW version. Received payload: {user}' ) + logger.error('DSW user validation returned an unsupported role payload', extra={'url.full': api_url}) raise ValueError(msg) @@ -99,9 +110,11 @@ def verify_authenticated( ) -> AuthenticatedUser: token = _parse_bearer_token(authorization) if token is None: + logger.error('Authentication failed: missing or invalid Authorization header') raise fastapi.HTTPException(status_code=401, detail='Unauthorized') if dsw_api_url is None or not dsw_api_url.strip(): + logger.error('Authentication failed: missing DSW API URL header') raise fastapi.HTTPException(status_code=401, detail='Unauthorized') config: Config = request.app.state.config @@ -110,13 +123,25 @@ def verify_authenticated( try: user_uuid, tenant_uuid = extract_identity_from_token(token) except ValueError as error: + logger.exception('Authentication failed: bearer token could not be parsed') raise fastapi.HTTPException(status_code=400, detail=str(error)) from error if not is_allowed_request(normalized_api_url, tenant_uuid, config.allowed_apis): + logger.error( + 'Authentication failed: request target is not allowed by configuration', + extra={ + 'tenant_uuid': str(tenant_uuid), + 'url.full': normalized_api_url, + }, + ) raise fastapi.HTTPException(status_code=401, detail='Unauthorized') user = _fetch_dsw_user(normalized_api_url, token) if user is None: + logger.error( + 'Authentication failed: DSW current-user validation rejected the token', + extra={'tenant_uuid': str(tenant_uuid), 'url.full': normalized_api_url}, + ) raise fastapi.HTTPException(status_code=401, detail='Unauthorized') return AuthenticatedUser( diff --git a/service/src/ai_document_plugin_service/api/jwt.py b/service/src/ai_document_plugin_service/api/jwt.py index d3cd431..4233635 100644 --- a/service/src/ai_document_plugin_service/api/jwt.py +++ b/service/src/ai_document_plugin_service/api/jwt.py @@ -1,15 +1,18 @@ import base64 import json +import logging import uuid from uuid import UUID JWT_PART_COUNT = 2 +logger = logging.getLogger(__name__) def decode_jwt_payload(token: str) -> dict[str, object]: parts = token.split('.') if len(parts) < JWT_PART_COUNT: msg = 'Invalid JWT token format.' + logger.error('JWT decode failed: invalid token format') raise ValueError(msg) payload = parts[1] @@ -20,10 +23,12 @@ def decode_jwt_payload(token: str) -> dict[str, object]: parsed = json.loads(decoded) except (ValueError, json.JSONDecodeError) as exc: msg = 'Invalid JWT token payload.' + logger.exception('JWT decode failed: invalid token payload', exc_info=exc) raise ValueError(msg) from exc if not isinstance(parsed, dict): msg = 'Invalid JWT token payload.' + logger.error('JWT decode failed: parsed payload is not an object') raise TypeError(msg) return parsed @@ -37,9 +42,11 @@ def _get_required_uuid_claim(payload: dict[str, object], *keys: str) -> str: # TODO: https://github.com/ds-wizard/ai-document-plugin/issues/61 return str(UUID(value)) except ValueError: + logger.exception('JWT claim is not a valid UUID', extra={'claim_name': key}) continue msg = f'Missing required JWT claim: {", ".join(keys)}' + logger.error('JWT decode failed: missing required UUID claim', extra={'claim_names': ','.join(keys)}) raise ValueError(msg) diff --git a/service/src/ai_document_plugin_service/api/request_logging.py b/service/src/ai_document_plugin_service/api/request_logging.py new file mode 100644 index 0000000..aa76953 --- /dev/null +++ b/service/src/ai_document_plugin_service/api/request_logging.py @@ -0,0 +1,83 @@ +import logging +import time +import uuid +from collections.abc import Awaitable, Callable + +import fastapi +from starlette.types import Message + +from ai_document_plugin_service.ai.common.logging_payloads import summarize_http_body +from ai_document_plugin_service.ai.common.trace_context import trace_context + +logger = logging.getLogger(__name__) + +TRACE_UUID_HEADER = 'X-Trace-UUID' + + +async def log_http_request_response( + request: fastapi.Request, + call_next: Callable[[fastapi.Request], Awaitable[fastapi.Response]], +) -> fastapi.Response: + trace_uuid = str(uuid.uuid4()) + request.state.trace_uuid = trace_uuid + request_body = await request.body() + _restore_request_body(request, request_body) + + with trace_context(trace_uuid): + logger.info( + 'HTTP request started', + extra={ + 'http.request.method': request.method, + 'url.path': request.url.path, + 'url.query': request.url.query, + }, + ) + request_body_summary = summarize_http_body( + request_body, + content_type=request.headers.get('content-type'), + ) + if request_body_summary is not None: + logger.debug( + 'HTTP request body', + extra={ + 'url.path': request.url.path, + 'http.request.body': request_body_summary, + }, + ) + + started_at = time.perf_counter() + try: + response = await call_next(request) + except Exception: + logger.exception( + 'HTTP request failed with unhandled exception', + extra={ + 'http.request.method': request.method, + 'url.path': request.url.path, + 'duration_ms': round((time.perf_counter() - started_at) * 1000, 3), + }, + ) + raise + + duration_ms = round((time.perf_counter() - started_at) * 1000, 3) + logger.info( + 'HTTP request completed', + extra={ + 'http.request.method': request.method, + 'url.path': request.url.path, + 'http.response.status_code': response.status_code, + 'duration_ms': duration_ms, + }, + ) + response.headers[TRACE_UUID_HEADER] = trace_uuid + return response + + +def _restore_request_body(request: fastapi.Request, body: bytes) -> None: + async def receive() -> Message: # noqa: RUF029 + return { + 'type': 'http.request', + 'body': body, + } + + request._receive = receive # type: ignore[method-assign] # noqa: SLF001 diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index e4ff95d..a14af78 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -1,3 +1,4 @@ +import logging from typing import Annotated from uuid import UUID @@ -25,6 +26,8 @@ from ai_document_plugin_service.service.errors import NotFoundError from ai_document_plugin_service.utils.docx_export import DOCX_MEDIA_TYPE +logger = logging.getLogger(__name__) + public_router = fastapi.APIRouter() protected_router = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_authenticated)]) @@ -68,6 +71,7 @@ async def delete_template(template_uuid: UUID, templates: TemplateServiceDI, aut @protected_router.post('/pipelines/run') async def start_pipeline( + request: fastapi.Request, payload: PipelineRunRequest, auth: AuthenticatedDI, config: ConfigDI, @@ -81,6 +85,7 @@ async def start_pipeline( template.title, auth, config, + getattr(request.state, 'trace_uuid', '-'), ) status = await pipeline.get_pipeline_status(run_id, auth) if status is None: @@ -113,6 +118,15 @@ async def get_pipeline_status( async def save_pipeline_result( run_id: UUID, save_request: PipelineSaveRequest, pipeline: PipelineServiceDI, auth: AuthenticatedDI ) -> PipelineStatusResponse: + logger.info( + 'Pipeline result update requested', + extra={ + 'run_id': run_id, + 'tenant_uuid': str(auth.tenant_uuid), + 'user_uuid': str(auth.user_uuid), + 'result_markdown_length': len(save_request.result_markdown), + }, + ) return await pipeline.update_pipeline_result(run_id, save_request, auth) diff --git a/service/src/ai_document_plugin_service/app.py b/service/src/ai_document_plugin_service/app.py index 72f3d98..a032097 100644 --- a/service/src/ai_document_plugin_service/app.py +++ b/service/src/ai_document_plugin_service/app.py @@ -1,3 +1,5 @@ +import logging + import fastapi import fastapi.middleware.cors from starlette.requests import Request @@ -6,10 +8,13 @@ from ai_document_plugin_service.ai.common import configure_logging from ai_document_plugin_service.ai.common.config import load_config, resolve_config_path from ai_document_plugin_service.ai.persistence.migrations import run_startup_migrations +from ai_document_plugin_service.api.request_logging import log_http_request_response from ai_document_plugin_service.api.routes import protected_router, public_router from ai_document_plugin_service.di import setup_app_state from ai_document_plugin_service.service.errors import ServiceError +logger = logging.getLogger(__name__) + def create_app(*, run_migrations: bool = True) -> fastapi.FastAPI: config_path = resolve_config_path() @@ -21,9 +26,20 @@ def create_app(*, run_migrations: bool = True) -> fastapi.FastAPI: app = fastapi.FastAPI(title='Plugin Service', version='1.0.0') setup_app_state(app, config) + app.middleware('http')(log_http_request_response) @app.exception_handler(ServiceError) - def service_error_handler(_request: Request, exc: ServiceError) -> JSONResponse: + def service_error_handler(request: Request, exc: ServiceError) -> JSONResponse: + logger.error( + 'API request rejected', + extra={ + 'http.request.method': request.method, + 'url.path': request.url.path, + 'http.response.status_code': exc.status_code, + 'error_type': type(exc).__name__, + 'error_message': exc.detail, + }, + ) return JSONResponse(status_code=exc.status_code, content={'detail': exc.detail}) app.add_middleware( diff --git a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py b/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py index 33862b4..7512621 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py +++ b/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py @@ -6,6 +6,8 @@ from typing import Any from uuid import UUID +from ai_document_plugin_service.ai.common.trace_context import trace_context + logger = logging.getLogger(__name__) JobFactory = Callable[[], Coroutine[Any, Any, None]] @@ -38,17 +40,23 @@ def __init__(self, max_concurrent_jobs: int) -> None: daemon=True, ) self._thread.start() + logger.info( + 'Initialized pipeline queue manager', + extra={'max_concurrent_jobs': max_concurrent_jobs, 'thread_name': self._thread.name}, + ) def _run_loop(self) -> None: asyncio.set_event_loop(self._loop) self._loop.run_forever() - def enqueue(self, run_id: UUID, job: JobFactory) -> None: + def enqueue(self, run_id: UUID, job: JobFactory, *, trace_id: str = '-') -> None: with self._order_lock: self._order.append(run_id) + queue_size = len(self._order) + logger.info('Enqueued pipeline job', extra={'run_id': run_id, 'queue_size': queue_size}) - future = asyncio.run_coroutine_threadsafe(self._run_job(run_id, job), self._loop) - future.add_done_callback(self._log_job_failure) + future = asyncio.run_coroutine_threadsafe(self._run_job(run_id, job, trace_id), self._loop) + future.add_done_callback(lambda done_future: self._log_job_failure(done_future, trace_id)) def progress_message(self, run_id: UUID) -> str | None: jobs_waiting_ahead = self._jobs_waiting_ahead(run_id) @@ -60,6 +68,7 @@ def remove(self, run_id: UUID) -> None: with self._order_lock: if run_id in self._order: self._order.remove(run_id) + logger.debug('Removed pipeline job from queue order tracking', extra={'run_id': run_id}) def _jobs_waiting_ahead(self, run_id: UUID) -> int | None: with self._order_lock: @@ -69,17 +78,21 @@ def _jobs_waiting_ahead(self, run_id: UUID) -> int | None: return None return queue_index - self._max_concurrent_jobs - async def _run_job(self, run_id: UUID, job: JobFactory) -> None: - try: - async with self._semaphore: - await job() - finally: - self.remove(run_id) + async def _run_job(self, run_id: UUID, job: JobFactory, trace_id: str) -> None: + with trace_context(trace_id): + try: + async with self._semaphore: + logger.info('Starting queued pipeline job', extra={'run_id': run_id}) + await job() + finally: + self.remove(run_id) + logger.info('Finished queued pipeline job', extra={'run_id': run_id}) @staticmethod - def _log_job_failure(future: Future[None]) -> None: + def _log_job_failure(future: Future[None], trace_id: str) -> None: # Jobs are expected to handle their own errors; this guards against an # unhandled exception being silently swallowed by the background loop. - error = future.exception() - if error is not None: - logger.error('Pipeline job crashed without handling its error', exc_info=error) + with trace_context(trace_id): + error = future.exception() + if error is not None: + logger.error('Pipeline job crashed without handling its error', exc_info=error) diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index 44e516e..e5299b6 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -10,6 +10,9 @@ Config, LLMConfig, ) +from ai_document_plugin_service.ai.common.execution_logging import ( + log_timing_event, +) from ai_document_plugin_service.ai.common.llm_client import LLMClient from ai_document_plugin_service.ai.knowledgemodel.dsw_client import DSWClient from ai_document_plugin_service.ai.persistence.assignment_saver_component import DBSaver @@ -131,6 +134,7 @@ async def get_pipeline_status(self, run_id: UUID, auth: AuthenticatedUser) -> Pi # Handle queued progress message (x dmps ahead in queue) progress_message = self.pipeline_queue_manager.progress_message(run_id) if progress_message is None: + logger.debug('Queued pipeline status has no queue progress message', extra={'run_id': run_id}) return status return status.model_copy(update={'progress_message': progress_message}) @@ -145,6 +149,7 @@ async def enqueue_pipeline_job( title: str, auth: AuthenticatedUser, config: Config, + trace_id: str, ) -> UUID: """Queue a pipeline job; concurrency is limited by ``pipeline_queue_manager``.""" run_id = await self.database.create_generation( @@ -172,6 +177,7 @@ async def enqueue_pipeline_job( llm_config, config, ), + trace_id=trace_id ) return run_id @@ -215,16 +221,22 @@ async def _run_pipeline_job( try: await self._run_pipeline(run_id, questionnaire_uuid, template_uuid, auth, llm_config, config) except Exception as error: - logger.exception('Pipeline run failed') + logger.exception('Pipeline run failed', extra={'run_id': run_id, 'tenant_uuid': str(auth.tenant_uuid)}) pipeline_error = _pipeline_error_from_exception(error) - await self.database.update_generation( - run_id, - auth.tenant_uuid, - status=PipelineStatus.FAILED, - error_type=pipeline_error.type, - error_message=pipeline_error.message, - progress_message=None, - ) + try: + await self.database.update_generation( + run_id, + auth.tenant_uuid, + status=PipelineStatus.FAILED, + error_type=pipeline_error.type, + error_message=pipeline_error.message, + progress_message=None, + ) + except Exception: + logger.exception( + 'Failed to persist pipeline failure status', + extra={'run_id': run_id, 'tenant_uuid': str(auth.tenant_uuid)}, + ) async def _run_pipeline( self, @@ -244,6 +256,10 @@ async def _run_pipeline( error_type=ErrorType.TEMPLATE_NOT_FOUND, error_message=TEMPLATE_NOT_FOUND_MESSAGE, ) + logger.error( + 'Pipeline run failed because template was not found', + extra={'run_id': run_id, 'template_uuid': str(template_uuid), 'tenant_uuid': str(auth.tenant_uuid)}, + ) return await self.database.update_generation( @@ -261,6 +277,15 @@ async def _run_pipeline( config=config, llm_client=llm_client, ) + logger.info( + 'Pipeline graph built and LLM client configured', + extra={ + 'run_id': run_id, + 'tenant_uuid': str(auth.tenant_uuid), + 'llm_model': llm_client.get_model_name(), + 'llm_max_workers': llm_client.get_max_workers(), + }, + ) def on_progress(message: str) -> None: # Called synchronously from deep inside the (async) pipeline; fire the DB @@ -283,6 +308,7 @@ def on_progress(message: str) -> None: model_name=llm_client.get_model_name(), dsw_client=DSWClient(auth.token, auth.api_url), ) + log_timing_event('pipeline_generation_finished', knowledge_model_uuid=str(knowledge_model_uuid)) await self.database.update_generation( run_id, @@ -292,3 +318,12 @@ def on_progress(message: str) -> None: result_markdown=result, progress_message=None, ) + logger.info( + 'Pipeline run status updated to succeeded', + extra={ + 'run_id': run_id, + 'tenant_uuid': str(auth.tenant_uuid), + 'knowledge_model_uuid': str(knowledge_model_uuid), + 'result_markdown_length': len(result), + }, + ) diff --git a/service/tests/common/test_pipeline_metrics.py b/service/tests/common/test_pipeline_metrics.py index 9fa2177..00cd06f 100644 --- a/service/tests/common/test_pipeline_metrics.py +++ b/service/tests/common/test_pipeline_metrics.py @@ -92,4 +92,6 @@ def test_get_stats_returns_json_summary_with_expected_shape() -> None: 'cost_per_mil_input': 0.25, 'cost_per_mil_output': 2.0, 'elapsed_seconds': 12.5, + 'total_llm_wait_ms': 0.0, + 'total_llm_response_ms': 0.0, } diff --git a/service/uv.lock b/service/uv.lock index 8ee4d18..33ffdc6 100644 --- a/service/uv.lock +++ b/service/uv.lock @@ -30,7 +30,6 @@ dependencies = [ { name = "sqlalchemy", extra = ["asyncio"] }, { name = "tabulate" }, { name = "tiktoken" }, - { name = "tqdm" }, ] [package.dev-dependencies] @@ -59,7 +58,6 @@ requires-dist = [ { name = "sqlalchemy", extras = ["asyncio"] }, { name = "tabulate" }, { name = "tiktoken" }, - { name = "tqdm" }, ] [package.metadata.requires-dev]