diff --git a/py/README.md b/py/README.md index b0b3242dc..c4a3b2afb 100644 --- a/py/README.md +++ b/py/README.md @@ -41,6 +41,39 @@ Then run: BRAINTRUST_API_KEY= braintrust eval tutorial_eval.py ``` +## Logging With an Ingestion Key + +Ingestion keys let apps that run on your users' machines, like a desktop app or a CLI, send traces to a single project without an API key. +Pass the ingestion key URL from your project settings to `init_logger`, or set `BRAINTRUST_INGESTION_KEY`: + +```python +import braintrust + +logger = braintrust.init_logger(ingestion_key="https://dp.example.com/ingest?ingestKey=bt-ik-...") +logger.log(input="question", output="answer") +``` + +The logger sends rows and attachments straight to the data plane in the URL, with the key in the `Authorization` header. +It never logs in, so it ignores `BRAINTRUST_API_KEY` and any earlier `braintrust.login()`, and it doesn't register or look up the project. +Passing an explicit `api_key` to `init_logger` uses that key instead of `BRAINTRUST_INGESTION_KEY`, and passing both `api_key` and `ingestion_key` is an error. +Feedback logged with an ingestion key can only include scores, expected values, and tags. +`Attachment` works as usual, but `ExternalAttachment` isn't supported, since its URL can point at content outside the project. + +Ingestion keys don't own the rows they write. +Anyone holding the key can update any row of the project whose ID they know, including rows logged with an API key or with another ingestion key. +Row IDs aren't secrets, so treat every row in the project as writable by whoever holds the key. + +To export OpenTelemetry spans with an ingestion key, point a standard OTLP exporter at the `/otel/v1/traces` path under the ingestion URL and drop the query string: + +```python +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + +exporter = OTLPSpanExporter( + endpoint="https://dp.example.com/ingest/otel/v1/traces", + headers={"Authorization": "Bearer bt-ik-..."}, +) +``` + ## Optional Extras Install extras as needed for specific workflows: diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index f08afb668..aff3b50f5 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -214,6 +214,7 @@ def __get__(self, instance: object, owner: type | None = None) -> bool: class BraintrustEnv: API_KEY = EnvVar("BRAINTRUST_API_KEY", EnvParser.STRING) + INGESTION_KEY = EnvVar("BRAINTRUST_INGESTION_KEY", EnvParser.STRING) API_URL = EnvVar("BRAINTRUST_API_URL", EnvParser.STRING) PROXY_URL = EnvVar("BRAINTRUST_PROXY_URL", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 21155a2dd..1f7996b6c 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -14,6 +14,7 @@ import json import logging import os +import re import sys import textwrap import threading @@ -37,7 +38,7 @@ cast, overload, ) -from urllib.parse import quote, urlencode +from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit import chevron import exceptiongroup @@ -47,7 +48,7 @@ from . import context, id_gen from .api._routing import normalize_proxy_url -from .api._transport import HTTPConnection +from .api._transport import HTTPConnection, _parse_retry_after from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter from .api.auth import LoginResult, OrganizationInfo from .api.client import BraintrustClient, BraintrustOpenApiClient @@ -57,12 +58,16 @@ BraintrustJSONDecodeError, BraintrustTransportError, ) -from .api.policies import DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception +from .api.policies import DEFAULT_MAX_ELAPSED_TIME, DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( + ARRAY_DELETE_FIELD, AUDIT_METADATA_FIELD, AUDIT_SOURCE_FIELD, + CREATED_FIELD, + ID_FIELD, IS_MERGE_FIELD, + MERGE_PATHS_FIELD, OBJECT_DELETE_FIELD, OBJECT_ID_KEYS, TRANSACTION_ID_FIELD, @@ -616,6 +621,12 @@ def default_get_api_conn(): lambda: _HTTPBackgroundLogger(LazyValue(default_get_api_conn, use_mutex=True)), use_mutex=True ) + # Ingestion key loggers publish independently of the private login, + # with one background logger per ingestion endpoint and key. + self._public_bg_loggers: dict[_IngestionEndpoint, _PublicHTTPBackgroundLogger] = {} + self._public_bg_loggers_lock = threading.Lock() + self._masking_function: Callable[[Any], Any] | None = None + self._id_generator = None # For unit-testing, tests may wish to temporarily override the global @@ -750,6 +761,9 @@ def copy_state(self, other: "BraintrustState"): "current_span", "_global_bg_logger", "_override_bg_logger", + "_public_bg_loggers", + "_public_bg_loggers_lock", + "_masking_function", "_context_manager", "_last_otel_setting", "_context_manager_lock", @@ -876,6 +890,15 @@ def loader_api_client(self, options: "_LoaderLoginOptions") -> "Iterator[Braintr def global_bg_logger(self) -> "_BackgroundLogger": return getattr(self._override_bg_logger, "logger", None) or self._global_bg_logger.get() + def public_bg_logger(self, endpoint: "_IngestionEndpoint") -> "_PublicHTTPBackgroundLogger": + with self._public_bg_loggers_lock: + bg_logger = self._public_bg_loggers.get(endpoint) + if bg_logger is None: + bg_logger = _PublicHTTPBackgroundLogger(endpoint) + bg_logger.set_masking_function(self._masking_function) + self._public_bg_loggers[endpoint] = bg_logger + return bg_logger + # Should only be called by the login function. def login_replace_api_conn(self, api_conn: "HTTPConnection"): self._global_bg_logger.get().internal_replace_api_conn(api_conn) @@ -891,8 +914,12 @@ def enforce_queue_size_limit(self, enforce: bool) -> None: bg_logger.enforce_queue_size_limit(enforce) def set_masking_function(self, masking_function: Callable[[Any], Any] | None) -> None: - """Set the masking function on the background logger.""" + """Set the masking function on the background loggers.""" self.global_bg_logger().set_masking_function(masking_function) + with self._public_bg_loggers_lock: + self._masking_function = masking_function + for bg_logger in self._public_bg_loggers.values(): + bg_logger.set_masking_function(masking_function) _state: BraintrustState = None # type: ignore @@ -922,6 +949,10 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: _state._api_conn._reset() if _state._client: _state._client.transport._set_adapter(adapter) + with _state._public_bg_loggers_lock: + for bg_logger in _state._public_bg_loggers.values(): + bg_logger.api_conn.get()._set_adapter(adapter=adapter) + bg_logger.api_conn.get()._reset() # Per-credential loader resources may have been created with the previous # adapter. Eviction closes them once any active requests release their lease; @@ -1210,35 +1241,7 @@ def flush(self, batch_size: int | None = None): if len(all_items) == 0: return - # Construct batches of records to flush in parallel. - all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] - max_request_size_result = self._get_max_request_size() - batches = batch_items( - items=all_items_with_meta, - batch_max_num_items=batch_size, - batch_max_num_bytes=max_request_size_result["max_request_size"] // 2, - get_byte_size=lambda item: len(item.str_value), - ) - - post_promises = [] - try: - post_promises = [ - HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) - for batch in batches - ] - except RuntimeError: - # If the thread pool has shut down, e.g. because the process - # is terminating, run the requests the old fashioned way. - for batch in batches: - self._submit_logs_request(batch, max_request_size_result) - - concurrent.futures.wait(post_promises) - # Raise any exceptions from the promises as one group. - post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] - if post_promise_exceptions: - raise exceptiongroup.BaseExceptionGroup( - f"Encountered the following errors while logging:", post_promise_exceptions - ) + self._publish_items(all_items, batch_size) attachment_errors: list[Exception] = [] for attachment in attachments: @@ -1257,12 +1260,45 @@ def flush(self, batch_size: int | None = None): attachment_errors, ) + def _publish_items(self, all_items: list[dict[str, Any]], batch_size: int) -> None: + # Construct batches of records to flush in parallel. + all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] + max_request_size_result = self._get_max_request_size() + batches = batch_items( + items=all_items_with_meta, + batch_max_num_items=batch_size, + batch_max_num_bytes=max_request_size_result["max_request_size"] // 2, + get_byte_size=lambda item: len(item.str_value), + ) + + post_promises = [] + try: + post_promises = [ + HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) + for batch in batches + ] + except RuntimeError: + # If the thread pool has shut down, e.g. because the process + # is terminating, run the requests the old fashioned way. + for batch in batches: + self._submit_logs_request(batch, max_request_size_result) + + concurrent.futures.wait(post_promises) + # Raise any exceptions from the promises as one group. + post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] + if post_promise_exceptions: + raise exceptiongroup.BaseExceptionGroup( + f"Encountered the following errors while logging:", post_promise_exceptions + ) + def _unwrap_lazy_values( - self, wrapped_items: Sequence[_LazyLogRecord] + self, wrapped_items: Sequence[_LazyLogRecord], extract_attachments: bool = True ) -> tuple[list[dict[str, Any]], list["BaseAttachment"]]: for i in range(self.num_tries): try: - unwrapped_items = [record for item in wrapped_items if (record := item.get()) is not None] + unwrapped_items = [ + self._prepare_record(record) for item in wrapped_items if (record := item.get()) is not None + ] merged_items = merge_row_batch(unwrapped_items) # Logger-local hooks run after span hooks and merging. @@ -1274,8 +1310,9 @@ def _unwrap_lazy_values( ] attachments: list["BaseAttachment"] = [] - for item in merged_items: - _extract_attachments(item, attachments) + if extract_attachments: + for item in merged_items: + _extract_attachments(item, attachments) return merged_items, attachments except Exception as e: @@ -1300,6 +1337,9 @@ def _unwrap_lazy_values( ) return [], [] + def _prepare_record(self, record: dict[str, Any]) -> dict[str, Any]: + return record + def _request_logs3_overflow_upload( self, conn: HTTPConnection, payload_size_bytes: int, rows: list[dict[str, Any]] ) -> dict[str, Any]: @@ -1488,6 +1528,316 @@ def set_masking_function(self, masking_function: Callable[[Any], Any] | None): self._export_customizers = (_MaskingCustomizer(masking_function),) if masking_function is not None else () +INGESTION_KEY_QUERY_PARAM = "ingestKey" +_INGESTION_KEY_PATTERN = re.compile(r"bt-ik-[A-Za-z0-9]{48}") +# Ingestion keys may only write trace rows, so rows never carry other objects' ids or writable provenance. +_INGESTION_ROW_FIELDS = frozenset( + { + ID_FIELD, + "span_id", + "root_span_id", + "span_parents", + CREATED_FIELD, + "org_id", + "project_id", + "log_id", + "input", + "output", + "expected", + "error", + "tags", + "scores", + "metadata", + "metrics", + "context", + "span_attributes", + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + ARRAY_DELETE_FIELD, + OBJECT_DELETE_FIELD, + } +) +_UPLOAD_REFERENCE_TYPES = {"attachment": "braintrust_attachment", "logs3_overflow": LOGS3_OVERFLOW_REFERENCE_TYPE} +_MAX_UPLOAD_GRANT_LIFETIME_MS = 300_000 +# Requests above this go through an overflow upload. Lambda deployments +# base64-encode binary bodies, and data plane ingress accepts 1 MiB by default. +DEFAULT_INGESTION_MAX_REQUEST_SIZE = 512 * 1024 + + +@dataclasses.dataclass(frozen=True) +class _IngestionEndpoint: + """An ingestion key URL split into the ingestion root and the key it carries.""" + + url: str + key: str = dataclasses.field(repr=False) + + +def _parse_ingestion_url(value: str) -> _IngestionEndpoint: + # Errors never include `value`, since it contains the key. + try: + parsed = urlsplit(value.strip()) + parsed.port # Raises for malformed ports. + # Matches URLSearchParams: empty segments are skipped and values are percent-decoded. + query = parse_qsl(parsed.query, keep_blank_values=True) + reason = None + except ValueError: + reason = "it is not a valid URL" + if reason is None: + path = parsed.path.rstrip("/") + if parsed.scheme not in ("http", "https") or not parsed.hostname: + reason = "expected an http or https URL" + elif parsed.username is not None or parsed.password is not None: + reason = "URL credentials are not allowed" + elif "#" in value: + reason = "URL fragments are not allowed" + elif not path.endswith("/ingest"): + reason = "expected the URL path to end with /ingest" + elif ( + len(query) != 1 + or query[0][0] != INGESTION_KEY_QUERY_PARAM + or not _INGESTION_KEY_PATTERN.fullmatch(query[0][1]) + ): + reason = f"expected a single {INGESTION_KEY_QUERY_PARAM} query parameter holding a bt-ik- key" + else: + return _IngestionEndpoint(url=urlunsplit((parsed.scheme, parsed.netloc, path, "", "")), key=query[0][1]) + raise ValueError(f"Invalid Braintrust ingestion key URL: {reason}") + + +class _UploadGrantExpired(Exception): + pass + + +class _PublicHTTPBackgroundLogger(_HTTPBackgroundLogger): + """Publishes rows to a data plane with an ingestion key. + + The key only allows trace ingestion into a single project, so this logger + never logs in, discovers projects, or falls back to private credentials. + Attachments and oversized batches go through the data plane's chunked + upload API before the rows that reference them are published. + """ + + def __init__(self, endpoint: _IngestionEndpoint): + conn = HTTPConnection(endpoint.url, adapter=_http_adapter) + conn.set_token(endpoint.key) + super().__init__(LazyValue(lambda: conn, use_mutex=False)) + self.endpoint = endpoint + self.timeout = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + self._stripped_fields: set[str] = set() + self.enforce_queue_size_limit(True) + + def _get_max_request_size(self) -> dict[str, Any]: + # Any batch can overflow through the upload API, so there is no server limit to look up. + max_request_size = coalesce(self._max_request_size_override, DEFAULT_INGESTION_MAX_REQUEST_SIZE) + return {"max_request_size": max_request_size, "can_use_overflow": True} + + def flush(self, batch_size: int | None = None): + if batch_size is None: + batch_size = self.default_batch_size + + with self.flush_lock: + items, _ = self._unwrap_lazy_values(self.queue.drain_all(), extract_attachments=False) + if not items: + return + + # Attachment references are issued by the data plane, so attachments + # upload before the rows that point at them. + uploads: dict[int, AttachmentReference | Exception] = {} + rows: list[dict[str, Any]] = [] + attachment_errors: list[Exception] = [] + for item in items: + try: + _extract_attachments(item, [], lambda attachment: self._upload_attachment(attachment, uploads)) + except Exception as e: + attachment_errors.append(e) + print( + f"Dropping row {item.get(ID_FIELD)} because of an attachment that can't be uploaded: {e}", + file=self.outfile, + ) + continue + rows.append(item) + + if rows: + self._publish_items(rows, batch_size) + + if len(attachment_errors) == 1: + raise attachment_errors[0] + elif len(attachment_errors) > 1: + raise exceptiongroup.ExceptionGroup( + "Encountered errors while uploading attachments", + attachment_errors, + ) + + def _upload_attachment( + self, attachment: "BaseAttachment", uploads: "dict[int, AttachmentReference | Exception]" + ) -> AttachmentReference: + if isinstance(attachment, ExternalAttachment): + # The data plane rejects these, since their URL can point at any object store content. + raise ValueError("ExternalAttachment is not supported with ingestion keys. Use Attachment instead") + # The same attachment object can appear in several rows of one flush. + if id(attachment) not in uploads: + reference = attachment.reference + assert reference["type"] == "braintrust_attachment" + try: + uploads[id(attachment)] = self._upload( + "attachment", attachment.data, reference["content_type"], filename=reference["filename"] + ) + except Exception as e: + uploads[id(attachment)] = e + result = uploads[id(attachment)] + if isinstance(result, Exception): + raise result + return result + + def _prepare_record(self, record: dict[str, Any]) -> dict[str, Any]: + # Strip before merging, since merging groups rows by their object ids. + for field in [field for field in record if field not in _INGESTION_ROW_FIELDS]: + if field not in self._stripped_fields: + self._stripped_fields.add(field) + self.logger.warning(f"Ingestion keys cannot write the {field!r} field. Dropping it from logged rows.") + del record[field] + return record + + def _submit_logs_request(self, items: Sequence[LogItemWithMeta], max_request_size_result: dict[str, Any]): + data_str = construct_logs3_data(items) + payload = data_str.encode("utf-8") + if self.all_publish_payloads_dir: + _HTTPBackgroundLogger._write_payload_to_dir(payload_dir=self.all_publish_payloads_dir, payload=data_str) + + start_time = time.time() + try: + if len(payload) > max_request_size_result["max_request_size"]: + reference = self._upload("logs3_overflow", payload, "application/json") + self._request("post", "/v1/logs", json={"api_version": DATA_API_VERSION, "rows": reference}) + self._overflow_upload_count += 1 + else: + self._request("post", "/v1/logs", data=payload, headers={"Content-Type": "application/json"}) + except Exception as e: + if self.failed_publish_payloads_dir: + _HTTPBackgroundLogger._write_payload_to_dir( + payload_dir=self.failed_publish_payloads_dir, payload=data_str + ) + self._log_failed_payloads_dir() + errmsg = f"log request failed. Elapsed time: {time.time() - start_time} seconds. Payload size: {len(payload)}. Error: {e}" + if self.sync_flush: + raise Exception(errmsg) from e + print(errmsg, file=self.outfile) + print("Dropping batch", file=self.outfile) + + def _upload(self, purpose: str, data: bytes, content_type: str, filename: str | None = None) -> Any: + """Upload `data` in the chunks the data plane asks for and return the committed reference.""" + sha256 = hashlib.sha256(data).hexdigest() + body: dict[str, Any] = { + "purpose": purpose, + "size_bytes": len(data), + "content_type": content_type, + "sha256": sha256, + } + if filename is not None: + body["filename"] = filename + + for _ in range(self.num_tries): + # The grant's lifetime is relative, so measure it from before the request goes out. + requested_at = time.monotonic() + grant = self._request("post", "/v1/uploads", json=body) + chunk_bytes, num_chunks, expires_in_ms = ( + grant.get(k) for k in ("chunk_bytes", "num_chunks", "expires_in_ms") + ) + try: + # Normalizing the id keeps a malformed response from changing the request path. + upload_id = str(uuid.UUID(grant.get("upload_id"))) + except (TypeError, ValueError, AttributeError): + upload_id = None + if ( + upload_id is None + or not all(type(v) is int for v in (chunk_bytes, num_chunks, expires_in_ms)) + or chunk_bytes <= 0 + or num_chunks != -(-len(data) // chunk_bytes) + or not 0 < expires_in_ms <= _MAX_UPLOAD_GRANT_LIFETIME_MS + ): + raise RuntimeError(f"Invalid upload grant from the data plane: {self._redact(str(grant))}") + deadline = requested_at + expires_in_ms / 1000 + + try: + for index in range(num_chunks): + chunk = data[index * chunk_bytes : (index + 1) * chunk_bytes] + ack = self._request( + "put", + f"/v1/uploads/{upload_id}/chunks/{index}", + deadline=deadline, + data=chunk, + headers={"Content-Type": "application/octet-stream"}, + ) + if ack.get("index") != index or ack.get("size_bytes") != len(chunk): + raise RuntimeError( + f"Invalid chunk acknowledgement from the data plane: {self._redact(str(ack))}" + ) + result = self._request("post", f"/v1/uploads/{upload_id}/complete", deadline=deadline, json={}) + except _UploadGrantExpired: + print(f"Upload {upload_id} expired before it completed. Starting a new upload", file=self.outfile) + continue + + reference = result.get("reference") + if ( + result.get("size_bytes") != len(data) + or result.get("sha256") != sha256 + or not isinstance(reference, dict) + or reference.get("type") != _UPLOAD_REFERENCE_TYPES[purpose] + ): + raise RuntimeError( + f"Upload {upload_id} completed with unexpected content: {self._redact(str(result))}" + ) + return reference + + raise RuntimeError(f"Uploads kept expiring before they completed after {self.num_tries} attempts") + + def _request(self, method: str, path: str, deadline: float | None = None, **kwargs: Any) -> Any: + """Send one request, retrying transient failures, and return its JSON body. + + Requests against an upload grant pass its `deadline`. They never wait + past it, and raise `_UploadGrantExpired` once it passes or the data + plane reports the grant as gone. + """ + error = "" + for i in range(self.num_tries): + timeout = self.timeout + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise _UploadGrantExpired() + timeout = min(timeout, remaining) + retry_after = None + try: + # Redirects could forward the key somewhere else. + resp = getattr(self.api_conn.get(), method)(path, timeout=timeout, allow_redirects=False, **kwargs) + except requests_exceptions.RequestException as e: + if not is_retryable_request_exception(e): + raise RuntimeError(f"{method.upper()} {path} failed with {self._redact(str(e))}") from None + error = self._redact(str(e)) + else: + if 200 <= resp.status_code < 300: + return resp.json() + if deadline is not None and resp.status_code == 410: + raise _UploadGrantExpired() + error = self._redact(f"{resp.status_code}: {resp.text}") + if resp.status_code not in DEFAULT_RETRYABLE_STATUSES: + raise RuntimeError(f"{method.upper()} {path} failed with {error}") + retry_after = _parse_retry_after(resp.headers.get("Retry-After"), time.time()) + + if i + 1 < self.num_tries: + sleep_time_s = max( + BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i), min(retry_after or 0, DEFAULT_MAX_ELAPSED_TIME) + ) + if deadline is not None and time.monotonic() + sleep_time_s >= deadline: + raise _UploadGrantExpired() + print(f"{method.upper()} {path} failed with {error}. Retrying in {sleep_time_s}s", file=self.outfile) + time.sleep(sleep_time_s) + + raise RuntimeError(f"{method.upper()} {path} failed after {self.num_tries} attempts with {error}") + + def _redact(self, text: str) -> str: + return text.replace(self.endpoint.key, "") + + def _internal_reset_global_state() -> None: global _state _state = BraintrustState() @@ -1927,6 +2277,7 @@ def init_logger( state: BraintrustState | None = None, environment: SpanOriginEnvironment | None = None, _create_in_project_group: str | None = None, + ingestion_key: str | None = None, ) -> "Logger": """ Create a new logger in a specified project. If the project does not exist, it will be created. @@ -1941,11 +2292,26 @@ def init_logger( :param force_login: Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) :param set_current: If true (the default), set the global current-experiment to the newly-created one. :param _create_in_project_group: Experimental: the name of an existing project group to create the project in. Existing projects must already belong to the group, otherwise registration fails with a 409. Ignored when project_id is provided. + :param ingestion_key: (Optional) The ingestion key URL of a project, for example `https://dp.example.com/ingest?ingestKey=bt-ik-...`. + If neither this nor `api_key` is specified, will try to use the `BRAINTRUST_INGESTION_KEY` environment variable. + Ingestion keys can only write traces to the project they were issued for, so the logger never logs in, ignores `BRAINTRUST_API_KEY` + and any previous login, and doesn't register or look up the project. `project` and `project_id` are optional and only used to + label exported spans. A `project_id` is also sent with each row, and the data plane rejects rows for a different project. + Anyone holding the key can update any row of the project whose id they know, including rows logged with an API key or another + ingestion key. `ExternalAttachment` isn't supported, and rows containing one are dropped. :returns: The newly created Logger. """ state = state or _state state.span_origin_environment = detect_environment(environment) + + if api_key is not None and ingestion_key is not None: + raise ValueError("Pass either api_key or ingestion_key to init_logger, not both") + if api_key is None and ingestion_key is None: + ingestion_key = BraintrustEnv.INGESTION_KEY.get(None, use_dotenv=True) + if ingestion_key is not None and _create_in_project_group is not None: + raise ValueError("Loggers using an ingestion key cannot create projects") + compute_metadata_args = dict(project_name=project, project_id=project_id) if project_id is None and _create_in_project_group is not None: compute_metadata_args["project_group_name"] = _create_in_project_group @@ -1957,20 +2323,36 @@ def init_logger( "project_id": project_id, } - def compute_metadata(): - state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login) - return _compute_logger_metadata(**compute_metadata_args, state=state) + if ingestion_key is not None: + # The data plane resolves the org and project from the key, so there is nothing to look up. + public_metadata = OrgProjectMetadata( + org_id=cast(str, None), + project=ObjectMetadata(id=cast(str, project_id), name=project or "", full_info={}), + ) + ret = Logger( + lazy_metadata=LazyValue(lambda: public_metadata, use_mutex=False), + async_flush=async_flush, + compute_metadata_args=compute_metadata_args, + link_args=link_args, + state=state, + public_bg_logger=state.public_bg_logger(_parse_ingestion_url(ingestion_key)), + ) + else: - # For loggers, enable queue size limit enforcement (bounded queue) - state.enforce_queue_size_limit(True) + def compute_metadata(): + state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login) + return _compute_logger_metadata(**compute_metadata_args, state=state) - ret = Logger( - lazy_metadata=LazyValue(compute_metadata, use_mutex=True), - async_flush=async_flush, - compute_metadata_args=compute_metadata_args, - link_args=link_args, - state=state, - ) + # For loggers, enable queue size limit enforcement (bounded queue) + state.enforce_queue_size_limit(True) + + ret = Logger( + lazy_metadata=LazyValue(compute_metadata, use_mutex=True), + async_flush=async_flush, + compute_metadata_args=compute_metadata_args, + link_args=link_args, + state=state, + ) if set_current: if _state is None: raise RuntimeError("_state is None in init_logger. This should never happen.") @@ -2568,6 +2950,15 @@ def current_logger() -> "Logger | None": return _state._cv_logger.get() or _state._local_logger +def _current_public_logger() -> "Logger | None": + """Return the current logger if it uses an ingestion key and no experiment takes precedence.""" + + logger = current_logger() + if logger is None or logger._public_bg_logger is None or current_experiment() is not None: + return None + return logger + + def current_span() -> Span: """Return the currently-active span for logging (set by running a span under a context manager). If there is no active span, returns a no-op span object, which supports the same interface as spans but does no logging. @@ -2851,18 +3242,26 @@ def _resolve_w3c_parent(context: dict) -> "tuple[str | None, PropagatedState | N braintrust_parent = parse_baggage(baggage_value).get(BRAINTRUST_PARENT_KEY) if not braintrust_parent: braintrust_parent = _current_braintrust_parent() - if not braintrust_parent: + public_logger = _current_public_logger() + if not braintrust_parent and public_logger is not None: + # Ingestion key spans always go to the key's project, so the trace only needs span ids. + object_type, object_id, compute_args = ( + SpanObjectTypeV3.PROJECT_LOGS, + None, + public_logger._compute_metadata_args, + ) + elif not braintrust_parent: logging.warning( "Received traceparent without a braintrust.parent and no active logger/experiment; " "cannot route the trace. Starting a fresh local span instead." ) return None, None - - parsed_parent = _braintrust_parent_to_components(braintrust_parent) - if parsed_parent is None: - logging.warning(f"Invalid braintrust.parent: {braintrust_parent!r}") - return None, None - object_type, object_id, compute_args = parsed_parent + else: + parsed_parent = _braintrust_parent_to_components(braintrust_parent) + if parsed_parent is None: + logging.warning(f"Invalid braintrust.parent: {braintrust_parent!r}") + return None, None + object_type, object_id, compute_args = parsed_parent tracestate = get_header(context, TRACESTATE_HEADER) or None @@ -3103,10 +3502,25 @@ def start_span( # No row id (object-only slug): there is no span to link to, so start # a fresh root, keeping the object routing from the slug. parent_span_ids = None + public_logger = _current_public_logger() + if public_logger is not None: + # Resolving the parent's object needs a private login. Ingestion key + # spans stay in the key's project anyway, so the parent only + # contributes its span ids. + object_args = dict( + parent_object_type=public_logger._parent_object_type(), + parent_object_id=public_logger._lazy_id, + parent_compute_object_metadata_args=public_logger._compute_metadata_args, + public_bg_logger=public_logger._public_bg_logger, + ) + else: + object_args = dict( + parent_object_type=parent_obj.object_type, + parent_object_id=LazyValue(_span_components_to_object_id_lambda(parent_obj), use_mutex=False), + parent_compute_object_metadata_args=parent_obj.compute_object_metadata_args, + ) return SpanImpl( - parent_object_type=parent_obj.object_type, - parent_object_id=LazyValue(_span_components_to_object_id_lambda(parent_obj), use_mutex=False), - parent_compute_object_metadata_args=parent_obj.compute_object_metadata_args, + **object_args, parent_span_ids=parent_span_ids, name=name, type=type, @@ -3138,6 +3552,10 @@ def flush(): """Flush any pending rows to the server.""" _state.global_bg_logger().flush() + with _state._public_bg_loggers_lock: + public_bg_loggers = list(_state._public_bg_loggers.values()) + for public_bg_logger in public_bg_loggers: + public_bg_logger.flush() def _internal_start_span_with_initial_merge( @@ -3219,7 +3637,11 @@ def validate_tags(tags: Sequence[str]) -> None: seen.add(tag) -def _extract_attachments(event: dict[str, Any], attachments: list["BaseAttachment"]) -> None: +def _extract_attachments( + event: dict[str, Any], + attachments: list["BaseAttachment"], + get_reference: "Callable[[BaseAttachment], AttachmentReference] | None" = None, +) -> None: """ Helper function for uploading attachments. Recursively extracts `Attachment` and `ExternalAttachment` values and replaces them with their associated @@ -3227,13 +3649,15 @@ def _extract_attachments(event: dict[str, Any], attachments: list["BaseAttachmen :param event: The event to filter. Will be modified in-place. :param attachments: Flat array of extracted attachments (output parameter). + :param get_reference: Returns the reference that replaces an attachment. Defaults to `attachment.reference`. """ def _helper(v: Any) -> Any: # Base case: Attachment or ExternalAttachment. if isinstance(v, BaseAttachment): attachments.append(v) - return v.reference # Attachment cannot be nested. + # Attachment cannot be nested. + return get_reference(v) if get_reference else v.reference # Recursive case: object. if isinstance(v, dict): @@ -3909,6 +4333,13 @@ def __str__(self) -> str: return f"data:{self.reference['content_type']};base64,{b64_content}" +def _object_id_fields(object_type: SpanObjectTypeV3, object_id: str | None) -> dict[str, str]: + # Ingestion key loggers may not know their project. The data plane resolves it from the key. + if object_id is None and object_type == SpanObjectTypeV3.PROJECT_LOGS: + return {"log_id": "g"} + return _get_exporter()(object_type=object_type, object_id=object_id).object_id_fields() + + def _log_feedback_impl( parent_object_type: SpanObjectTypeV3, parent_object_id: LazyValue[str], @@ -3919,7 +4350,12 @@ def _log_feedback_impl( comment: str | None = None, metadata: Metadata | None = None, source: Literal["external", "app", "api", None] = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): + if public_bg_logger is not None and (comment is not None or metadata is not None or source is not None): + # Comments and audit fields are writable provenance, which ingestion keys cannot write. + raise ValueError("Feedback logged with an ingestion key can only include scores, expected, and tags") + if source is None: source = "external" elif source not in VALID_SOURCES: @@ -3943,13 +4379,11 @@ def _log_feedback_impl( update_event = {k: v for k, v in update_event.items() if v is not None} update_event = bt_safe_deep_copy(update_event) + audit_fields = {} if public_bg_logger else {AUDIT_SOURCE_FIELD: source, AUDIT_METADATA_FIELD: metadata} + bg_logger = public_bg_logger or _state.global_bg_logger() def parent_ids(): - exporter = _get_exporter() - return exporter( - object_type=parent_object_type, - object_id=parent_object_id.get(), - ).object_id_fields() + return _object_id_fields(parent_object_type, parent_object_id.get()) if len(update_event) > 0: @@ -3958,14 +4392,11 @@ def compute_update_record(): id=id, **update_event, **parent_ids(), - **{ - AUDIT_SOURCE_FIELD: source, - AUDIT_METADATA_FIELD: metadata, - IS_MERGE_FIELD: True, - }, + **audit_fields, + **{IS_MERGE_FIELD: True}, ) - _state.global_bg_logger().log(LazyValue(compute_update_record, use_mutex=False)) + bg_logger.log(LazyValue(compute_update_record, use_mutex=False)) if comment is not None: # pylint: disable=function-redefined @@ -3985,7 +4416,7 @@ def compute_comment_record(): **{AUDIT_SOURCE_FIELD: source, AUDIT_METADATA_FIELD: metadata}, ) - _state.global_bg_logger().log(LazyValue(compute_comment_record, use_mutex=False)) + bg_logger.log(LazyValue(compute_comment_record, use_mutex=False)) def _update_span_impl( @@ -3994,6 +4425,7 @@ def _update_span_impl( id: str, root_span_id: str | None, span_id: str | None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, **event: Any, ): if (root_span_id is None) != (span_id is None): @@ -4011,25 +4443,18 @@ def _update_span_impl( update_event = bt_safe_deep_copy(update_event) customizers = _get_span_customizers() - def parent_ids(): - exporter = _get_exporter() - return exporter( - object_type=parent_object_type, - object_id=parent_object_id.get(), - ).object_id_fields() - def compute_record() -> dict[str, Any] | None: record = dict( id=id, **update_event, - **parent_ids(), + **_object_id_fields(parent_object_type, parent_object_id.get()), **{ IS_MERGE_FIELD: True, }, ) return _customize_span_export(record, customizers) if customizers else record - _state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False)) + (public_bg_logger or _state.global_bg_logger()).log(LazyValue(compute_record, use_mutex=False)) def update_span(exported: str, **event: Any) -> None: @@ -4054,6 +4479,17 @@ def update_span(exported: str, **event: Any) -> None: event_without_span_ids.pop("span_id", None) event_without_span_ids.pop("root_span_id", None) + public_logger = _current_public_logger() + if public_logger is not None: + # Resolving the slug's object needs a private login, and ingestion key + # rows stay in the key's project anyway. + return public_logger.update_span( + id=components.row_id, + root_span_id=components.root_span_id, + span_id=components.span_id, + **event_without_span_ids, + ) + return _update_span_impl( parent_object_type=components.object_type, parent_object_id=LazyValue(_span_components_to_object_id_lambda(components), use_mutex=False), @@ -4247,6 +4683,7 @@ def _start_span_parent_args( parent_span_ids: ParentSpanIds | None, propagated_event: dict[str, Any] | None, propagated_state: "PropagatedState | None" = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ) -> dict[str, Any]: # `parent` may be an exported slug string or an opaque W3C trace-context dict. parent_slug, parent_propagated_state = _normalize_parent(parent) @@ -4257,16 +4694,22 @@ def _start_span_parent_args( f"Mismatch between expected span parent object type {parent_object_type} and provided type {parent_components.object_type}" ) - parent_components_object_id_lambda = _span_components_to_object_id_lambda(parent_components) + if public_bg_logger is not None: + # Resolving the parent's object needs a private login. Ingestion key + # spans stay in the key's project anyway, so the parent only + # contributes its span ids. + arg_parent_object_id = parent_object_id + else: + parent_components_object_id_lambda = _span_components_to_object_id_lambda(parent_components) - def compute_parent_object_id(): - parent_components_object_id = parent_components_object_id_lambda() - assert parent_object_id.get() == parent_components_object_id, ( - f"Mismatch between expected span parent object id {parent_object_id.get()} and provided id {parent_components_object_id}" - ) - return parent_object_id.get() + def compute_parent_object_id(): + parent_components_object_id = parent_components_object_id_lambda() + assert parent_object_id.get() == parent_components_object_id, ( + f"Mismatch between expected span parent object id {parent_object_id.get()} and provided id {parent_components_object_id}" + ) + return parent_object_id.get() - arg_parent_object_id = LazyValue(compute_parent_object_id, use_mutex=False) + arg_parent_object_id = LazyValue(compute_parent_object_id, use_mutex=False) if parent_components.row_id and _parent_span_ids_usable( parent_components.span_id, parent_components.root_span_id ): @@ -4292,6 +4735,7 @@ def compute_parent_object_id(): parent_span_ids=arg_parent_span_ids, propagated_event=arg_propagated_event, propagated_state=arg_propagated_state, + public_bg_logger=public_bg_logger, ) @@ -4777,6 +5221,7 @@ def __init__( state: BraintrustState | None = None, lookup_span_parent: bool = True, internal: SpanInternalOptions | None = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): if span_attributes is None: span_attributes = SpanAttributes() @@ -4797,6 +5242,7 @@ def __init__( self.parent_object_type = parent_object_type self.parent_object_id = parent_object_id self.parent_compute_object_metadata_args = parent_compute_object_metadata_args + self.public_bg_logger = public_bg_logger # Merge propagated_event into event. The propagated_event data will get # propagated-and-merged into every subspan. @@ -4967,14 +5413,10 @@ def write_span_cache(record: dict[str, Any]) -> None: write_span_cache(serializable_partial_record) def compute_record() -> dict[str, Any] | None: - exporter = _get_exporter() record = dict( **serializable_partial_record, **{k: v.get() for k, v in lazy_partial_record.items()}, - **exporter( - object_type=self.parent_object_type, - object_id=self.parent_object_id.get(), - ).object_id_fields(), + **_object_id_fields(self.parent_object_type, self.parent_object_id.get()), ) # Resolve and customize inside the cached LazyValue: every incremental # span record is transformed once, before the background @@ -4993,13 +5435,14 @@ def compute_record() -> dict[str, Any] | None: lazy_record = LazyValue(compute_record, use_mutex=pending_cache_key is not None) if pending_cache_key is not None: self.state.span_cache._track_pending_record(self.root_span_id, pending_cache_key, lazy_record) - self.state.global_bg_logger().log(lazy_record) + (self.public_bg_logger or self.state.global_bg_logger()).log(lazy_record) def log_feedback(self, **event: Any) -> None: return _log_feedback_impl( parent_object_type=self.parent_object_type, parent_object_id=self.parent_object_id, id=self.id, + public_bg_logger=self.public_bg_logger, **event, ) @@ -5034,6 +5477,7 @@ def start_span( parent_span_ids=parent_span_ids, propagated_event=coalesce(propagated_event, self.propagated_event), propagated_state=self.propagated_state, + public_bg_logger=self.public_bg_logger, ), name=name, type=type, @@ -5059,10 +5503,10 @@ def end(self, end_time: float | None = None) -> float: def export(self) -> str: if self.parent_compute_object_metadata_args and not self.parent_object_id.has_succeeded: object_id = None - compute_object_metadata_args = self.parent_compute_object_metadata_args else: object_id = self.parent_object_id.get() - compute_object_metadata_args = None + # Ingestion key loggers may not know their project id. + compute_object_metadata_args = None if object_id else self.parent_compute_object_metadata_args # Choose SpanComponents version based on the active ID format (hex -> V4, # legacy UUID -> V3). Coupled via _get_exporter() so the two never desync. @@ -5142,6 +5586,9 @@ def link(self) -> str: return NOOP_SPAN_PERMALINK def permalink(self) -> str: + if self.public_bg_logger is not None: + # Permalinks resolve the org through a private login. + return self.link() try: return permalink(self.export()) except Exception as e: @@ -5156,7 +5603,7 @@ def close(self, end_time=None) -> float: def flush(self) -> None: """Flush any pending rows to the server.""" - self.state.global_bg_logger().flush() + (self.public_bg_logger or self.state.global_bg_logger()).flush() def set_current(self): if self.can_set_current: @@ -5942,8 +6389,10 @@ def __init__( compute_metadata_args: dict | None = None, link_args: dict | None = None, state: BraintrustState | None = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): self._lazy_metadata = lazy_metadata + self._public_bg_logger = public_bg_logger self.async_flush = async_flush self._compute_metadata_args = compute_metadata_args self.last_start_time = time.time() @@ -6186,6 +6635,7 @@ def log_feedback( comment=comment, metadata=metadata, source=source, + public_bg_logger=self._public_bg_logger, ) def start_span( @@ -6237,6 +6687,7 @@ def update_span(self, id: str, **event: Any) -> None: id=id, root_span_id=root_span_id, span_id=span_id, + public_bg_logger=self._public_bg_logger, **event, ) @@ -6262,6 +6713,7 @@ def _start_span_impl( parent_compute_object_metadata_args=self._compute_metadata_args, parent_span_ids=None, propagated_event=propagated_event, + public_bg_logger=self._public_bg_logger, ) return SpanImpl( **parent_args, @@ -6287,10 +6739,10 @@ def export(self) -> str: # this `_lazy_id` object specifically will also be marked as computed. if self._compute_metadata_args and not self._lazy_id.has_succeeded: object_id = None - compute_object_metadata_args = self._compute_metadata_args else: object_id = self._lazy_id.get() - compute_object_metadata_args = None + # Ingestion key loggers may not know their project id. + compute_object_metadata_args = None if object_id else self._compute_metadata_args exporter = _get_exporter() return exporter( @@ -6322,7 +6774,7 @@ def flush(self) -> None: """ Flush any pending logs to the server. """ - self.state.global_bg_logger().flush() + (self._public_bg_logger or self.state.global_bg_logger()).flush() @dataclasses.dataclass diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py new file mode 100644 index 000000000..7e49fd0ee --- /dev/null +++ b/py/src/braintrust/test_ingestion_key.py @@ -0,0 +1,885 @@ +"""Tests for loggers that publish with a public ingestion key. + +The fake data plane below implements the ingestion key contract over real HTTP, +so these tests exercise the actual logger, batcher, and upload flow. +""" + +import hashlib +import http.server +import json +import logging +import re +import threading +import time +import uuid +from contextlib import contextmanager + +import braintrust +import pytest +from braintrust import logger +from braintrust.logger import Attachment, ExternalAttachment, _parse_ingestion_url +from braintrust.test_helpers import memory_logger, simulate_login # noqa: F401 + + +KEY = "bt-ik-" + "a" * 48 +OTHER_KEY = "bt-ik-" + "b" * 48 +BASE_PATH = "/deployment/base/ingest" +UPLOAD_PATH = re.compile(rf"^{BASE_PATH}/v1/uploads/([0-9a-f-]+)/(chunks/(\d+)|complete)$") + + +class FakeDataPlane: + """Records requests and answers them like a data plane serving one ingestion key.""" + + def __init__(self, key: str = KEY, chunk_bytes: int = 4): + self.key = key + self.chunk_bytes = chunk_bytes + self.requests: list[dict] = [] + self.rows: list[dict] = [] + self.uploads: dict[str, dict] = {} + # Responses to return before handling requests normally, keyed by (method, route). + self.scripted: dict[tuple[str, str], list[tuple[int, dict, bytes]]] = {} + self.grant_lifetimes_ms: list[int] = [] + # Fields to override in the next grants, to simulate malformed responses. + self.grant_overrides: list[dict] = [] + # Delay before answering upload creation, to expire short grants regardless of clock resolution. + self.create_delay_s = 0.0 + self.project_id: str | None = None + self.lock = threading.Lock() + + def script(self, method: str, route: str, *responses: tuple[int, dict, bytes]) -> None: + self.scripted.setdefault((method, route), []).extend(responses) + + def paths(self) -> list[tuple[str, str]]: + return [(r["method"], r["path"]) for r in self.requests] + + def handle(self, method: str, path: str, headers, body: bytes) -> tuple[int, dict, bytes]: + with self.lock: + self.requests.append({"method": method, "path": path, "headers": dict(headers), "body": body}) + route = self._route(path) + scripted = self.scripted.get((method, route)) + if scripted: + return scripted.pop(0) + if headers.get("Authorization") != f"Bearer {self.key}": + return 401, {}, b'{"error": "invalid ingestion key"}' + if route == "logs" and method == "POST": + return self._logs(json.loads(body)) + if route == "create" and method == "POST": + return self._create(json.loads(body)) + if route == "chunk" and method == "PUT": + return self._chunk(path, headers, body) + if route == "complete" and method == "POST": + return self._complete(path) + return 404, {}, b'{"error": "not found"}' + + @staticmethod + def _route(path: str) -> str: + if path == f"{BASE_PATH}/v1/logs": + return "logs" + if path == f"{BASE_PATH}/v1/uploads": + return "create" + match = UPLOAD_PATH.match(path) + if match: + return "chunk" if match.group(3) is not None else "complete" + return "unknown" + + def _logs(self, body: dict) -> tuple[int, dict, bytes]: + assert body["api_version"] == 2 + rows = body["rows"] + if isinstance(rows, dict): + assert rows == {"type": "logs3_overflow", "key": rows["key"]} + upload = self.uploads[rows["key"]] + assert upload["purpose"] == "logs3_overflow" and upload["completed"] + rows = json.loads(upload["data"])["rows"] + for row in rows: + if self.project_id is not None and row.get("project_id", self.project_id) != self.project_id: + return 403, {}, b'{"error": "row project does not match the ingestion key"}' + # External references can point at any object store content, so ingestion keys can't write them. + if '"external_attachment"' in json.dumps(row): + return 400, {}, b'{"error": "external_attachment references are not allowed"}' + self.rows.extend(rows) + return 200, {}, json.dumps({"ids": [row["id"] for row in rows]}).encode() + + def _create(self, body: dict) -> tuple[int, dict, bytes]: + time.sleep(self.create_delay_s) + self.create_delay_s = 0.0 + upload_id = str(uuid.uuid4()) + num_chunks = -(-body["size_bytes"] // self.chunk_bytes) + self.uploads[upload_id] = {**body, "chunks": {}, "completed": False} + lifetime = self.grant_lifetimes_ms.pop(0) if self.grant_lifetimes_ms else 300000 + response = { + "upload_id": upload_id, + "chunk_bytes": self.chunk_bytes, + "num_chunks": num_chunks, + "expires_in_ms": lifetime, + **(self.grant_overrides.pop(0) if self.grant_overrides else {}), + } + return 201, {}, json.dumps(response).encode() + + def _chunk(self, path: str, headers, body: bytes) -> tuple[int, dict, bytes]: + match = UPLOAD_PATH.match(path) + assert match is not None + upload = self.uploads[match.group(1)] + index = int(match.group(3)) + assert headers.get("Content-Type") == "application/octet-stream" + assert "Content-Encoding" not in headers + upload["chunks"][index] = body + return 200, {}, json.dumps({"index": index, "size_bytes": len(body)}).encode() + + def _complete(self, path: str) -> tuple[int, dict, bytes]: + match = UPLOAD_PATH.match(path) + assert match is not None + upload_id = match.group(1) + upload = self.uploads[upload_id] + data = b"".join(upload["chunks"][i] for i in sorted(upload["chunks"])) + assert len(data) == upload["size_bytes"] + assert hashlib.sha256(data).hexdigest() == upload["sha256"] + upload["data"] = data + upload["completed"] = True + if upload["purpose"] == "attachment": + reference = { + "type": "braintrust_attachment", + "key": upload_id, + "filename": upload["filename"], + "content_type": upload["content_type"], + } + else: + reference = {"type": "logs3_overflow", "key": upload_id} + response = {"reference": reference, "size_bytes": len(data), "sha256": upload["sha256"]} + return 200, {}, json.dumps(response).encode() + + +@contextmanager +def serve(data_plane: FakeDataPlane): + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, format, *args): + pass + + def _handle(self): + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + status, headers, response_body = data_plane.handle(self.command, self.path, self.headers, body) + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response_body))) + self.end_headers() + self.wfile.write(response_body) + + do_GET = do_POST = do_PUT = do_PATCH = do_DELETE = _handle + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}{BASE_PATH}" + finally: + server.shutdown() + server.server_close() + + +def ingestion_url(root: str, key: str = KEY) -> str: + return f"{root}?ingestKey={key}" + + +@pytest.fixture(autouse=True) +def deterministic_flush(monkeypatch): + # Rows only publish on explicit flushes, failures raise, and retries don't sleep. + monkeypatch.setenv("BRAINTRUST_SYNC_FLUSH", "true") + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "true") + monkeypatch.delenv("BRAINTRUST_INGESTION_KEY", raising=False) + monkeypatch.setattr(logger, "BACKGROUND_LOGGER_BASE_SLEEP_TIME_S", 0) + logger._state.reset_parent_state() + yield + logger._state.reset_parent_state() + + +@pytest.fixture +def data_plane(): + plane = FakeDataPlane() + with serve(plane) as root: + plane.root = root + yield plane + + +@pytest.fixture +def no_private_login(monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("ingestion key loggers must not log in") + + monkeypatch.setattr(logger, "login_to_state", fail) + monkeypatch.setattr(logger, "_compute_logger_metadata", fail) + + +def flush_error(target) -> str: + """Flush and return the error, which publishing raises inside an exception group.""" + with pytest.raises(BaseException) as exc_info: + target.flush() + return "; ".join(str(e) for e in exc_info.value.exceptions) + + +def assert_no_key_in_urls(plane: FakeDataPlane) -> None: + for request in plane.requests: + assert "ingestKey" not in request["path"] + assert plane.key not in request["path"] + + +@pytest.mark.parametrize( + "url,expected", + [ + (f"https://dp.example/ingest?ingestKey={KEY}", "https://dp.example/ingest"), + (f"https://dp.example/deployment/base/ingest/?ingestKey={KEY}", "https://dp.example/deployment/base/ingest"), + (f"http://localhost:8000/ingest?ingestKey={KEY}\n", "http://localhost:8000/ingest"), + (f"https://dp.example/base/ingest///?ingestKey={KEY}", "https://dp.example/base/ingest"), + (f"https://dp.example/ingest?&ingestKey={KEY}&&", "https://dp.example/ingest"), + (f"https://dp.example/ingest?ingest%4Bey=bt%2Dik%2D{'a' * 48}", "https://dp.example/ingest"), + ], +) +def test_parse_ingestion_url(url, expected): + endpoint = _parse_ingestion_url(url) + + assert endpoint.url == expected + assert endpoint.key == KEY + assert KEY not in repr(endpoint) + + +@pytest.mark.parametrize( + "url", + [ + KEY, + f"ftp://dp.example/ingest?ingestKey={KEY}", + f"https://user:pass@dp.example/ingest?ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}#fragment", + f"https://dp.example/ingest?ingestKey={KEY}#", + f"https://dp.example/ingest#?ingestKey={KEY}", + "https://dp.example/ingest?ingestKey", + f"https://dp.example/ingest?ingestKey={KEY}&ingestKey=", + f"https://dp.example/ingest?ingestKey={KEY}=x", + f"https://dp.example/ingest/v1/logs?ingestKey={KEY}", + f"https://dp.example/not-ingest?ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}&ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}&project=foo", + f"https://dp.example/ingest?key={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}x", + "https://dp.example/ingest?ingestKey=sk-" + "a" * 48, + f"https://dp.example:notaport/ingest?ingestKey={KEY}", + ], +) +def test_parse_ingestion_url_rejects_invalid_urls_without_leaking_the_key(url): + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL") as exc_info: + _parse_ingestion_url(url) + + assert "a" * 48 not in str(exc_info.value) + assert exc_info.value.__cause__ is None and exc_info.value.__context__ is None + + +def test_first_flush_only_posts_rows_with_the_key(data_plane, monkeypatch, no_private_login): + # Ambient private credentials must be ignored. + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + + public_logger = braintrust.init_logger(project="my-project", ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="question", output="answer", scores={"good": 1}, metadata={"a": 1}) + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + request = data_plane.requests[0] + assert request["headers"]["Authorization"] == f"Bearer {KEY}" + assert request["headers"]["Content-Type"] == "application/json" + [row] = data_plane.rows + assert row["id"] == row_id + assert row["log_id"] == "g" + assert "project_id" not in row + assert { + "input": "question", + "output": "answer", + "scores": {"good": 1}, + "metadata": {"a": 1}, + }.items() <= row.items() + assert set(row) <= set(logger._INGESTION_ROW_FIELDS) + assert public_logger.id is None + assert not logger._state.logged_in + assert_no_key_in_urls(data_plane) + + +def test_trace_context_round_trips_without_project_metadata(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + with public_logger.start_span(name="root") as root: + exported = root.export() + headers = root.inject() + assert root.permalink() + # A parent exported elsewhere only contributes its span ids. + with braintrust.start_span(name="from-slug", parent=exported) as from_slug: + pass + with braintrust.start_span(name="from-headers", parent=braintrust.extract_trace_context(headers)) as from_headers: + pass + public_logger.flush() + + rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + assert rows["from-slug"]["root_span_id"] == root.root_span_id + assert rows["from-slug"]["span_parents"] == [root.span_id] + assert rows["from-headers"]["root_span_id"] == root.root_span_id + assert rows["from-headers"]["span_parents"] == [root.span_id] + assert all("project_id" not in row for row in data_plane.rows) + assert from_slug.public_bg_logger is public_logger._public_bg_logger + assert from_headers.public_bg_logger is public_logger._public_bg_logger + + +def test_explicit_project_id_is_sent_and_mismatches_are_not_retried(data_plane, no_private_login): + data_plane.project_id = "key-project" + public_logger = braintrust.init_logger(project_id="other-project", ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + + assert "403" in flush_error(public_logger) + + [request] = data_plane.requests + [row] = json.loads(request["body"])["rows"] + assert row["project_id"] == "other-project" + assert data_plane.rows == [] + + +def test_retries_transient_failures_with_the_same_key(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "2") + data_plane.script( + "POST", + "logs", + (429, {"Retry-After": "0"}, b'{"error": "rate limited"}'), + (503, {}, b'{"error": "unavailable"}'), + ) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] * 3 + assert {r["headers"]["Authorization"] for r in data_plane.requests} == {f"Bearer {KEY}"} + assert len(data_plane.rows) == 1 + + +def test_rejected_writes_are_not_retried_or_sent_elsewhere(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "2") + data_plane.script("POST", "logs", (401, {}, f'{{"error": "bad key {KEY}"}}'.encode())) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + + error = flush_error(public_logger) + + assert "401" in error + assert KEY not in error + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + + +def test_attachments_upload_in_advertised_chunks_before_rows(data_plane, no_private_login): + data_plane.chunk_bytes = 4 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + attachment = Attachment(data=b"hello world", filename="greeting.txt", content_type="text/plain") + empty = Attachment(data=b"", filename="empty.txt", content_type="text/plain") + # A reference committed earlier is plain data, so it passes through unchanged. + committed = { + "type": "braintrust_attachment", + "key": str(uuid.uuid4()), + "filename": "old.txt", + "content_type": "text/plain", + } + public_logger.log(input={"file": attachment, "empty": empty, "committed": committed, "again": attachment}) + public_logger.flush() + + [attachment_id, empty_id] = list(data_plane.uploads) + assert data_plane.paths() == [ + ("POST", f"{BASE_PATH}/v1/uploads"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/0"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/1"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/2"), + ("POST", f"{BASE_PATH}/v1/uploads/{attachment_id}/complete"), + ("POST", f"{BASE_PATH}/v1/uploads"), + ("POST", f"{BASE_PATH}/v1/uploads/{empty_id}/complete"), + ("POST", f"{BASE_PATH}/v1/logs"), + ] + assert [len(r["body"]) for r in data_plane.requests if r["method"] == "PUT"] == [4, 4, 3] + assert json.loads(data_plane.requests[0]["body"]) == { + "purpose": "attachment", + "size_bytes": 11, + "content_type": "text/plain", + "sha256": hashlib.sha256(b"hello world").hexdigest(), + "filename": "greeting.txt", + } + assert json.loads(data_plane.requests[4]["body"]) == {} + assert data_plane.uploads[empty_id]["size_bytes"] == 0 + [row] = data_plane.rows + file_reference = { + "type": "braintrust_attachment", + "key": attachment_id, + "filename": "greeting.txt", + "content_type": "text/plain", + } + assert row["input"] == { + "file": file_reference, + "empty": { + "type": "braintrust_attachment", + "key": empty_id, + "filename": "empty.txt", + "content_type": "text/plain", + }, + "committed": committed, + "again": file_reference, + } + + +def test_external_attachments_are_rejected_before_publishing(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + external = ExternalAttachment(url="s3://private-bucket/file.txt", filename="file.txt", content_type="text/plain") + public_logger.log(input={"file": external}) + kept_id = public_logger.log(input="no attachment") + + with pytest.raises(ValueError, match="ExternalAttachment is not supported with ingestion keys"): + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + assert [row["id"] for row in data_plane.rows] == [kept_id] + assert b"private-bucket" not in data_plane.requests[0]["body"] + + +def test_raw_external_attachment_references_are_not_a_silent_success(data_plane, no_private_login): + # Plain dicts aren't validated client side, so the data plane's rejection must surface. + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input={"type": "external_attachment", "url": "s3://private-bucket/file.txt"}) + + assert "400" in flush_error(public_logger) + assert data_plane.rows == [] + + +def test_rows_with_failed_attachments_are_dropped(data_plane, no_private_login): + data_plane.script("POST", "create", (400, {}, b'{"error": "uploads disabled"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abc", filename="a.txt", content_type="text/plain")) + kept_id = public_logger.log(input="no attachment") + + with pytest.raises(RuntimeError, match="400"): + public_logger.flush() + + assert [row["id"] for row in data_plane.rows] == [kept_id] + assert not any("attachment" in path for _, path in data_plane.paths()) + + +def test_expired_upload_grant_starts_a_new_upload(data_plane, no_private_login): + # The grant expires while its creation response is in flight. + data_plane.grant_lifetimes_ms = [50] + data_plane.create_delay_s = 0.2 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + [expired_id, upload_id] = list(data_plane.uploads) + # Nothing is sent against the expired grant, so its id never shows up in a path. + assert data_plane.paths() == [ + ("POST", f"{BASE_PATH}/v1/uploads"), + ("POST", f"{BASE_PATH}/v1/uploads"), + ("PUT", f"{BASE_PATH}/v1/uploads/{upload_id}/chunks/0"), + ("PUT", f"{BASE_PATH}/v1/uploads/{upload_id}/chunks/1"), + ("POST", f"{BASE_PATH}/v1/uploads/{upload_id}/complete"), + ("POST", f"{BASE_PATH}/v1/logs"), + ] + assert not data_plane.uploads[expired_id]["chunks"] + assert data_plane.uploads[upload_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_oversized_batches_overflow_through_uploads(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_MAX_REQUEST_SIZE", "2000") + data_plane.chunk_bytes = 1024 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x" * 5000) + public_logger.flush() + + [upload_id] = list(data_plane.uploads) + upload = data_plane.uploads[upload_id] + assert upload["purpose"] == "logs3_overflow" + assert upload["content_type"] == "application/json" + overflow_payload = json.loads(upload["data"]) + assert set(overflow_payload) == {"api_version", "rows"} + assert overflow_payload["api_version"] == 2 + assert [row["id"] for row in overflow_payload["rows"]] == [row_id] + [_, *chunk_requests, _, logs_request] = data_plane.requests + assert [len(r["body"]) for r in chunk_requests] == [1024] * (len(chunk_requests) - 1) + [ + len(upload["data"]) - 1024 * (len(chunk_requests) - 1) + ] + assert json.loads(logs_request["body"]) == { + "api_version": 2, + "rows": {"type": "logs3_overflow", "key": upload_id}, + } + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_feedback_and_disallowed_fields(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + with public_logger.start_span(name="root") as span: + span.log(input="x", experiment_id="someone-elses-experiment", dataset_id="d") + public_logger.log_feedback(id=span.id, scores={"thumbs_up": 1}, tags=["a"]) + braintrust.update_span(span.export(), output="updated") + with pytest.raises(ValueError, match="ingestion key"): + public_logger.log_feedback(id=span.id, comment="nice") + with pytest.raises(ValueError, match="ingestion key"): + span.log_feedback(scores={"thumbs_up": 0}, metadata={"user": "u"}) + public_logger.flush() + + [row] = data_plane.rows + assert row["id"] == span.id + assert row["scores"] == {"thumbs_up": 1} + assert row["output"] == "updated" + assert set(row) <= set(logger._INGESTION_ROW_FIELDS) + + +def test_masking_and_global_flush_cover_public_loggers(data_plane, no_private_login): + braintrust.set_masking_function(lambda value: "masked" if value == "secret" else value) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="secret") + braintrust.flush() + + assert data_plane.rows[0]["input"] == "masked" + + +def test_env_ingestion_key_takes_precedence_over_ambient_private_credentials( + data_plane, monkeypatch, no_private_login +): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + public_logger = braintrust.init_logger(project="p") + public_logger.log(input="x") + public_logger.flush() + + assert public_logger._public_bg_logger is not None + assert len(data_plane.rows) == 1 + + +def test_env_ingestion_key_ignores_previous_login(data_plane, monkeypatch, memory_logger): + simulate_login() + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + public_logger = braintrust.init_logger(project="p") + public_logger.log(input="x") + public_logger.flush() + + assert len(data_plane.rows) == 1 + assert memory_logger.pop() == [] + + +def test_explicit_ingestion_key_overrides_env(data_plane, monkeypatch, no_private_login): + other_plane = FakeDataPlane(key=OTHER_KEY) + with serve(other_plane) as other_root: + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(other_root, OTHER_KEY)) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + public_logger.flush() + + assert len(data_plane.rows) == 1 + assert other_plane.requests == [] + + +def test_explicit_api_key_overrides_env_ingestion_key(data_plane, monkeypatch, memory_logger): + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + private_logger = braintrust.init_logger(project="p", project_id="private-project", api_key=logger.TEST_API_KEY) + private_logger.log(input="x") + private_logger.flush() + + assert private_logger._public_bg_logger is None + assert [row["input"] for row in memory_logger.pop()] == ["x"] + assert data_plane.requests == [] + + +def test_explicit_api_key_and_ingestion_key_conflict(data_plane): + with pytest.raises(ValueError, match="not both"): + braintrust.init_logger(api_key="sk-private", ingestion_key=ingestion_url(data_plane.root)) + + +def test_public_and_private_loggers_do_not_mix(data_plane, memory_logger): + simulate_login() + other_plane = FakeDataPlane(key=OTHER_KEY) + with serve(other_plane) as other_root: + private_logger = braintrust.init_logger(project="private", project_id="private-project", set_current=False) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + other_logger = braintrust.init_logger(ingestion_key=ingestion_url(other_root, OTHER_KEY), set_current=False) + + def log_many(target, name): + for i in range(20): + target.log(input=f"{name}-{i}") + + threads = [ + threading.Thread(target=log_many, args=(target, name)) + for target, name in [(public_logger, "public"), (other_logger, "other")] + ] + for thread in threads: + thread.start() + # The memory logger override is thread local, so private rows log from this thread. + log_many(private_logger, "private") + for thread in threads: + thread.join() + braintrust.flush() + + assert sorted(row["input"] for row in data_plane.rows) == sorted(f"public-{i}" for i in range(20)) + assert sorted(row["input"] for row in other_plane.rows) == sorted(f"other-{i}" for i in range(20)) + assert {r["headers"]["Authorization"] for r in other_plane.requests} == {f"Bearer {OTHER_KEY}"} + assert sorted(row["input"] for row in memory_logger.pop()) == sorted(f"private-{i}" for i in range(20)) + + +def test_key_does_not_leak_through_repr(data_plane): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + assert KEY not in repr(public_logger) + assert KEY not in repr(public_logger._public_bg_logger.endpoint) + assert KEY not in public_logger.export() + + +def test_explicit_empty_or_malformed_ingestion_key_never_falls_back(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + for value in ["", " ", "https://dp.example/ingest", KEY]: + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL"): + braintrust.init_logger(ingestion_key=value) + + assert braintrust.current_logger() is None + assert data_plane.requests == [] + + +def test_malformed_env_ingestion_key_never_falls_back(monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", "https://dp.example/ingest?ingestKey=not-a-key") + + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL"): + braintrust.init_logger(project="p") + + +def test_default_overflow_threshold_is_512_kib(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_MAX_REQUEST_SIZE", raising=False) + data_plane.chunk_bytes = 512 * 1024 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + small_id = public_logger.log(input="x" * (200 * 1024)) + public_logger.flush() + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + assert len(data_plane.requests[0]["body"]) < 512 * 1024 + + large_id = public_logger.log(input="x" * (600 * 1024)) + public_logger.flush() + [upload_id] = list(data_plane.uploads) + assert [method for method, _ in data_plane.paths()[1:]] == ["POST", "PUT", "PUT", "POST", "POST"] + assert json.loads(data_plane.requests[-1]["body"]) == { + "api_version": 2, + "rows": {"type": "logs3_overflow", "key": upload_id}, + } + assert [row["id"] for row in data_plane.rows] == [small_id, large_id] + assert not any(path.endswith("/version") for _, path in data_plane.paths()) + + +@pytest.mark.parametrize( + "override", + [ + {"expires_in_ms": 0}, + {"expires_in_ms": -1}, + {"expires_in_ms": 300001}, + {"expires_in_ms": "300000"}, + {"expires_in_ms": True}, + {"expires_in_ms": None}, + {"num_chunks": 3}, + {"num_chunks": 1}, + {"chunk_bytes": 0}, + {"chunk_bytes": 2.5}, + {"upload_id": "../../v1/logs"}, + {"upload_id": None}, + ], +) +def test_malformed_upload_grants_are_rejected(data_plane, override, no_private_login): + # 6 bytes in 4 byte chunks is 2 chunks. + data_plane.grant_overrides = [override] + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + with pytest.raises(RuntimeError, match="Invalid upload grant") as exc_info: + public_logger.flush() + + assert KEY not in str(exc_info.value) + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/uploads")] + assert data_plane.rows == [] + + +def test_maximum_grant_lifetime_is_accepted(data_plane, no_private_login): + data_plane.grant_lifetimes_ms = [300000] + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + assert len(data_plane.rows) == 1 + + +def test_gone_upload_grant_starts_a_new_upload(data_plane, no_private_login): + data_plane.script("PUT", "chunk", (410, {}, b'{"error": "upload grant expired"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + [gone_id, upload_id] = list(data_plane.uploads) + assert not data_plane.uploads[gone_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_gone_upload_grants_are_retried_a_bounded_number_of_times(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "1") + data_plane.script("PUT", "chunk", *[(410, {}, b'{"error": "upload grant expired"}')] * 5) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + with pytest.raises(RuntimeError, match="kept expiring"): + public_logger.flush() + + assert len(data_plane.uploads) == 2 + assert data_plane.rows == [] + + +def test_retry_waits_never_outlive_the_upload_grant(data_plane, no_private_login): + data_plane.grant_lifetimes_ms = [2000] + data_plane.script("PUT", "chunk", (503, {"Retry-After": "30"}, b'{"error": "unavailable"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + started_at = time.monotonic() + public_logger.flush() + + assert time.monotonic() - started_at < 5 + [expired_id, upload_id] = list(data_plane.uploads) + assert not data_plane.uploads[expired_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_transport_failures_never_reveal_the_key(data_plane, monkeypatch, capsys, caplog, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "1") + # Reach the closed port directly, even if the environment configures a proxy. + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setenv("no_proxy", "127.0.0.1") + caplog.set_level(logging.DEBUG) + with serve(FakeDataPlane()) as closed_root: + pass + errors = [] + for root in [closed_root, data_plane.root]: + data_plane.script("POST", "logs", *[(500, {}, f'{{"error": "bad {KEY}"}}'.encode())] * 2) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(root)) + public_logger.log(input="x") + errors.append(flush_error(public_logger)) + + # The closed port fails without an HTTP response, and is retried like any transport error. + assert "POST /v1/logs failed after 2 attempts with " in errors[0] + assert re.search(r"failed after 2 attempts with \d{3}: ", errors[0]) is None + assert "POST /v1/logs failed after 2 attempts with 500: " in errors[1] + output = capsys.readouterr() + for text in [*errors, output.out, output.err, caplog.text]: + assert KEY not in text + assert "ingestKey" not in text + + +def test_redirects_are_not_followed(data_plane, no_private_login): + other_plane = FakeDataPlane() + with serve(other_plane) as other_root: + data_plane.script("POST", "logs", (307, {"Location": f"{other_root}/v1/logs"}, b"")) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + error = flush_error(public_logger) + + assert "307" in error + assert other_plane.requests == [] + + +def test_background_thread_publishes_without_explicit_flush(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_SYNC_FLUSH") + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x") + + deadline = time.monotonic() + 10 + while not data_plane.rows and time.monotonic() < deadline: + time.sleep(0.01) + + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_atexit_flush_publishes_queued_rows(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH") + registered = [] + monkeypatch.setattr(logger.atexit, "register", registered.append) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x") + + bg_logger = public_logger._public_bg_logger + assert registered == [bg_logger._finalize] + bg_logger._finalize() + + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_public_and_private_spans_nest_without_crossing_loggers(data_plane, memory_logger): + simulate_login() + private_logger = braintrust.init_logger(project="private", project_id="private-project", set_current=False) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + + with private_logger.start_span(name="private-root") as private_root: + with public_logger.start_span(name="public-child") as public_child: + # Global helpers follow the current span. + with braintrust.start_span(name="public-grandchild"): + pass + with private_logger.start_span(name="private-grandchild"): + pass + braintrust.flush() + + public_rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + private_rows = {row["span_attributes"]["name"]: row for row in memory_logger.pop()} + assert set(public_rows) == {"public-child", "public-grandchild"} + assert set(private_rows) == {"private-root", "private-grandchild"} + assert public_rows["public-child"]["span_parents"] == [private_root.span_id] + assert public_rows["public-grandchild"]["span_parents"] == [public_child.span_id] + assert private_rows["private-grandchild"]["span_parents"] == [public_child.span_id] + assert {row["root_span_id"] for row in [*public_rows.values(), *private_rows.values()]} == { + private_root.root_span_id + } + assert all(row["project_id"] == "private-project" for row in private_rows.values()) + assert all("project_id" not in row for row in public_rows.values()) + + +def test_public_trace_context_continues_in_a_private_service(data_plane, memory_logger): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + with public_logger.start_span(name="client") as client_span: + headers = client_span.inject() + public_logger.flush() + + simulate_login() + # Resolving the id up front keeps the private logger from looking up its project by name. + braintrust.init_logger(project="server", project_id="server-project")._lazy_id.get() + with braintrust.start_span(name="handler", parent=braintrust.extract_trace_context(headers)): + pass + braintrust.flush() + + [handler] = memory_logger.pop() + assert handler["project_id"] == "server-project" + assert handler["root_span_id"] == client_span.root_span_id + assert handler["span_parents"] == [client_span.span_id] + assert [row["span_attributes"]["name"] for row in data_plane.rows] == ["client"] + + +def test_traced_functions_under_a_public_logger_stay_public(data_plane, memory_logger): + simulate_login() + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + span_ids = [] + + @braintrust.traced + def answer(question): + span_ids.append(braintrust.current_span().id) + with braintrust.start_span(name="lookup") as lookup: + lookup.log(metadata={"q": question}) + return "42" + + answer("why") + public_logger.log_feedback(id=span_ids[0], scores={"ok": 1}) + braintrust.flush() + + rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + assert set(rows) == {"answer", "lookup"} + assert rows["answer"]["output"] == "42" + assert rows["answer"]["scores"] == {"ok": 1} + assert memory_logger.pop() == []