feat: add IngestTraces client for dedicated ingest service#501
feat: add IngestTraces client for dedicated ingest service#501pratduv wants to merge 2 commits intosc-58230/local-ingest-modelsfrom
Conversation
pratduv
requested review from
anurag-lang
and removed request for
a team and
anurag-lang
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| async def ingest_traces(self, traces_ingest_request: TracesIngestRequest) -> dict[str, Any]: | ||
| if self.experiment_id: | ||
| traces_ingest_request.experiment_id = UUID(self.experiment_id) | ||
| elif self.log_stream_id: | ||
| traces_ingest_request.log_stream_id = UUID(self.log_stream_id) |
There was a problem hiding this comment.
IngestTraces.ingest_traces repeats almost the same flow as ingest_spans: set the experiment/log stream IDs, flip on LoggingMethod.python_client, build the ingest URL, log the payload size, and post via an httpx.AsyncClient. The only real differences between the two methods are which route/attribute they use and the log message. Any future change to headers, timeout, metrics, or logging therefore needs to be made twice in blocks that are otherwise identical. Can we extract a shared helper (e.g. _post_to_ingest(endpoint: str, payload_attr: str, log_message: str, request: BaseModel)) that handles the ID injection, LoggingMethod assignment, URL construction, httpx request, and logging, and then let ingest_traces/ingest_spans simply call it with the different route and attribute? That keeps the helper consistent (shown below) and leaves the call sites tiny:
async def _post_to_ingest(...):
if self.experiment_id:
request.experiment_id = UUID(self.experiment_id)
elif self.log_stream_id:
request.log_stream_id = UUID(self.log_stream_id)
request.logging_method = LoggingMethod.python_client
url = f"{self._get_ingest_base_url()}{endpoint}"
_logger.info(log_message, extra={"url": url, "project_id": self.project_id, "num_items": len(getattr(request, payload_attr))})
async with httpx.AsyncClient(...) as client:
response = await client.post(url, json=request.model_dump(mode="json"), headers=self._get_auth_headers())
response.raise_for_status()
return response.json()
async def ingest_traces(...):
return await self._post_to_ingest(Routes.ingest_traces.format(...), "traces", "Sending traces to ingest service", traces_ingest_request)
async def ingest_spans(...):
return await self._post_to_ingest(Routes.ingest_spans.format(...), "spans", "Sending spans to ingest service", spans_ingest_request)
This avoids the same control flow being duplicated across the two methods.
Finding type: Code Dedup and Conventions | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
| _logger.info( | ||
| "Sending traces to ingest service", | ||
| extra={"url": url, "project_id": self.project_id, "num_traces": len(traces_ingest_request.traces)}, | ||
| ) | ||
|
|
There was a problem hiding this comment.
ingest_traces logs the "Sending traces to ingest service" start event but never logs success or failure after httpx.AsyncClient.post/response.raise_for_status. This violates AGENTS.md's "Logging & sensitive-data handling" lifecycle requirement and leaves ingest writes unobservable. Can we add a completion log (success or failure with non-sensitive context) immediately after the POST/response.raise_for_status?
Finding type: AI Coding Guidelines | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In src/galileo/traces.py around lines 227 to 231, the ingest_traces method only logs the
start of the request but never logs success or failure after the HTTP POST. Modify
ingest_traces to wrap the httpx post/response.raise_for_status sequence in a try/except:
after a successful response.raise_for_status() emit an _logger.info() completion log
with non-sensitive context (url, project_id, num_traces, response.status_code); on
exceptions (catch httpx.HTTPStatusError and a generic Exception) emit an _logger.error()
with safe context (url, project_id, num_traces, status code or exception message
trimmed) without logging request/response bodies or auth headers, then re-raise the
exception to preserve existing error behavior.
| _logger.info( | ||
| "Sending spans to ingest service", | ||
| extra={"url": url, "project_id": self.project_id, "num_spans": len(spans_ingest_request.spans)}, | ||
| ) | ||
|
|
There was a problem hiding this comment.
ingest_spans only logs 'Sending spans to ingest service' before the HTTP POST and never emits a completion log. AGENTS.md requires start+completion lifecycle logs so operators lack success/failure context; can we add a completion log after response.raise_for_status()/on success that doesn't include sensitive data?
Finding type: AI Coding Guidelines | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In src/galileo/traces.py around lines 252 to 256, the ingest_spans method only logs the
start of the HTTP request and never emits a completion/success log. Add a safe
completion log immediately after response.raise_for_status() and before returning
response.json(), e.g. call _logger.info("Finished sending spans to ingest service",
extra={"url": url, "project_id": self.project_id, "num_spans":
len(spans_ingest_request.spans), "status_code": response.status_code}). Do not log
response body, headers, auth, or any sensitive data. Keep the rest of the method
unchanged.
pratduv
force-pushed
the
sc-58230/local-ingest-models
branch
from
1a29b4f to
2405b0e
Compare
Adds a new IngestTraces client that talks directly to the Go ingest service via httpx, activated when GALILEO_INGEST_URL is set. Falls back to the existing Traces client (main API) otherwise. [sc-58541]
pratduv
force-pushed
the
sc-58230/ingest-client
branch
from
cce7d18 to
4d196b7
Compare
| client = self._ingest_client or self._traces_client | ||
| await client.ingest_traces(traces_ingest_request) |
There was a problem hiding this comment.
Flush in batch mode can keep using the traces client instead of the dedicated ingest client when GALILEO_INGEST_URL is set after constructing GalileoLogger. Flush selects client via client = self._ingest_client or self._traces_client and immediately calls await client.ingest_traces(...) (lines 1929-1930), but _ingest_client is only created in __init__ so it stays None; can we call _create_ingest_client() before selecting the client (as async_ingest_traces does around 2124-2130) and apply the same lazy-creation fix to other similar call sites?
Finding type: Logical Bugs | Severity: 🔴 High
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In src/galileo/logger/logger.py around lines 1929-1930 (method _flush_batch), the code
picks client = self._ingest_client or self._traces_client and then calls await
client.ingest_traces(...). This fails if GALILEO_INGEST_URL was set after the
GalileoLogger was constructed because _ingest_client was never created. Change the
selection to lazily create the ingest client first (e.g. if self._ingest_client is None
and os.environ.get('GALILEO_INGEST_URL'): self._ingest_client =
self._create_ingest_client()), then set client = self._ingest_client or
self._traces_client (and ensure _traces_client is created if needed). Also update the
similar client-selection logic in _ingest_span_streaming (lines ~550-593) and
_ingest_trace_streaming (lines ~516-542) to lazily create _ingest_client before using it
so the dedicated ingest service is used if the env var is present.
| @patch("galileo.traces.GalileoPythonConfig") | ||
| def test_accepts_log_stream_id(self, mock_config_class) -> None: | ||
| mock_config_class.get.return_value = Mock() | ||
|
|
||
| client = IngestTraces(project_id=PROJECT_ID, log_stream_id=LOG_STREAM_ID) |
There was a problem hiding this comment.
test_accepts_log_stream_id and other methods in TestIngestTracesInit omit the sentence-case # Given:/When:/Then: comments required by AGENTS.md, so these tests deviate from the documented testing convention and reduce clarity. Can we add non-empty sentence-case # Given:/When:/Then: comments to each test in tests/?
Finding type: AI Coding Guidelines | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In tests/test_ingest_traces_client.py around lines 35 to 39 (and similarly for the other
test methods in this file such as those in TestIngestTracesInit and the request test
classes), the test methods lack the required sentence-case `# Given:/When:/Then:`
comments per AGENTS.md. Edit each test function to add three short, sentence-case
comments: `# Given:` describing the test setup, `# When:` describing the action under
test, and `# Then:` describing the expected outcome (do not leave any of them empty).
Keep existing code and assertions unchanged aside from inserting these comments in the
appropriate places immediately above the relevant code blocks.
| @respx.mock | ||
| @pytest.mark.asyncio | ||
| async def test_ingest_traces_posts_to_correct_url(self, client, monkeypatch) -> None: | ||
| # Given: an ingest URL is configured | ||
| monkeypatch.setenv("GALILEO_INGEST_URL", INGEST_URL) | ||
|
|
There was a problem hiding this comment.
TestIngestTracesRequest uses @respx.mock/httpx instead of the shared mock_request fixture required by AGENTS.md (lines 207–235) and tests/conftest.py. This bypasses the shared setup (including --disable-socket) and breaks the documented mocking convention; can we refactor these tests to use mock_request and the accompanying mocks instead of wiring respx manually?
Finding type: AI Coding Guidelines | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In tests/test_ingest_traces_client.py around lines 148 to 153, the tests in
TestIngestTracesRequest use @respx.mock and respx.post to stub HTTP calls which violates
AGENTS.md and bypasses the shared mock_request fixture. Refactor these tests to remove
@respx.mock and respx usage: have each async test accept the mock_request fixture (add
mock_request to the parameter list), use that fixture to register the expected ingest
URL and canned httpx response (matching the JSON bodies used today), and assert route
calls via the mock_request tracking API instead of respx. Apply the same replacement to
the other tests in TestIngestTracesRequest and TestIngestSpansRequest that use respx so
all HTTP interactions reuse mock_request and respect --disable-socket.
c9eb90e to
f2c95c6
Compare
1 participant
User description
Summary
IngestTracesclient class that communicates directly with the Go ingest service viahttpx, bypassing the auto-generated API clientGALILEO_INGEST_URLis set; falls back to the existingTracesclient otherwiseGalileoLoggeringestion paths (batch flush, streaming traces, streaming spans)Changes
New:
IngestTracesclient (src/galileo/traces.py)httpx.AsyncClient-based client targetingPOST /ingest/traces/{project_id}andPOST /ingest/spans/{project_id}GALILEO_INGEST_URLenv var, falling back toapi_url@async_warn_catch_exceptionfor resilient telemetry semanticsModified:
GalileoLogger(src/galileo/logger/logger.py)_ingest_clientfield, created only whenGALILEO_INGEST_URLis setself._ingest_client or self._traces_clientfallback pattern_ensure_project_and_log_stream,_create_traces_client, and_create_ingest_clientfor lazy creationNew routes (
src/galileo/constants/routes.py)ingest_traces = "/ingest/traces/{project_id}"ingest_spans = "/ingest/spans/{project_id}"Tests (
tests/test_ingest_traces_client.py)respx-mocked HTTP tests foringest_tracesandingest_spansWhy
The Go ingest service runs at a separate URL from the main API and exposes raw HTTP endpoints (no OpenAPI client). This PR adds a thin
httpxclient so the SDK can route ingestion traffic directly to it when configured, without changing any behavior for users who don't set the env var.Test plan
test_ingest_traces_client.pypassFixes sc-58541
Made with Cursor
Generated description
Below is a concise technical summary of the changes proposed in this PR:
Enable direct ingestion through the new
IngestTracesclient that targets dedicated ingest endpoints viahttpx, respectsGALILEO_INGEST_URL, and handles auth/logging metadata so traffic can bypass the autogenerated API client when needed. WireGalileoLoggeringestion paths to prefer this client, lazily creating it while preserving existingTracesfallbacks for batch and streaming flows.GalileoLoggerso every ingestion path prefers the ingest client when configured, lazily creates clients via shared helpers, and keeps theTracesclient as a fallback.Modified files (1)
Latest Contributors(2)
IngestTracesto call the dedicated ingest endpoints, resolve the service URL fromGALILEO_INGEST_URL, attach auth and SDK headers, and cover its behavior with HTTP mocks and constructor/auth URL tests.Modified files (3)
Latest Contributors(2)