diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index 261817f5..404882ba 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -1290,6 +1290,22 @@ class ToolContext: parent_category: str | None = None # Parent category for CloudEvents +def _normalize_observation_id(obs_id: str) -> str: + """Strip the display-format ``id:`` prefix from a model-supplied observation ID. + + Observations are presented to agents as ``[id:xxx]`` (see + ``Representation.str_with_ids``), and despite tool-schema instructions to + pass the bare ID, models sometimes copy the prefix verbatim. Since document + IDs are nanoids whose alphabet includes ``-`` and ``_``, only the ``id:`` + prefix and surrounding whitespace are stripped — anything more aggressive + could mangle legitimate IDs. + """ + obs_id = obs_id.strip() + if obs_id.lower().startswith("id:"): + obs_id = obs_id[3:] + return obs_id.strip() + + async def _handle_create_observations_impl( ctx: ToolContext, tool_input: dict[str, Any], @@ -1309,6 +1325,13 @@ async def _handle_create_observations_impl( obs["level"] = forced_level else: obs.setdefault("level", default_level) + # Models sometimes copy the display-format "id:" prefix into source_ids; + # normalize so provenance links reference real document IDs. + source_ids = obs.get("source_ids") + if isinstance(source_ids, list): + obs["source_ids"] = [ + _normalize_observation_id(s) for s in source_ids if isinstance(s, str) + ] # Validate observations individually so valid ones are still processed observations: list[schemas.ObservationInput] = [] @@ -2203,6 +2226,7 @@ async def _handle_get_reasoning_chain( observation_id = tool_input.get("observation_id") if not observation_id: return "ERROR: 'observation_id' is required" + observation_id = _normalize_observation_id(observation_id) direction = tool_input.get("direction", "both") if direction not in ("premises", "conclusions", "both"): diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index 8ddab1bc..12037bf7 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -34,6 +34,7 @@ from src.utils.agent_tools import ( _handle_search_messages, # pyright: ignore[reportPrivateUsage] _handle_search_messages_temporal, # pyright: ignore[reportPrivateUsage] _handle_update_peer_card, # pyright: ignore[reportPrivateUsage] + _normalize_observation_id, # pyright: ignore[reportPrivateUsage] _validate_peer_card_entry, # pyright: ignore[reportPrivateUsage] create_observations, create_tool_executor, @@ -247,6 +248,40 @@ class TestCreateObservations: assert doc.level == "deductive" assert doc.source_ids == ["premise1", "premise2"] + async def test_source_ids_display_prefix_is_stripped( + self, + db_session: AsyncSession, + make_tool_context: Callable[..., ToolContext], + ): + """Models sometimes copy the '[id:xxx]' display format into source_ids; + the prefix must be stripped so provenance links reference real IDs.""" + ctx = make_tool_context(current_messages=None) + + result = await _handle_create_observations( + ctx, + { + "observations": [ + { + "content": "Inferred preference for early mornings", + "source_ids": ["id:premise1", "ID:premise2"], + "premises": [ + "User schedules meetings before 9am", + "User mentions waking at 5:30", + ], + }, + ] + }, + ) + + assert "Created 1 observations" in result + + stmt = select(models.Document).where( + models.Document.content == "Inferred preference for early mornings" + ) + doc = (await db_session.execute(stmt)).scalar_one_or_none() + assert doc is not None + assert doc.source_ids == ["premise1", "premise2"] + async def test_empty_observations_list_returns_error( self, make_tool_context: Callable[..., ToolContext] ): @@ -475,6 +510,27 @@ class TestCreateObservations: create_documents.assert_not_awaited() +class TestNormalizeObservationId: + """Unit tests for _normalize_observation_id.""" + + @pytest.mark.parametrize( + "raw,expected", + [ + ("doc_abc123", "doc_abc123"), + ("id:doc_abc123", "doc_abc123"), + ("ID:doc_abc123", "doc_abc123"), + (" id:doc_abc123 ", "doc_abc123"), + ("id: doc_abc123", "doc_abc123"), + # nanoid alphabet includes '-' and '_'; these must survive untouched + ("3-bwp1hxCRkRbUh_nrqn0", "3-bwp1hxCRkRbUh_nrqn0"), + ("id:3-bwp1hxCRkRbUh_nrqn0", "3-bwp1hxCRkRbUh_nrqn0"), + ("_leading_underscore", "_leading_underscore"), + ], + ) + def test_normalization(self, raw: str, expected: str): + assert _normalize_observation_id(raw) == expected + + @pytest.mark.asyncio class TestDeleteObservations: """Tests for _handle_delete_observations."""