mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-04-08 03:00:28 -04:00
updating get library agents functions
This commit is contained in:
@@ -25,6 +25,8 @@ async def get_library_agents(
|
||||
"Search query is too long."
|
||||
)
|
||||
|
||||
prisma.types.AgentGraphRelationFilter
|
||||
|
||||
where_clause = prisma.types.LibraryAgentWhereInput(
|
||||
userId=user_id,
|
||||
isDeleted=False,
|
||||
@@ -48,14 +50,17 @@ async def get_library_agents(
|
||||
]
|
||||
|
||||
try:
|
||||
|
||||
library_agents = await prisma.models.LibraryAgent.prisma().find_many(
|
||||
where=where_clause,
|
||||
include={
|
||||
"Agent": {
|
||||
"include": {
|
||||
"AgentNodes": {"include": {"Input": True, "Output": True}}
|
||||
}
|
||||
}
|
||||
"AgentNodes": {"include": {"Input": True, "Output": True}},
|
||||
"AgentGraphExecution": {"where": {"userId": user_id}},
|
||||
},
|
||||
},
|
||||
"Creator": True,
|
||||
},
|
||||
order=[{"updatedAt": "desc"}],
|
||||
)
|
||||
@@ -186,6 +191,8 @@ async def add_store_agent_to_library(
|
||||
f"Store listing version {store_listing_version_id} not found"
|
||||
)
|
||||
|
||||
# We need the agent object to be able to check if
|
||||
# the user_id is the same as the agent's user_id
|
||||
agent = store_listing_version.Agent
|
||||
|
||||
if agent.userId == user_id:
|
||||
@@ -353,3 +360,27 @@ async def delete_preset(user_id: str, preset_id: str) -> None:
|
||||
raise backend.server.v2.store.exceptions.DatabaseError(
|
||||
"Failed to delete preset"
|
||||
) from e
|
||||
|
||||
|
||||
async def main():
|
||||
import time
|
||||
|
||||
import backend.data.db
|
||||
|
||||
await backend.data.db.connect()
|
||||
|
||||
try:
|
||||
|
||||
time.sleep(2)
|
||||
library_agents = await get_library_agents(
|
||||
"5f8edafd-27ab-4343-88b3-6dbc02be7b06"
|
||||
)
|
||||
print(library_agents)
|
||||
finally:
|
||||
await backend.data.db.disconnect()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
|
||||
@@ -80,12 +80,8 @@ async def test_get_library_agents(mocker):
|
||||
assert result[0].id == "ua1"
|
||||
assert result[0].name == "Test Agent 2"
|
||||
assert result[0].description == "Test Description 2"
|
||||
assert result[0].is_created_by_user is False
|
||||
assert result[0].is_latest_version is True
|
||||
assert result[0].is_favorite is False
|
||||
assert result[0].agent_id == "agent2"
|
||||
assert result[0].agent_version == 1
|
||||
assert result[0].preset_id is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
import datetime
|
||||
import enum
|
||||
import json
|
||||
import typing
|
||||
|
||||
import prisma.enums
|
||||
import prisma.models
|
||||
import pydantic
|
||||
|
||||
@@ -10,25 +12,40 @@ import backend.data.graph
|
||||
import backend.server.model
|
||||
|
||||
|
||||
class AgentStatus(str, enum.Enum):
|
||||
# The agent has completed all runs
|
||||
COMPLETED = "COMPLETED"
|
||||
# An agent is running, but not all runs have completed
|
||||
HEALTHY = "HEALTHY"
|
||||
# An agent is waiting to start or waiting for another reason
|
||||
WAITING = "WAITING"
|
||||
# An agent is in an error state
|
||||
ERROR = "ERROR"
|
||||
|
||||
|
||||
class LibraryAgent(pydantic.BaseModel):
|
||||
id: str # Changed from agent_id to match GraphMeta
|
||||
|
||||
agent_id: str
|
||||
agent_version: int # Changed from agent_version to match GraphMeta
|
||||
|
||||
preset_id: str | None
|
||||
image_url: str
|
||||
|
||||
creator_name: str # from profile
|
||||
creator_image_url: str # from profile
|
||||
|
||||
status: AgentStatus
|
||||
|
||||
updated_at: datetime.datetime
|
||||
|
||||
name: str
|
||||
description: str
|
||||
name: str # from graph
|
||||
description: str # from graph
|
||||
|
||||
# Made input_schema and output_schema match GraphMeta's type
|
||||
input_schema: dict[str, typing.Any] # Should be BlockIOObjectSubSchema in frontend
|
||||
output_schema: dict[str, typing.Any] # Should be BlockIOObjectSubSchema in frontend
|
||||
|
||||
is_favorite: bool
|
||||
is_created_by_user: bool
|
||||
new_output: bool
|
||||
can_access_graph: bool
|
||||
|
||||
is_latest_version: bool
|
||||
|
||||
@@ -42,6 +59,18 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
agent_updated_at = agent.Agent.updatedAt
|
||||
lib_agent_updated_at = agent.updatedAt
|
||||
|
||||
name = graph.name
|
||||
description = graph.description
|
||||
image_url = agent.image_url if agent.image_url else ""
|
||||
if agent.Creator:
|
||||
creator_name = agent.Creator.name
|
||||
creator_image_url = (
|
||||
agent.Creator.avatarUrl if agent.Creator.avatarUrl else ""
|
||||
)
|
||||
else:
|
||||
creator_name = "Unknown"
|
||||
creator_image_url = ""
|
||||
|
||||
# Take the latest updated_at timestamp either when the graph was updated or the library agent was updated
|
||||
updated_at = (
|
||||
max(agent_updated_at, lib_agent_updated_at)
|
||||
@@ -49,22 +78,56 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
else lib_agent_updated_at
|
||||
)
|
||||
|
||||
# Getting counts as expecting more refined logic for determining status
|
||||
status_counts = {status: 0 for status in prisma.enums.AgentExecutionStatus}
|
||||
new_output = False
|
||||
|
||||
runs_since = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=7)
|
||||
if not agent.Agent.AgentGraphExecution:
|
||||
status = AgentStatus.COMPLETED
|
||||
else:
|
||||
for execution in agent.Agent.AgentGraphExecution:
|
||||
if runs_since > execution.createdAt:
|
||||
if (
|
||||
execution.executionStatus
|
||||
== prisma.enums.AgentExecutionStatus.COMPLETED
|
||||
):
|
||||
new_output = True
|
||||
status_counts[execution.executionStatus] += 1
|
||||
|
||||
if status_counts[prisma.enums.AgentExecutionStatus.FAILED] > 0:
|
||||
status = AgentStatus.ERROR
|
||||
elif status_counts[prisma.enums.AgentExecutionStatus.QUEUED] > 0:
|
||||
status = AgentStatus.WAITING
|
||||
elif status_counts[prisma.enums.AgentExecutionStatus.RUNNING] > 0:
|
||||
status = AgentStatus.HEALTHY
|
||||
else:
|
||||
status = AgentStatus.COMPLETED
|
||||
|
||||
return LibraryAgent(
|
||||
id=agent.id,
|
||||
agent_id=agent.agentId,
|
||||
agent_version=agent.agentVersion,
|
||||
image_url=image_url,
|
||||
creator_name=creator_name,
|
||||
creator_image_url=creator_image_url,
|
||||
name=name,
|
||||
description=description,
|
||||
status=status,
|
||||
updated_at=updated_at,
|
||||
name=graph.name,
|
||||
description=graph.description,
|
||||
input_schema=graph.input_schema,
|
||||
output_schema=graph.output_schema,
|
||||
is_favorite=agent.isFavorite,
|
||||
is_created_by_user=agent.isCreatedByUser,
|
||||
is_latest_version=graph.is_active,
|
||||
preset_id=agent.AgentPreset.id if agent.AgentPreset else None,
|
||||
new_output=new_output,
|
||||
can_access_graph=agent.Agent.userId == agent.userId,
|
||||
# TODO: work out how to calculate this efficiently
|
||||
is_latest_version=True,
|
||||
)
|
||||
|
||||
|
||||
class LibraryAgentResponse:
|
||||
agents: typing.List[LibraryAgent]
|
||||
pagination: backend.server.model.Pagination # info
|
||||
|
||||
|
||||
class LibraryAgentPreset(pydantic.BaseModel):
|
||||
id: str
|
||||
updated_at: datetime.datetime
|
||||
|
||||
@@ -2,154 +2,9 @@ import datetime
|
||||
|
||||
import prisma.models
|
||||
|
||||
import backend.data.block
|
||||
import backend.server.model
|
||||
import backend.server.v2.library.model
|
||||
|
||||
|
||||
def test_library_agent():
|
||||
agent = backend.server.v2.library.model.LibraryAgent(
|
||||
id="test-agent-123",
|
||||
agent_id="agent-123",
|
||||
agent_version=1,
|
||||
preset_id=None,
|
||||
updated_at=datetime.datetime.now(),
|
||||
name="Test Agent",
|
||||
description="Test description",
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
output_schema={"type": "object", "properties": {}},
|
||||
is_favorite=False,
|
||||
is_created_by_user=False,
|
||||
is_latest_version=True,
|
||||
)
|
||||
assert agent.id == "test-agent-123"
|
||||
assert agent.agent_id == "agent-123"
|
||||
assert agent.agent_version == 1
|
||||
assert agent.name == "Test Agent"
|
||||
assert agent.description == "Test description"
|
||||
assert agent.is_favorite is False
|
||||
assert agent.is_created_by_user is False
|
||||
assert agent.is_latest_version is True
|
||||
assert agent.input_schema == {"type": "object", "properties": {}}
|
||||
assert agent.output_schema == {"type": "object", "properties": {}}
|
||||
|
||||
|
||||
def test_library_agent_with_user_created():
|
||||
agent = backend.server.v2.library.model.LibraryAgent(
|
||||
id="user-agent-456",
|
||||
agent_id="agent-456",
|
||||
agent_version=2,
|
||||
preset_id=None,
|
||||
updated_at=datetime.datetime.now(),
|
||||
name="User Created Agent",
|
||||
description="An agent created by the user",
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
output_schema={"type": "object", "properties": {}},
|
||||
is_favorite=False,
|
||||
is_created_by_user=True,
|
||||
is_latest_version=True,
|
||||
)
|
||||
assert agent.id == "user-agent-456"
|
||||
assert agent.agent_id == "agent-456"
|
||||
assert agent.agent_version == 2
|
||||
assert agent.name == "User Created Agent"
|
||||
assert agent.description == "An agent created by the user"
|
||||
assert agent.is_favorite is False
|
||||
assert agent.is_created_by_user is True
|
||||
assert agent.is_latest_version is True
|
||||
assert agent.input_schema == {"type": "object", "properties": {}}
|
||||
assert agent.output_schema == {"type": "object", "properties": {}}
|
||||
|
||||
|
||||
def test_library_agent_preset():
|
||||
preset = backend.server.v2.library.model.LibraryAgentPreset(
|
||||
id="preset-123",
|
||||
name="Test Preset",
|
||||
description="Test preset description",
|
||||
agent_id="test-agent-123",
|
||||
agent_version=1,
|
||||
is_active=True,
|
||||
inputs={
|
||||
"input1": backend.data.block.BlockInput(
|
||||
name="input1",
|
||||
data={"type": "string", "value": "test value"},
|
||||
)
|
||||
},
|
||||
updated_at=datetime.datetime.now(),
|
||||
)
|
||||
assert preset.id == "preset-123"
|
||||
assert preset.name == "Test Preset"
|
||||
assert preset.description == "Test preset description"
|
||||
assert preset.agent_id == "test-agent-123"
|
||||
assert preset.agent_version == 1
|
||||
assert preset.is_active is True
|
||||
assert preset.inputs == {
|
||||
"input1": backend.data.block.BlockInput(
|
||||
name="input1", data={"type": "string", "value": "test value"}
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
def test_library_agent_preset_response():
|
||||
preset = backend.server.v2.library.model.LibraryAgentPreset(
|
||||
id="preset-123",
|
||||
name="Test Preset",
|
||||
description="Test preset description",
|
||||
agent_id="test-agent-123",
|
||||
agent_version=1,
|
||||
is_active=True,
|
||||
inputs={
|
||||
"input1": backend.data.block.BlockInput(
|
||||
name="input1",
|
||||
data={"type": "string", "value": "test value"},
|
||||
)
|
||||
},
|
||||
updated_at=datetime.datetime.now(),
|
||||
)
|
||||
|
||||
pagination = backend.server.model.Pagination(
|
||||
total_items=1, total_pages=1, current_page=1, page_size=10
|
||||
)
|
||||
|
||||
response = backend.server.v2.library.model.LibraryAgentPresetResponse(
|
||||
presets=[preset], pagination=pagination
|
||||
)
|
||||
|
||||
assert len(response.presets) == 1
|
||||
assert response.presets[0].id == "preset-123"
|
||||
assert response.pagination.total_items == 1
|
||||
assert response.pagination.total_pages == 1
|
||||
assert response.pagination.current_page == 1
|
||||
assert response.pagination.page_size == 10
|
||||
|
||||
|
||||
def test_create_library_agent_preset_request():
|
||||
request = backend.server.v2.library.model.CreateLibraryAgentPresetRequest(
|
||||
name="New Preset",
|
||||
description="New preset description",
|
||||
agent_id="agent-123",
|
||||
agent_version=1,
|
||||
is_active=True,
|
||||
inputs={
|
||||
"input1": backend.data.block.BlockInput(
|
||||
name="input1",
|
||||
data={"type": "string", "value": "test value"},
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
assert request.name == "New Preset"
|
||||
assert request.description == "New preset description"
|
||||
assert request.agent_id == "agent-123"
|
||||
assert request.agent_version == 1
|
||||
assert request.is_active is True
|
||||
assert request.inputs == {
|
||||
"input1": backend.data.block.BlockInput(
|
||||
name="input1", data={"type": "string", "value": "test value"}
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
def test_library_agent_from_db():
|
||||
# Create mock DB agent
|
||||
db_agent = prisma.models.AgentPreset(
|
||||
|
||||
@@ -39,29 +39,33 @@ def test_get_library_agents_success(mocker: pytest_mock.MockFixture):
|
||||
id="test-agent-1",
|
||||
agent_id="test-agent-1",
|
||||
agent_version=1,
|
||||
preset_id="preset-1",
|
||||
image_url="",
|
||||
creator_name="Test Creator",
|
||||
creator_image_url="",
|
||||
status=backend.server.v2.library.model.AgentStatus.COMPLETED,
|
||||
updated_at=datetime.datetime(2023, 1, 1, 0, 0, 0),
|
||||
is_favorite=False,
|
||||
is_created_by_user=True,
|
||||
is_latest_version=True,
|
||||
name="Test Agent 1",
|
||||
description="Test Description 1",
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
output_schema={"type": "object", "properties": {}},
|
||||
new_output=False,
|
||||
can_access_graph=True,
|
||||
is_latest_version=True,
|
||||
),
|
||||
backend.server.v2.library.model.LibraryAgent(
|
||||
id="test-agent-2",
|
||||
agent_id="test-agent-2",
|
||||
agent_version=1,
|
||||
preset_id="preset-2",
|
||||
image_url="",
|
||||
creator_name="Test Creator",
|
||||
creator_image_url="",
|
||||
status=backend.server.v2.library.model.AgentStatus.COMPLETED,
|
||||
updated_at=datetime.datetime(2023, 1, 1, 0, 0, 0),
|
||||
is_favorite=False,
|
||||
is_created_by_user=False,
|
||||
is_latest_version=True,
|
||||
name="Test Agent 2",
|
||||
description="Test Description 2",
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
output_schema={"type": "object", "properties": {}},
|
||||
new_output=False,
|
||||
can_access_graph=False,
|
||||
is_latest_version=True,
|
||||
),
|
||||
]
|
||||
mock_db_call = mocker.patch("backend.server.v2.library.db.get_library_agents")
|
||||
@@ -76,9 +80,9 @@ def test_get_library_agents_success(mocker: pytest_mock.MockFixture):
|
||||
]
|
||||
assert len(data) == 2
|
||||
assert data[0].agent_id == "test-agent-1"
|
||||
assert data[0].is_created_by_user is True
|
||||
assert data[0].can_access_graph is True
|
||||
assert data[1].agent_id == "test-agent-2"
|
||||
assert data[1].is_created_by_user is False
|
||||
assert data[1].can_access_graph is False
|
||||
mock_db_call.assert_called_once_with("test-user-id")
|
||||
|
||||
|
||||
|
||||
@@ -813,21 +813,21 @@ async def get_agent(
|
||||
)
|
||||
)
|
||||
|
||||
if not store_listing_version or not store_listing_version.Agent:
|
||||
if not store_listing_version:
|
||||
raise fastapi.HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Store listing version {store_listing_version_id} not found",
|
||||
)
|
||||
|
||||
graph_id = store_listing_version.agentId
|
||||
graph_version = store_listing_version.agentVersion
|
||||
graph = await backend.data.graph.get_graph(graph_id, graph_version)
|
||||
graph = await backend.data.graph.get_graph(
|
||||
store_listing_version.agentId, store_listing_version.agentVersion
|
||||
)
|
||||
|
||||
if not graph:
|
||||
raise fastapi.HTTPException(
|
||||
status_code=404,
|
||||
detail=(
|
||||
f"Agent #{graph_id} not found "
|
||||
f"Agent #{store_listing_version.agentId} not found "
|
||||
f"for store listing version #{store_listing_version_id}"
|
||||
),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
-- AlterTable
|
||||
ALTER TABLE "LibraryAgent" ADD COLUMN "creatorId" TEXT,
|
||||
ADD COLUMN "image_url" TEXT;
|
||||
|
||||
-- AddForeignKey
|
||||
ALTER TABLE "LibraryAgent" ADD CONSTRAINT "LibraryAgent_creatorId_fkey" FOREIGN KEY ("creatorId") REFERENCES "Profile"("id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
@@ -53,6 +53,7 @@ model AgentGraph {
|
||||
|
||||
name String?
|
||||
description String?
|
||||
|
||||
isActive Boolean @default(true)
|
||||
isTemplate Boolean @default(false)
|
||||
|
||||
@@ -122,6 +123,8 @@ model LibraryAgent {
|
||||
userId String
|
||||
User User @relation(fields: [userId], references: [id], onDelete: Cascade)
|
||||
|
||||
image_url String?
|
||||
|
||||
agentId String
|
||||
agentVersion Int
|
||||
Agent AgentGraph @relation(fields: [agentId, agentVersion], references: [id, version])
|
||||
@@ -129,6 +132,9 @@ model LibraryAgent {
|
||||
agentPresetId String?
|
||||
AgentPreset AgentPreset? @relation(fields: [agentPresetId], references: [id])
|
||||
|
||||
creatorId String?
|
||||
Creator Profile? @relation(fields: [creatorId], references: [id])
|
||||
|
||||
useGraphIsActiveVersion Boolean @default(false)
|
||||
|
||||
isFavorite Boolean @default(false)
|
||||
@@ -429,6 +435,8 @@ model Profile {
|
||||
|
||||
isFeatured Boolean @default(false)
|
||||
|
||||
LibraryAgent LibraryAgent[]
|
||||
|
||||
@@index([username])
|
||||
@@index([userId])
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user