updating get library agents functions

This commit is contained in:
SwiftyOS
2025-02-03 15:37:27 +01:00
parent 8dacdd16f2
commit 582097221f
8 changed files with 145 additions and 182 deletions

View File

@@ -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())

View File

@@ -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

View File

@@ -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

View File

@@ -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(

View File

@@ -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")

View File

@@ -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}"
),
)

View File

@@ -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;

View File

@@ -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])
}