diff --git a/src/agentex/lib/adk/__init__.py b/src/agentex/lib/adk/__init__.py index 25b858485..d5be0ac52 100644 --- a/src/agentex/lib/adk/__init__.py +++ b/src/agentex/lib/adk/__init__.py @@ -29,6 +29,10 @@ from agentex.lib.adk._modules.tasks import TasksModule from agentex.lib.adk._modules.tracing import TracingModule, TurnSpan +# Data-source refs for lineage (SGP-6513); implementation lives in core.tracing +from agentex.lib.core.tracing import lineage +from agentex.lib.core.tracing.lineage import DataSourceRef, data_sources + # Unified harness surface (AGX1-375) from agentex.lib.core.harness import ( UnifiedEmitter, @@ -67,6 +71,10 @@ "events", "agent_task_tracker", "TurnSpan", + # Lineage data-source refs (SGP-6513) + "lineage", + "DataSourceRef", + "data_sources", # Checkpointing / LangGraph "create_checkpointer", "stream_langgraph_events", diff --git a/src/agentex/lib/adk/providers/_modules/sync_provider.py b/src/agentex/lib/adk/providers/_modules/sync_provider.py index 86696a2b5..120915eec 100644 --- a/src/agentex/lib/adk/providers/_modules/sync_provider.py +++ b/src/agentex/lib/adk/providers/_modules/sync_provider.py @@ -19,6 +19,7 @@ from agentex import AsyncAgentex from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items logger = make_logger(__name__) @@ -185,6 +186,9 @@ async def get_response( "new_items": new_items, "final_output": final_output, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return response else: @@ -303,6 +307,9 @@ async def stream_response( "new_items": new_items, "final_output": final_response_text if final_response_text else None, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) finally: # End the span after all events have been yielded await trace.end_span(span) diff --git a/src/agentex/lib/core/harness/tracer.py b/src/agentex/lib/core/harness/tracer.py index bf37bad30..34cd95616 100644 --- a/src/agentex/lib/core/harness/tracer.py +++ b/src/agentex/lib/core/harness/tracer.py @@ -6,6 +6,17 @@ from agentex.lib.core.harness.types import OpenSpan, CloseSpan, SpanSignal +try: + from agentex.lib.core.tracing.lineage import resolve_refs, merge_refs_into_data +except Exception: # keep the harness importable without optional tracing deps + + def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: # noqa: ARG001 + return [] + + def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: # noqa: ARG001 + return dict(data or {}) + + try: from agentex.lib.utils.logging import make_logger @@ -80,6 +91,11 @@ async def handle(self, signal: SpanSignal) -> None: task_id=self.task_id, ) if span is not None: + if signal.kind == "tool": + refs = resolve_refs(signal.name, signal.input if isinstance(signal.input, dict) else {}) + if refs: + data = span.data if isinstance(span.data, dict) else {} + span.data = merge_refs_into_data(data, refs) self._open[signal.key] = span elif isinstance(signal, CloseSpan): span = self._open.pop(signal.key, None) diff --git a/src/agentex/lib/core/services/adk/providers/openai.py b/src/agentex/lib/core/services/adk/providers/openai.py index a2513ea01..cc411dc30 100644 --- a/src/agentex/lib/core/services/adk/providers/openai.py +++ b/src/agentex/lib/core/services/adk/providers/openai.py @@ -25,6 +25,7 @@ from agentex.lib.utils.temporal import heartbeat_if_in_workflow from agentex.lib.core.tracing.tracer import AsyncTracer from agentex.lib.core.harness.emitter import UnifiedEmitter +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_update import StreamTaskMessageFull from agentex.types.task_message_content import ( TextContent, @@ -286,13 +287,17 @@ async def run_agent( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -431,13 +436,17 @@ async def run_agent_auto_send( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) tool_call_map: dict[str, Any] = {} @@ -646,13 +655,17 @@ async def run_agent_streamed( result = Runner.run_streamed(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -906,12 +919,16 @@ async def run_agent_streamed_auto_send( raise if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py index 7c8690f21..c985d5e65 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py @@ -64,6 +64,7 @@ from agentex.lib import adk from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_delta import TextDelta, ToolRequestDelta, ReasoningContentDelta, ReasoningSummaryDelta from agentex.types.task_message_update import StreamTaskMessageFull, StreamTaskMessageDelta from agentex.types.task_message_content import TextContent, ReasoningContent, ToolRequestContent, ToolResponseContent @@ -1257,6 +1258,9 @@ async def get_response( output_data["tool_outputs"] = tool_outputs span.output = output_data + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) # Streaming-only metrics. Token counters and the success request # counter are emitted by LLMMetricsHooks.on_llm_end so they fire diff --git a/src/agentex/lib/core/tracing/lineage.py b/src/agentex/lib/core/tracing/lineage.py new file mode 100644 index 000000000..75eaffdc0 --- /dev/null +++ b/src/agentex/lib/core/tracing/lineage.py @@ -0,0 +1,174 @@ +"""Data-source reference capture for lineage: tools declare which sources they +touch and the refs land in span data under the ``sgp.lineage.refs`` key.""" + +from __future__ import annotations + +import re +import json +from typing import Any, Literal, Callable, Iterable + +from pydantic import Field, BaseModel, field_validator + +try: + from agentex.lib.utils.logging import make_logger + + logger = make_logger(__name__) +except Exception: # ddtrace may be absent in some envs; fall back to stdlib + import logging + + logger = logging.getLogger(__name__) + +LINEAGE_REFS_KEY = "sgp.lineage.refs" + +# The URI arm of the lineage namespace identifier rule (namespace-conventions.md): +# lowercase scheme and host (dots/hyphens only — normalize `_` to `-`), one optional path segment. +_URI_NAMESPACE_RE = re.compile(r"^[a-z][a-z0-9._-]*://[a-z0-9.-]+(/[a-zA-Z0-9._-]*)?$") + +RefResolver = Callable[[dict[str, Any]], "list[DataSourceRef]"] + + +class DataSourceRef(BaseModel): + """One data source a tool call touched, as a lineage coordinate.""" + + namespace: str = Field(max_length=512) + name: str = Field(min_length=1, max_length=512) + version: str | None = Field(default=None, max_length=256) + role: Literal["input", "output"] = "input" + + def __init__(self, namespace: str | None = None, name: str | None = None, **kwargs: Any) -> None: + if namespace is not None: + kwargs["namespace"] = namespace + if name is not None: + kwargs["name"] = name + super().__init__(**kwargs) + + @field_validator("namespace") + @classmethod + def _namespace_is_uri_form(cls, value: str) -> str: + if not _URI_NAMESPACE_RE.match(value): + raise ValueError(f"namespace must be URI-form (scheme://system), got: {value!r}") + return value + + +class _ToolSources(BaseModel): + refs: list[DataSourceRef] = Field(default_factory=list) + resolver: RefResolver | None = None + + model_config = {"arbitrary_types_allowed": True} + + +_tool_sources: dict[str, _ToolSources] = {} + + +def register_tool_sources( + tool_name: str, + refs: Iterable[DataSourceRef] | None = None, + resolver: RefResolver | None = None, +) -> None: + """Declare the data sources a tool touches, keyed by its tool name. + + Use for tools the agent does not own (e.g. MCP proxy tools). Static refs and + a resolver over the tool's parsed arguments may be combined; repeated + registration for the same name replaces the prior entry. The registry is + process-wide: co-located agents sharing a tool name share (and overwrite) + one entry, so disambiguate shared names before co-locating agent types. + """ + _tool_sources[tool_name] = _ToolSources(refs=list(refs or []), resolver=resolver) + + +def data_sources(*refs: DataSourceRef, resolver: RefResolver | None = None) -> Callable[[Any], Any]: + """Decorator form of ``register_tool_sources`` for tools the agent owns. + + Works below or above ``@function_tool``: the tool name is taken from the + decorated object's ``name`` attribute when present, else ``__name__``. + """ + + def _register(obj: Any) -> Any: + tool_name = getattr(obj, "name", None) or getattr(obj, "__name__", None) + if isinstance(tool_name, str) and tool_name: + register_tool_sources(tool_name, refs=refs, resolver=resolver) + else: + logger.warning("data_sources could not determine a tool name for %r; refs not registered", obj) + return obj + + return _register + + +def clear_tool_sources() -> None: + """Reset the registry (test isolation).""" + _tool_sources.clear() + + +def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: + """Resolve registered refs for one tool call to serialized, deduplicated dicts. + + Resolver failures are logged and swallowed: ref capture must never break a + tool call or its tracing. + """ + entry = _tool_sources.get(tool_name) + if entry is None: + return [] + refs = list(entry.refs) + if entry.resolver is not None: + try: + refs.extend(entry.resolver(arguments or {})) + except Exception: + logger.warning("data-source resolver for tool %s failed; static refs kept", tool_name, exc_info=True) + return _dedupe(refs) + + +def resolve_refs_from_items(items: Iterable[Any]) -> list[dict[str, Any]]: + """Resolve refs across serialized run items, matching ``function_call`` entries. + + Accepts the item dicts the providers already build for span output; string + ``arguments`` are parsed as JSON for resolver-based registrations. + """ + refs: list[dict[str, Any]] = [] + for item in items: + if not isinstance(item, dict) or item.get("type") != "function_call": + continue + tool_name = item.get("name") + if not isinstance(tool_name, str) or not tool_name: + continue + arguments = item.get("arguments") + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + arguments = {} + refs.extend(resolve_refs(tool_name, arguments if isinstance(arguments, dict) else {})) + return _dedupe_dicts(refs) + + +def record(span: Any, refs: Iterable[DataSourceRef]) -> None: + """Attach refs to a manually managed span (no-op when the span is None).""" + if span is None: + return + merged = merge_refs_into_data(getattr(span, "data", None), _dedupe(list(refs))) + span.data = merged + + +def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: + """Merge serialized refs into a span data dict, deduplicating with any present.""" + out = dict(data) if isinstance(data, dict) else {} + if refs: + existing = out.get(LINEAGE_REFS_KEY) + combined = list(existing) if isinstance(existing, list) else [] + combined.extend(refs) + out[LINEAGE_REFS_KEY] = _dedupe_dicts(combined) + return out + + +def _dedupe(refs: list[DataSourceRef]) -> list[dict[str, Any]]: + return _dedupe_dicts([ref.model_dump(exclude_none=True) for ref in refs]) + + +def _dedupe_dicts(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + seen: set[tuple[Any, ...]] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + key = (ref.get("namespace"), ref.get("name"), ref.get("version"), ref.get("role")) + if key not in seen: + seen.add(key) + out.append(ref) + return out diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 6d186de5f..32b7bae73 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -65,6 +65,8 @@ def _add_source_to_span(span: Span, env_vars: EnvironmentVariables) -> None: span.data["__agent_name__"] = env_vars.AGENT_NAME if env_vars.AGENT_ID is not None: span.data["__agent_id__"] = env_vars.AGENT_ID + if env_vars.AGENT_VERSION is not None: + span.data["__agent_version__"] = env_vars.AGENT_VERSION def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan: diff --git a/src/agentex/lib/environment_variables.py b/src/agentex/lib/environment_variables.py index 3113b78f4..cbad0f2d8 100644 --- a/src/agentex/lib/environment_variables.py +++ b/src/agentex/lib/environment_variables.py @@ -24,6 +24,7 @@ class EnvVarKeys(str, Enum): AGENT_NAME = "AGENT_NAME" AGENT_DESCRIPTION = "AGENT_DESCRIPTION" AGENT_ID = "AGENT_ID" + AGENT_VERSION = "AGENT_VERSION" AGENT_API_KEY = "AGENT_API_KEY" # ACP Configuration ACP_URL = "ACP_URL" @@ -66,6 +67,8 @@ class EnvironmentVariables(BaseModel): AGENT_NAME: str AGENT_DESCRIPTION: str | None = None AGENT_ID: str | None = None + # Build/version discriminator (image tag or git sha), set by the deployment + AGENT_VERSION: str | None = None AGENT_API_KEY: str | None = None ACP_TYPE: str | None = "async" AGENT_INPUT_TYPE: str | None = None diff --git a/tests/lib/core/harness/test_tracer_lineage.py b/tests/lib/core/harness/test_tracer_lineage.py new file mode 100644 index 000000000..75799caee --- /dev/null +++ b/tests/lib/core/harness/test_tracer_lineage.py @@ -0,0 +1,53 @@ +"""SpanTracer stamps registered data-source refs onto tool spans (SGP-6513).""" + +import pytest + +from agentex.lib.core.harness.types import OpenSpan, CloseSpan +from agentex.lib.core.harness.tracer import SpanTracer +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + clear_tool_sources, + register_tool_sources, +) + +from ._fakes import FakeTracing + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +@pytest.mark.asyncio +async def test_tool_open_span_carries_registered_refs(): + register_tool_sources( + "query_guidance", + refs=[DataSourceRef("databricks://ey-tax", "guidance.rulings")], + resolver=lambda args: [DataSourceRef("elasticsearch://ey", args["index"])], + ) + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id="p1", tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="query_guidance", input={"index": "filings"})) + await tracer.handle(CloseSpan(key="c1", output={"ok": True}, is_complete=True)) + + (span,) = fake.ended_spans + namespaces = {ref["namespace"] for ref in span.data[LINEAGE_REFS_KEY]} + assert namespaces == {"databricks://ey-tax", "elasticsearch://ey"} + + +@pytest.mark.asyncio +async def test_unregistered_tool_and_reasoning_spans_carry_no_refs(): + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id=None, tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="unregistered", input={})) + await tracer.handle(CloseSpan(key="c1", output=None, is_complete=True)) + await tracer.handle(OpenSpan(key="reasoning:0", kind="reasoning", name="reasoning", input={})) + await tracer.handle(CloseSpan(key="reasoning:0", output="thought", is_complete=True)) + + for span in fake.ended_spans: + assert not (isinstance(span.data, dict) and LINEAGE_REFS_KEY in span.data) diff --git a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py index dc8bab127..4a233fb72 100644 --- a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py +++ b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py @@ -39,6 +39,30 @@ def _make_mock_sgp_span() -> MagicMock: return sgp_span +class TestSourceStamps: + def test_agent_identity_and_version_stamped_into_span_data(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE="async", AGENT_NAME="emu-tax", AGENT_ID="a1", AGENT_VERSION="sha-abc123") + span = _make_span() + _add_source_to_span(span, env) + assert span.data == { + "__source__": "agentex", + "__acp_type__": "async", + "__agent_name__": "emu-tax", + "__agent_id__": "a1", + "__agent_version__": "sha-abc123", + } + + def test_unset_identity_fields_are_omitted(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) + span = _make_span() + _add_source_to_span(span, env) + assert span.data == {"__source__": "agentex"} + + # --------------------------------------------------------------------------- # Sync processor tests # --------------------------------------------------------------------------- @@ -48,7 +72,7 @@ class TestSGPSyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.SGPClient"), patch( @@ -150,7 +174,7 @@ class TestSGPAsyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) mock_async_client = MagicMock() @@ -319,11 +343,9 @@ async def test_get_client_caches_per_event_loop(self): keepalive instead of paying a TLS handshake per span. """ mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: mock_sgp_cls.side_effect = lambda **kwargs: MagicMock() from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -365,11 +387,11 @@ def capture_limits(*args, **kwargs): return original_async_client(*args, **kwargs) mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ), patch("httpx.AsyncClient", side_effect=capture_limits): + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"), patch( + "httpx.AsyncClient", side_effect=capture_limits + ): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) @@ -380,8 +402,7 @@ def capture_limits(*args, **kwargs): assert len(captured_limits) == 1 max_keepalive = captured_limits[0].max_keepalive_connections assert max_keepalive is not None and max_keepalive > 0, ( - f"SGP async client should have keepalive enabled, got " - f"max_keepalive_connections={max_keepalive}" + f"SGP async client should have keepalive enabled, got max_keepalive_connections={max_keepalive}" ) def test_cache_is_weakkeydict_and_evicts_dead_loops(self): @@ -395,7 +416,7 @@ def test_cache_is_weakkeydict_and_evicts_dead_loops(self): import weakref mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -428,18 +449,14 @@ async def test_disabled_processor_returns_none_client(self): from agentex.lib.types.tracing import SGPTracingProcessorConfig mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) - processor = SGPAsyncTracingProcessor( - SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="") - ) + processor = SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) assert processor._get_client() is None assert mock_sgp_cls.call_count == 0 diff --git a/tests/lib/core/tracing/test_lineage.py b/tests/lib/core/tracing/test_lineage.py new file mode 100644 index 000000000..c0fc3ebb9 --- /dev/null +++ b/tests/lib/core/tracing/test_lineage.py @@ -0,0 +1,147 @@ +"""Unit tests for the data-source ref module (sgp.lineage.refs capture).""" + +import json + +import pytest +from pydantic import ValidationError + +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + record, + data_sources, + resolve_refs, + clear_tool_sources, + merge_refs_into_data, + register_tool_sources, + resolve_refs_from_items, +) + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +ES_REF = DataSourceRef("elasticsearch://ey-embryonic", "companies_v3") +DBX_REF = DataSourceRef("databricks://ey-tax", "guidance.rulings", role="input") + + +class TestDataSourceRef: + def test_positional_construction(self): + ref = DataSourceRef("s3://bucket", "key", version="v1", role="output") + assert ref.namespace == "s3://bucket" + assert ref.name == "key" + assert ref.version == "v1" + assert ref.role == "output" + + def test_non_uri_namespace_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("not-a-uri", "name") + + def test_underscore_host_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("mcp://ey_tax_server", "competitive-edge") + + def test_host_with_path_segment_allowed(self): + DataSourceRef("confluence://ey-tax/TAX", "page-123") + + def test_empty_name_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "") + + def test_bad_role_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "key", role="sideways") + + +class TestRegistryAndResolve: + def test_unregistered_tool_resolves_empty(self): + assert resolve_refs("unknown_tool", {}) == [] + + def test_static_refs(self): + register_tool_sources("search", refs=[ES_REF]) + refs = resolve_refs("search", {"q": "acme"}) + assert refs == [{"namespace": "elasticsearch://ey-embryonic", "name": "companies_v3", "role": "input"}] + + def test_resolver_refs_combined_with_static(self): + register_tool_sources( + "query_table", + refs=[ES_REF], + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + refs = resolve_refs("query_table", {"table": "guidance.rulings"}) + assert {r["namespace"] for r in refs} == {"elasticsearch://ey-embryonic", "databricks://ey-tax"} + + def test_resolver_failure_keeps_static_refs(self): + register_tool_sources("flaky", refs=[ES_REF], resolver=lambda args: args["missing"]) + refs = resolve_refs("flaky", {}) + assert len(refs) == 1 + + def test_reregistration_replaces(self): + register_tool_sources("search", refs=[ES_REF]) + register_tool_sources("search", refs=[DBX_REF]) + assert resolve_refs("search", {})[0]["namespace"] == "databricks://ey-tax" + + def test_dedupe(self): + register_tool_sources("search", refs=[ES_REF, ES_REF]) + assert len(resolve_refs("search", {})) == 1 + + +class TestDecorator: + def test_registers_by_function_name(self): + @data_sources(ES_REF) + def search_companies(q: str) -> str: + return q + + assert search_companies("x") == "x" + assert resolve_refs("search_companies", {}) != [] + + def test_registers_by_name_attribute(self): + class FakeFunctionTool: + name = "mcp_search" + + data_sources(DBX_REF)(FakeFunctionTool()) + assert resolve_refs("mcp_search", {}) != [] + + +class TestResolveFromItems: + def test_matches_function_call_items_and_parses_string_arguments(self): + register_tool_sources( + "query_table", + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + items = [ + {"type": "message", "content": []}, + {"type": "function_call", "name": "query_table", "arguments": json.dumps({"table": "t1"})}, + {"type": "function_call", "name": "unregistered", "arguments": "{}"}, + "not-a-dict", + ] + refs = resolve_refs_from_items(items) + assert refs == [{"namespace": "databricks://ey-tax", "name": "t1", "role": "input"}] + + def test_malformed_arguments_fall_back_to_static(self): + register_tool_sources("search", refs=[ES_REF]) + items = [{"type": "function_call", "name": "search", "arguments": "{not json"}] + assert len(resolve_refs_from_items(items)) == 1 + + +class TestRecordAndMerge: + def test_record_on_none_span_is_noop(self): + record(None, [ES_REF]) + + def test_record_merges_into_span_data(self): + class Span: + data = {"__span_type__": "CUSTOM"} + + span = Span() + record(span, [ES_REF]) + assert span.data["__span_type__"] == "CUSTOM" + assert span.data[LINEAGE_REFS_KEY][0]["name"] == "companies_v3" + + def test_merge_dedupes_against_existing(self): + data = merge_refs_into_data(None, [ES_REF.model_dump(exclude_none=True)]) + data = merge_refs_into_data(data, [ES_REF.model_dump(exclude_none=True)]) + assert len(data[LINEAGE_REFS_KEY]) == 1