From f00678fd1c6c8f1ca5ea97c7e5f621535cc0b48c Mon Sep 17 00:00:00 2001 From: Zamil Majdy Date: Thu, 29 Jan 2026 22:02:46 -0600 Subject: [PATCH] fix: support lookup by library agent ID in addition to graph_id When users paste library URLs (e.g., /library/agents/{id}), the ID is the LibraryAgent primary key, not the graph_id. The previous code only looked up by graph_id, causing "agent not found" errors. Now get_library_agent_by_id() tries both lookup strategies: 1. First by graph_id (AgentGraph primary key) 2. Then by library agent ID (LibraryAgent primary key) This fixes the issue where users couldn't reference agents by pasting their library URLs in chat. --- .../chat/tools/agent_generator/__init__.py | 2 + .../chat/tools/agent_generator/core.py | 66 ++++++++++---- .../agent_generator/test_library_agents.py | 90 +++++++++++++++---- 3 files changed, 122 insertions(+), 36 deletions(-) diff --git a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/__init__.py b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/__init__.py index 42887d9ff4..1a636c41f7 100644 --- a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/__init__.py +++ b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/__init__.py @@ -16,6 +16,7 @@ from .core import ( get_agent_as_json, get_all_relevant_agents_for_generation, get_library_agent_by_graph_id, + get_library_agent_by_id, get_library_agents_for_generation, json_to_graph, save_agent_to_library, @@ -42,6 +43,7 @@ __all__ = [ "get_agent_as_json", "get_all_relevant_agents_for_generation", "get_library_agent_by_graph_id", + "get_library_agent_by_id", "get_library_agents_for_generation", "get_user_message_for_error", "is_external_service_configured", diff --git a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py index 950aa37924..940360adf6 100644 --- a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py +++ b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py @@ -126,33 +126,65 @@ def extract_uuids_from_text(text: str) -> list[str]: return list({m.lower() for m in matches}) -async def get_library_agent_by_graph_id( - user_id: str, graph_id: str +async def get_library_agent_by_id( + user_id: str, agent_id: str ) -> LibraryAgentSummary | None: - """Fetch a specific library agent by its graph_id. + """Fetch a specific library agent by its ID (library agent ID or graph_id). + + This function tries multiple lookup strategies: + 1. First tries to find by graph_id (AgentGraph primary key) + 2. If not found, tries to find by library agent ID (LibraryAgent primary key) + + This handles both cases: + - User provides graph_id (e.g., from AgentExecutorBlock) + - User provides library agent ID (e.g., from library URL) Args: user_id: The user ID - graph_id: The graph ID to look up + agent_id: The ID to look up (can be graph_id or library agent ID) Returns: LibraryAgentSummary if found, None otherwise """ + # Try 1: Look up by graph_id try: - agent = await library_db.get_library_agent_by_graph_id(user_id, graph_id) - if not agent: - return None - return LibraryAgentSummary( - graph_id=agent.graph_id, - graph_version=agent.graph_version, - name=agent.name, - description=agent.description, - input_schema=agent.input_schema, - output_schema=agent.output_schema, - ) + agent = await library_db.get_library_agent_by_graph_id(user_id, agent_id) + if agent: + logger.debug(f"Found library agent by graph_id: {agent.name}") + return LibraryAgentSummary( + graph_id=agent.graph_id, + graph_version=agent.graph_version, + name=agent.name, + description=agent.description, + input_schema=agent.input_schema, + output_schema=agent.output_schema, + ) except Exception as e: - logger.debug(f"Could not fetch library agent by graph_id {graph_id}: {e}") - return None + logger.debug(f"Could not fetch library agent by graph_id {agent_id}: {e}") + + # Try 2: Look up by library agent ID (primary key) + try: + agent = await library_db.get_library_agent(agent_id, user_id) + if agent: + logger.debug(f"Found library agent by library_id: {agent.name}") + return LibraryAgentSummary( + graph_id=agent.graph_id, + graph_version=agent.graph_version, + name=agent.name, + description=agent.description, + input_schema=agent.input_schema, + output_schema=agent.output_schema, + ) + except NotFoundError: + logger.debug(f"Library agent not found by library_id: {agent_id}") + except Exception as e: + logger.debug(f"Could not fetch library agent by library_id {agent_id}: {e}") + + return None + + +# Alias for backward compatibility +get_library_agent_by_graph_id = get_library_agent_by_id async def get_library_agents_for_generation( diff --git a/autogpt_platform/backend/test/agent_generator/test_library_agents.py b/autogpt_platform/backend/test/agent_generator/test_library_agents.py index ef179c6d55..a8f69fc09c 100644 --- a/autogpt_platform/backend/test/agent_generator/test_library_agents.py +++ b/autogpt_platform/backend/test/agent_generator/test_library_agents.py @@ -687,11 +687,11 @@ class TestExtractUuidsFromText: assert len(result) == 0 -class TestGetLibraryAgentByGraphId: - """Test get_library_agent_by_graph_id function.""" +class TestGetLibraryAgentById: + """Test get_library_agent_by_id function (and its alias get_library_agent_by_graph_id).""" @pytest.mark.asyncio - async def test_returns_agent_when_found(self): + async def test_returns_agent_when_found_by_graph_id(self): """Test that agent is returned when found by graph_id.""" mock_agent = MagicMock() mock_agent.graph_id = "agent-123" @@ -707,38 +707,90 @@ class TestGetLibraryAgentByGraphId: new_callable=AsyncMock, return_value=mock_agent, ): - result = await core.get_library_agent_by_graph_id("user-123", "agent-123") + result = await core.get_library_agent_by_id("user-123", "agent-123") assert result is not None assert result["graph_id"] == "agent-123" assert result["name"] == "Test Agent" @pytest.mark.asyncio - async def test_returns_none_when_not_found(self): - """Test that None is returned when agent not found.""" - with patch.object( - core.library_db, - "get_library_agent_by_graph_id", - new_callable=AsyncMock, - return_value=None, + async def test_falls_back_to_library_agent_id(self): + """Test that lookup falls back to library agent ID when graph_id not found.""" + mock_agent = MagicMock() + mock_agent.graph_id = "graph-456" # Different from the lookup ID + mock_agent.graph_version = 1 + mock_agent.name = "Library Agent" + mock_agent.description = "Found by library ID" + mock_agent.input_schema = {"properties": {}} + mock_agent.output_schema = {"properties": {}} + + with ( + patch.object( + core.library_db, + "get_library_agent_by_graph_id", + new_callable=AsyncMock, + return_value=None, # Not found by graph_id + ), + patch.object( + core.library_db, + "get_library_agent", + new_callable=AsyncMock, + return_value=mock_agent, # Found by library ID + ), ): - result = await core.get_library_agent_by_graph_id("user-123", "nonexistent") + result = await core.get_library_agent_by_id("user-123", "library-id-123") + + assert result is not None + assert result["graph_id"] == "graph-456" + assert result["name"] == "Library Agent" + + @pytest.mark.asyncio + async def test_returns_none_when_not_found_by_either_method(self): + """Test that None is returned when agent not found by either method.""" + with ( + patch.object( + core.library_db, + "get_library_agent_by_graph_id", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + core.library_db, + "get_library_agent", + new_callable=AsyncMock, + side_effect=core.NotFoundError("Not found"), + ), + ): + result = await core.get_library_agent_by_id("user-123", "nonexistent") assert result is None @pytest.mark.asyncio async def test_returns_none_on_exception(self): - """Test that None is returned when exception occurs.""" - with patch.object( - core.library_db, - "get_library_agent_by_graph_id", - new_callable=AsyncMock, - side_effect=Exception("Database error"), + """Test that None is returned when exception occurs in both lookups.""" + with ( + patch.object( + core.library_db, + "get_library_agent_by_graph_id", + new_callable=AsyncMock, + side_effect=Exception("Database error"), + ), + patch.object( + core.library_db, + "get_library_agent", + new_callable=AsyncMock, + side_effect=Exception("Database error"), + ), ): - result = await core.get_library_agent_by_graph_id("user-123", "agent-123") + result = await core.get_library_agent_by_id("user-123", "agent-123") assert result is None + @pytest.mark.asyncio + async def test_alias_works(self): + """Test that get_library_agent_by_graph_id is an alias for get_library_agent_by_id.""" + assert core.get_library_agent_by_graph_id is core.get_library_agent_by_id + class TestGetAllRelevantAgentsWithUuids: """Test UUID extraction in get_all_relevant_agents_for_generation."""