diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index e99d5f23a4c4..a598124f6121 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -1,18 +1,29 @@ +import json import os -from typing import Any, Dict, List, Optional +import re +from typing import Any, Dict, List, Optional, Tuple, cast from pydantic import BaseModel, Field import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, + get_content_from_model_response, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.llms.openai import AllMessageValues + +GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai" +# Cap the in-memory buffer so persistent flush failures (e.g. Galileo +# unavailable, invalid credentials) cannot leak memory unboundedly. +GALILEO_MAX_IN_MEMORY_RECORDS = 1000 -# from here: https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#structuring-your-records class LLMResponse(BaseModel): latency_ms: int status_code: int @@ -37,65 +48,190 @@ class GalileoObserve(CustomLogger): def __init__(self) -> None: self.in_memory_records: List[dict] = [] self.batch_size = 1 - self.base_url = os.getenv("GALILEO_BASE_URL", None) - self.project_id = os.getenv("GALILEO_PROJECT_ID", None) + self.api_key = os.getenv("GALILEO_API_KEY") + self.project_id = os.getenv("GALILEO_PROJECT_ID") + self.log_stream_id = os.getenv("GALILEO_LOG_STREAM_ID") + self.username = os.getenv("GALILEO_USERNAME") + self.password = os.getenv("GALILEO_PASSWORD") + self.base_url = self._normalize_base_url(os.getenv("GALILEO_BASE_URL")) + if self.api_key and not self.base_url: + self.base_url = GALILEO_CLOUD_API_BASE_URL + self.use_v2_api = bool(self.api_key) self.headers: Optional[Dict[str, str]] = None self.async_httpx_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) - pass - def set_galileo_headers(self): - # following https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#logging-your-records + @staticmethod + def _normalize_base_url(base_url: Optional[str]) -> Optional[str]: + if base_url: + return base_url.rstrip("/") + return None - headers = { - "accept": "application/json", - "Content-Type": "application/x-www-form-urlencoded", - } - galileo_login_response = litellm.module_level_client.post( + def _is_configured(self) -> bool: + if not self.project_id or not self.base_url: + return False + if self.use_v2_api: + return bool(self.api_key) + return bool(self.username and self.password) + + async def async_set_galileo_headers(self) -> None: + galileo_login_response = await self.async_httpx_handler.post( url=f"{self.base_url}/login", - headers=headers, + headers={ + "accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, data={ - "username": os.getenv("GALILEO_USERNAME"), - "password": os.getenv("GALILEO_PASSWORD"), + "username": self.username, + "password": self.password, }, ) - + galileo_login_response.raise_for_status() access_token = galileo_login_response.json()["access_token"] - self.headers = { "accept": "application/json", "Content-Type": "application/json", "Authorization": f"Bearer {access_token}", } - def get_output_str_from_response(self, response_obj, kwargs): - output = None - if response_obj is not None and ( - kwargs.get("call_type", None) == "embedding" - or isinstance(response_obj, litellm.EmbeddingResponse) - ): - output = None - elif response_obj is not None and isinstance( - response_obj, litellm.ModelResponse - ): - output = response_obj["choices"][0]["message"].json() - elif response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): - output = response_obj.choices[0].text - elif response_obj is not None and isinstance( - response_obj, litellm.ImageResponse - ): - output = response_obj["data"] + async def _ensure_headers(self) -> bool: + if self.headers is not None: + return True + + if self.use_v2_api: + if not self.api_key: + return False + self.headers = { + "accept": "application/json", + "Content-Type": "application/json", + "Galileo-API-Key": self.api_key, + } + return True + + if not (self.username and self.password and self.base_url): + return False + + try: + await self.async_set_galileo_headers() + return True + except Exception as e: + verbose_logger.debug("Galileo Logger: failed to authenticate: %s", e) + return False + + @staticmethod + def _galileo_input_messages( + messages: Optional[List[Any]], input_text: str + ) -> List[Dict[str, str]]: + if not messages: + return [{"role": "user", "content": input_text}] + + galileo_messages: List[Dict[str, str]] = [] + for message in messages: + if not isinstance(message, dict): + continue + role = message.get("role") + if not role: + continue + galileo_messages.append( + { + "role": str(role), + "content": convert_content_list_to_str( + message=cast(AllMessageValues, message) + ), + } + ) + + if galileo_messages: + return galileo_messages + return [{"role": "user", "content": input_text}] + + @staticmethod + def _record_to_v2_span(record: Dict[str, Any]) -> Dict[str, Any]: + created_at = record.get("created_at", "") + if created_at and not re.search(r"(Z|[+-]\d{2}:?\d{2})$", created_at): + created_at = f"{created_at}Z" + + span: Dict[str, Any] = { + "type": "llm", + "name": record.get("node_type", "litellm"), + "created_at": created_at, + "input": GalileoObserve._galileo_input_messages( + record.get("messages"), record.get("input_text", "") + ), + "output": { + "role": "assistant", + "content": record.get("output_text", ""), + }, + "status_code": record.get("status_code", 200), + "model": record.get("model"), + "metrics": { + "duration_ns": int(record.get("latency_ms", 0)) * 1_000_000, + "num_input_tokens": record.get("num_input_tokens"), + "num_output_tokens": record.get("num_output_tokens"), + }, + } + if record.get("tags"): + span["tags"] = record["tags"] + return span + + def _get_ingest_request(self) -> Optional[Tuple[str, Dict[str, Any]]]: + if not self.base_url or not self.project_id: + return None + + # Snapshot the records to be sent into a new list so concurrent appends + # during the network round-trip (across the await points in + # flush_in_memory_records) aren't silently dropped when we later clear + # the in-memory buffer. + records = list(self.in_memory_records) + + if self.use_v2_api: + payload: Dict[str, Any] = { + "spans": [self._record_to_v2_span(record) for record in records], + "reliable": False, + } + if self.log_stream_id: + payload["log_stream_id"] = self.log_stream_id + return ( + f"{self.base_url}/v2/projects/{self.project_id}/spans", + payload, + ) + + return ( + f"{self.base_url}/projects/{self.project_id}/observe/ingest", + {"records": records}, + ) - return output + def get_output_str_from_response( + self, response_obj: Any, kwargs: Dict[str, Any] + ) -> Optional[str]: + if response_obj is None: + return None + if kwargs.get("call_type", None) == "embedding" or isinstance( + response_obj, litellm.EmbeddingResponse + ): + return None + if isinstance(response_obj, litellm.TextCompletionResponse): + return response_obj.choices[0].text + if isinstance(response_obj, litellm.ImageResponse): + return json.dumps(response_obj["data"], default=str) + if isinstance(response_obj, (litellm.ModelResponse, dict)): + return get_content_from_model_response(response_obj) + return None async def async_log_success_event( self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any ): verbose_logger.debug("On Async Success") + if not self._is_configured(): + verbose_logger.debug( + "Galileo Logger: skipping flush — set GALILEO_PROJECT_ID and " + "either GALILEO_API_KEY (hosted) or GALILEO_USERNAME/GALILEO_PASSWORD " + "(enterprise Observe)." + ) + return + _latency_ms = int((end_time - start_time).total_seconds() * 1000) _call_type = kwargs.get("call_type", "litellm") input_text = litellm.utils.get_formatted_prompt( @@ -125,26 +261,69 @@ async def async_log_success_event( ), # timestamp str constructed in "%Y-%m-%dT%H:%M:%S" format ) - # dump to dict request_dict = request_record.model_dump() + messages = kwargs.get("messages") + if messages: + request_dict["messages"] = messages self.in_memory_records.append(request_dict) + # Bound the buffer so persistent flush failures cannot grow it + # without limit. Drop the oldest records once we exceed the cap. + if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS: + dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS + self.in_memory_records = self.in_memory_records[ + -GALILEO_MAX_IN_MEMORY_RECORDS: + ] + verbose_logger.warning( + "Galileo Logger: in-memory buffer exceeded %s records; " + "dropped %s oldest record(s). Check Galileo connectivity/credentials.", + GALILEO_MAX_IN_MEMORY_RECORDS, + dropped, + ) + if len(self.in_memory_records) >= self.batch_size: await self.flush_in_memory_records() async def flush_in_memory_records(self): - verbose_logger.debug("flushing in memory records") - response = await self.async_httpx_handler.post( - url=f"{self.base_url}/projects/{self.project_id}/observe/ingest", - headers=self.headers, - json={"records": self.in_memory_records}, - ) + if not self.in_memory_records: + return + + # Capture the number of records that will be sent BEFORE any await so + # that concurrent appends made by other asyncio tasks during the + # network round-trip aren't silently dropped on the success-clear. + records_in_payload = len(self.in_memory_records) + + ingest_request = self._get_ingest_request() + if ingest_request is None: + verbose_logger.debug( + "Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID" + ) + return + + if not await self._ensure_headers(): + verbose_logger.debug("Galileo Logger: could not set request headers") + return + + url, payload = ingest_request + verbose_logger.debug("flushing in memory records to %s", url) + + try: + response = await self.async_httpx_handler.post( + url=url, + headers=self.headers, + json=payload, + ) + except Exception as e: + verbose_logger.debug( + "Galileo Logger: failed to flush in memory records: %s", e + ) + return - if response.status_code == 200: + if response.is_success: verbose_logger.debug( - "Galileo Logger:successfully flushed in memory records" + "Galileo Logger: successfully flushed in memory records" ) - self.in_memory_records = [] + del self.in_memory_records[:records_in_payload] else: verbose_logger.debug("Galileo Logger: failed to flush in memory records") verbose_logger.debug( @@ -152,6 +331,13 @@ async def flush_in_memory_records(self): response.text, response.status_code, ) + # Legacy enterprise auth caches a bearer token obtained from + # /login. If the request was rejected for auth reasons, drop the + # cached headers so the next flush re-authenticates instead of + # silently failing forever on a stale token. The v2 API key path + # uses a long-lived static key, so leave its headers in place. + if not self.use_v2_api and response.status_code in (401, 403): + self.headers = None async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("On Async Failure") diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 549d9e469f9a..888ccc781884 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -568,14 +568,22 @@ async def wrapper(*args, **kwargs): ) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: - await _emit_management_endpoint_otel_span( - func=func, - kwargs=kwargs, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - exception=e, - ) + try: + await _emit_management_endpoint_otel_span( + func=func, + kwargs=kwargs, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + exception=e, + ) + except Exception as otel_exc: + # Non-Blocking Exception - never let OTEL failures swallow + # the original management-endpoint exception. + verbose_logger.debug( + "Error emitting OTEL span in management endpoint wrapper failure path: %s", + str(otel_exc), + ) raise e diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/test_litellm/integrations/test_galileo.py new file mode 100644 index 000000000000..aab220f46be4 --- /dev/null +++ b/tests/test_litellm/integrations/test_galileo.py @@ -0,0 +1,397 @@ +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.integrations.galileo import GalileoObserve +from litellm.types.utils import ( + Choices, + EmbeddingResponse, + ImageObject, + ImageResponse, + Message, + ModelResponse, + TextCompletionResponse, +) + + +@pytest.fixture +def galileo_v2_env(monkeypatch): + monkeypatch.setenv("GALILEO_API_KEY", "test-api-key") + monkeypatch.setenv("GALILEO_PROJECT_ID", "86ff8ebe-a297-4134-b167-748bdd8d2c20") + monkeypatch.setenv("GALILEO_LOG_STREAM_ID", "76c4ea50-8aa3-4771-a0d7-8567b112210f") + monkeypatch.setenv("GALILEO_BASE_URL", "https://api.galileo.ai") + + +@pytest.mark.asyncio +async def test_galileo_v2_ingest_url_and_headers(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 100, + "status_code": 200, + "input_text": "hi", + "output_text": "hello", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 1, + "num_output_tokens": 2, + "created_at": "2026-05-25T12:00:00", + } + ] + + url, payload = logger._get_ingest_request() + assert ( + url + == "https://api.galileo.ai/v2/projects/86ff8ebe-a297-4134-b167-748bdd8d2c20/spans" + ) + assert payload["log_stream_id"] == "76c4ea50-8aa3-4771-a0d7-8567b112210f" + assert payload["spans"][0]["type"] == "llm" + assert payload["spans"][0]["output"]["content"] == "hello" + + assert await logger._ensure_headers() is True + assert logger.headers["Galileo-API-Key"] == "test-api-key" + + +def test_galileo_v2_span_preserves_message_roles(galileo_v2_env): + record = { + "latency_ms": 1, + "status_code": 200, + "input_text": "fallback", + "output_text": "ok", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hello"}, + ], + } + span = GalileoObserve._record_to_v2_span(record) + assert span["input"] == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hello"}, + ] + + +def test_galileo_output_text_from_model_response(galileo_v2_env): + logger = GalileoObserve() + response = ModelResponse( + choices=[ + Choices( + message=Message( + content="assistant reply", + role="assistant", + annotations=[], + ) + ) + ] + ) + + output = logger.get_output_str_from_response(response, {"call_type": "acompletion"}) + assert output == "assistant reply" + + +@pytest.mark.asyncio +async def test_galileo_flush_swallows_http_errors(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 1, + "status_code": 200, + "input_text": "a", + "output_text": "b", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + } + ] + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("404 Not Found") + await logger.flush_in_memory_records() + + assert len(logger.in_memory_records) == 1 + + +@pytest.mark.asyncio +async def test_galileo_flush_clears_records_on_201(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 1, + "status_code": 200, + "input_text": "a", + "output_text": "b", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + } + ] + + mock_response = AsyncMock() + mock_response.is_success = True + mock_response.status_code = 201 + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_response + await logger.flush_in_memory_records() + + assert logger.in_memory_records == [] + + +def test_galileo_normalize_base_url_none(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + assert logger.base_url is None + assert logger._normalize_base_url(None) is None + assert logger._normalize_base_url("https://x.example/") == "https://x.example" + + +def test_galileo_is_configured_branches(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + monkeypatch.delenv("GALILEO_USERNAME", raising=False) + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + + no_env = GalileoObserve() + assert no_env._is_configured() is False + + monkeypatch.setenv("GALILEO_API_KEY", "k") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + v2 = GalileoObserve() + assert v2._is_configured() is True + + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + legacy = GalileoObserve() + assert legacy._is_configured() is True + + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + no_pw = GalileoObserve() + assert no_pw._is_configured() is False + + +def test_galileo_input_messages_fallbacks(): + assert GalileoObserve._galileo_input_messages(None, "hi") == [ + {"role": "user", "content": "hi"} + ] + assert GalileoObserve._galileo_input_messages( + ["not-a-dict", {"content": "no role"}], "fallback" + ) == [{"role": "user", "content": "fallback"}] + + +def test_galileo_record_to_v2_span_with_tags_and_offset(): + span = GalileoObserve._record_to_v2_span( + { + "latency_ms": 5, + "status_code": 200, + "input_text": "in", + "output_text": "out", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 1, + "num_output_tokens": 2, + "created_at": "2026-05-25T12:00:00", + "tags": ["t1"], + } + ) + assert span["tags"] == ["t1"] + assert span["created_at"].endswith("Z") + + offset = GalileoObserve._record_to_v2_span( + {"created_at": "2026-05-25T12:00:00-05:00"} + ) + assert offset["created_at"] == "2026-05-25T12:00:00-05:00" + + +def test_galileo_get_output_str_variants(galileo_v2_env): + logger = GalileoObserve() + assert logger.get_output_str_from_response(None, {}) is None + assert ( + logger.get_output_str_from_response( + EmbeddingResponse(), {"call_type": "embedding"} + ) + is None + ) + + text_resp = TextCompletionResponse() + text_resp.choices = [MagicMock(text="text-completion-output")] + assert ( + logger.get_output_str_from_response(text_resp, {"call_type": "text_completion"}) + == "text-completion-output" + ) + + image_resp = ImageResponse(data=[ImageObject(url="https://x/y.png")]) + assert "y.png" in logger.get_output_str_from_response(image_resp, {}) + + assert logger.get_output_str_from_response("not-a-supported-type", {}) is None + + +def test_galileo_get_ingest_request_unconfigured(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + assert logger._get_ingest_request() is None + + +def test_galileo_get_ingest_request_legacy(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example/") + monkeypatch.setenv("GALILEO_PROJECT_ID", "proj") + logger = GalileoObserve() + logger.in_memory_records = [{"foo": "bar"}] + url, payload = logger._get_ingest_request() + assert url == "https://galileo.example/projects/proj/observe/ingest" + assert payload == {"records": [{"foo": "bar"}]} + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_v2_missing_key(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://x") + logger = GalileoObserve() + logger.use_v2_api = True + logger.api_key = None + assert await logger._ensure_headers() is False + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_cached(galileo_v2_env): + logger = GalileoObserve() + logger.headers = {"Galileo-API-Key": "already-set"} + assert await logger._ensure_headers() is True + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_legacy_login(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + + login_resp = MagicMock() + login_resp.raise_for_status = MagicMock() + login_resp.json = MagicMock(return_value={"access_token": "tok"}) + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = login_resp + assert await logger._ensure_headers() is True + + assert logger.headers["Authorization"] == "Bearer tok" + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_legacy_login_failure(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("boom") + assert await logger._ensure_headers() is False + + +@pytest.mark.asyncio +async def test_galileo_flush_noop_when_unconfigured(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + logger.in_memory_records = [{"foo": "bar"}] + await logger.flush_in_memory_records() + assert logger.in_memory_records == [{"foo": "bar"}] + + +@pytest.mark.asyncio +async def test_galileo_flush_resets_headers_on_401(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + logger.headers = {"Authorization": "Bearer stale"} + logger.in_memory_records = [{"records": "x"}] + + mock_response = MagicMock() + mock_response.is_success = False + mock_response.status_code = 401 + mock_response.text = "unauthorized" + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_response + await logger.flush_in_memory_records() + + assert logger.headers is None + assert logger.in_memory_records == [{"records": "x"}] + + +@pytest.mark.asyncio +async def test_galileo_async_log_success_appends_and_flushes(galileo_v2_env): + import datetime + + logger = GalileoObserve() + response = ModelResponse( + choices=[ + Choices(message=Message(content="reply", role="assistant", annotations=[])) + ], + usage={"prompt_tokens": 1, "completion_tokens": 2}, + ) + + flushed_url: dict = {} + mock_response = MagicMock() + mock_response.is_success = True + mock_response.status_code = 200 + + async def fake_post(**kwargs): + flushed_url["url"] = kwargs.get("url") + return mock_response + + with patch.object(logger.async_httpx_handler, "post", side_effect=fake_post): + await logger.async_log_success_event( + kwargs={ + "call_type": "acompletion", + "model": "gpt", + "messages": [{"role": "user", "content": "hi"}], + }, + response_obj=response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert "/v2/projects/" in flushed_url["url"] + assert logger.in_memory_records == []