feat: add agent publishing foundation
This commit is contained in:
a.tolmachev
@@ -1,4 +1,7 @@
|
||||
use crank_core::{AuthProfile, OperationId, OperationStatus, Workspace, WorkspaceId};
|
||||
use crank_core::{
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, OperationId,
|
||||
OperationStatus, Workspace, WorkspaceId,
|
||||
};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
use sqlx::{
|
||||
@@ -11,11 +14,13 @@ use crate::{
|
||||
error::RegistryError,
|
||||
migrations,
|
||||
model::{
|
||||
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest,
|
||||
DescriptorMetadata, OperationSampleMetadata, OperationSummary, OperationVersionRecord,
|
||||
PublishRequest, RegistryOperation, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
|
||||
SaveSampleMetadataRequest, UpdateWorkspaceRequest, WorkspaceRecord, YamlImportJob,
|
||||
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
AgentSummary, AgentVersionRecord, CreateAgentRequest, CreateVersionRequest,
|
||||
CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorMetadata,
|
||||
OperationSampleMetadata, OperationSummary, OperationVersionRecord, PublishAgentRequest,
|
||||
PublishRequest, PublishedAgentTool, RegistryOperation, SaveAgentBindingsRequest,
|
||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||
UpdateWorkspaceRequest, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
|
||||
YamlImportJobId, YamlImportJobStatus,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -158,6 +163,312 @@ impl PostgresRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_agents(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
) -> Result<Vec<AgentSummary>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
slug,
|
||||
display_name,
|
||||
description,
|
||||
status,
|
||||
current_draft_version,
|
||||
latest_published_version,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
|
||||
to_char(published_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as published_at
|
||||
from agents
|
||||
where workspace_id = $1
|
||||
order by slug asc",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_agent_summary).collect()
|
||||
}
|
||||
|
||||
pub async fn get_agent_summary(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
agent_id: &AgentId,
|
||||
) -> Result<Option<AgentSummary>, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
id,
|
||||
workspace_id,
|
||||
slug,
|
||||
display_name,
|
||||
description,
|
||||
status,
|
||||
current_draft_version,
|
||||
latest_published_version,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
|
||||
to_char(published_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as published_at
|
||||
from agents
|
||||
where workspace_id = $1 and id = $2",
|
||||
)
|
||||
.bind(workspace_id.as_str())
|
||||
.bind(agent_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.as_ref().map(map_agent_summary).transpose()
|
||||
}
|
||||
|
||||
pub async fn create_agent(&self, request: CreateAgentRequest<'_>) -> Result<(), RegistryError> {
|
||||
let mut tx = self.pool.begin().await?;
|
||||
|
||||
sqlx::query(
|
||||
"insert into agents (
|
||||
id,
|
||||
workspace_id,
|
||||
slug,
|
||||
display_name,
|
||||
description,
|
||||
status,
|
||||
current_draft_version,
|
||||
latest_published_version,
|
||||
created_at,
|
||||
updated_at,
|
||||
published_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9::timestamptz, $10::timestamptz, $11::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(request.agent.id.as_str())
|
||||
.bind(request.agent.workspace_id.as_str())
|
||||
.bind(&request.agent.slug)
|
||||
.bind(&request.agent.display_name)
|
||||
.bind(&request.agent.description)
|
||||
.bind(serialize_enum_text(&request.agent.status, "status")?)
|
||||
.bind(to_db_version(request.agent.current_draft_version))
|
||||
.bind(request.agent.latest_published_version.map(to_db_version))
|
||||
.bind(&request.agent.created_at)
|
||||
.bind(&request.agent.updated_at)
|
||||
.bind(request.agent.published_at.as_deref())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
insert_agent_version_row(&mut tx, request.version).await?;
|
||||
replace_agent_bindings_rows(
|
||||
&mut tx,
|
||||
&request.agent.id,
|
||||
request.version.version,
|
||||
request.bindings,
|
||||
)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_agent_version(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
agent_id: &AgentId,
|
||||
version: u32,
|
||||
) -> Result<Option<AgentVersionRecord>, RegistryError> {
|
||||
let summary = match self.get_agent_summary(workspace_id, agent_id).await? {
|
||||
Some(value) => value,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let row = sqlx::query(
|
||||
"select
|
||||
agent_id,
|
||||
version,
|
||||
status,
|
||||
instructions_json,
|
||||
tool_selection_policy_json,
|
||||
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at
|
||||
from agent_versions
|
||||
where agent_id = $1 and version = $2",
|
||||
)
|
||||
.bind(agent_id.as_str())
|
||||
.bind(to_db_version(version))
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
let Some(row) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let bindings = self.list_agent_bindings(agent_id, version).await?;
|
||||
Ok(Some(map_agent_version_record(&summary, &row, bindings)?))
|
||||
}
|
||||
|
||||
pub async fn save_agent_bindings(
|
||||
&self,
|
||||
request: SaveAgentBindingsRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
if self
|
||||
.get_agent_summary(request.workspace_id, request.agent_id)
|
||||
.await?
|
||||
.is_none()
|
||||
{
|
||||
return Err(RegistryError::AgentNotFound {
|
||||
agent_id: request.agent_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut tx = self.pool.begin().await?;
|
||||
replace_agent_bindings_rows(
|
||||
&mut tx,
|
||||
request.agent_id,
|
||||
request.agent_version,
|
||||
request.bindings,
|
||||
)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn publish_agent(
|
||||
&self,
|
||||
request: PublishAgentRequest<'_>,
|
||||
) -> Result<(), RegistryError> {
|
||||
if self
|
||||
.get_agent_version(request.workspace_id, request.agent_id, request.version)
|
||||
.await?
|
||||
.is_none()
|
||||
{
|
||||
return Err(RegistryError::AgentNotFound {
|
||||
agent_id: request.agent_id.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut tx = self.pool.begin().await?;
|
||||
|
||||
sqlx::query(
|
||||
"insert into published_agents (
|
||||
agent_id,
|
||||
version,
|
||||
published_at,
|
||||
published_by
|
||||
) values ($1, $2, $3::timestamptz, $4)
|
||||
on conflict(agent_id) do update set
|
||||
version = excluded.version,
|
||||
published_at = excluded.published_at,
|
||||
published_by = excluded.published_by",
|
||||
)
|
||||
.bind(request.agent_id.as_str())
|
||||
.bind(to_db_version(request.version))
|
||||
.bind(request.published_at)
|
||||
.bind(request.published_by)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"update agent_versions
|
||||
set status = $1
|
||||
where agent_id = $2 and version = $3",
|
||||
)
|
||||
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
|
||||
.bind(request.agent_id.as_str())
|
||||
.bind(to_db_version(request.version))
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"update agents
|
||||
set status = $1,
|
||||
latest_published_version = $2,
|
||||
published_at = $3::timestamptz,
|
||||
updated_at = $4::timestamptz
|
||||
where id = $5 and workspace_id = $6",
|
||||
)
|
||||
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
|
||||
.bind(to_db_version(request.version))
|
||||
.bind(request.published_at)
|
||||
.bind(request.published_at)
|
||||
.bind(request.agent_id.as_str())
|
||||
.bind(request.workspace_id.as_str())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_published_agent_tools_by_slug(
|
||||
&self,
|
||||
workspace_slug: &str,
|
||||
agent_slug: &str,
|
||||
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
w.id as workspace_id,
|
||||
w.slug as workspace_slug,
|
||||
a.id as agent_id,
|
||||
a.slug as agent_slug,
|
||||
b.tool_name,
|
||||
b.tool_title,
|
||||
coalesce(b.tool_description_override, ov.tool_description_json->>'description') as tool_description,
|
||||
o.id,
|
||||
o.name,
|
||||
o.display_name,
|
||||
o.protocol,
|
||||
to_char(o.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as operation_created_at,
|
||||
to_char(o.updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as operation_updated_at,
|
||||
to_char(o.published_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as operation_published_at,
|
||||
ov.version,
|
||||
ov.status,
|
||||
ov.target_json,
|
||||
ov.input_schema_json,
|
||||
ov.output_schema_json,
|
||||
ov.input_mapping_json,
|
||||
ov.output_mapping_json,
|
||||
ov.execution_config_json,
|
||||
ov.tool_description_json,
|
||||
ov.samples_json,
|
||||
ov.generated_draft_json,
|
||||
ov.config_export_json,
|
||||
ov.change_note,
|
||||
to_char(ov.created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
|
||||
ov.created_by
|
||||
from workspaces w
|
||||
join agents a on a.workspace_id = w.id
|
||||
join published_agents pa on pa.agent_id = a.id
|
||||
join agent_operation_bindings b on b.agent_id = a.id and b.agent_version = pa.version
|
||||
join operation_versions ov on ov.operation_id = b.operation_id and ov.version = b.operation_version
|
||||
join operations o on o.id = ov.operation_id and o.workspace_id = w.id
|
||||
where w.slug = $1 and a.slug = $2 and b.enabled = true
|
||||
order by b.tool_name asc",
|
||||
)
|
||||
.bind(workspace_slug)
|
||||
.bind(agent_slug)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
if rows.is_empty() {
|
||||
let exists = sqlx::query(
|
||||
"select 1
|
||||
from workspaces w
|
||||
join agents a on a.workspace_id = w.id
|
||||
join published_agents pa on pa.agent_id = a.id
|
||||
where w.slug = $1 and a.slug = $2",
|
||||
)
|
||||
.bind(workspace_slug)
|
||||
.bind(agent_slug)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
if exists.is_none() {
|
||||
return Err(RegistryError::PublishedAgentNotFound {
|
||||
workspace_slug: workspace_slug.to_owned(),
|
||||
agent_slug: agent_slug.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
rows.iter().map(map_published_agent_tool).collect()
|
||||
}
|
||||
|
||||
async fn connect_in_schema(
|
||||
database_url: &str,
|
||||
schema: Option<&str>,
|
||||
@@ -915,6 +1226,31 @@ impl PostgresRegistry {
|
||||
|
||||
row.as_ref().map(map_yaml_import_job).transpose()
|
||||
}
|
||||
|
||||
async fn list_agent_bindings(
|
||||
&self,
|
||||
agent_id: &AgentId,
|
||||
version: u32,
|
||||
) -> Result<Vec<AgentOperationBinding>, RegistryError> {
|
||||
let rows = sqlx::query(
|
||||
"select
|
||||
operation_id,
|
||||
operation_version,
|
||||
tool_name,
|
||||
tool_title,
|
||||
tool_description_override,
|
||||
enabled
|
||||
from agent_operation_bindings
|
||||
where agent_id = $1 and agent_version = $2
|
||||
order by tool_name asc",
|
||||
)
|
||||
.bind(agent_id.as_str())
|
||||
.bind(to_db_version(version))
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_agent_binding).collect()
|
||||
}
|
||||
}
|
||||
|
||||
async fn insert_version_row(
|
||||
@@ -968,6 +1304,77 @@ async fn insert_version_row(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_agent_version_row(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
version: &AgentVersion,
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"insert into agent_versions (
|
||||
agent_id,
|
||||
version,
|
||||
status,
|
||||
instructions_json,
|
||||
tool_selection_policy_json,
|
||||
created_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6::timestamptz
|
||||
)",
|
||||
)
|
||||
.bind(version.agent_id.as_str())
|
||||
.bind(to_db_version(version.version))
|
||||
.bind(serialize_enum_text(&version.status, "status")?)
|
||||
.bind(Json(version.instructions.clone()))
|
||||
.bind(Json(version.tool_selection_policy.clone()))
|
||||
.bind(&version.created_at)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn replace_agent_bindings_rows(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
agent_id: &AgentId,
|
||||
version: u32,
|
||||
bindings: &[AgentOperationBinding],
|
||||
) -> Result<(), RegistryError> {
|
||||
sqlx::query(
|
||||
"delete from agent_operation_bindings
|
||||
where agent_id = $1 and agent_version = $2",
|
||||
)
|
||||
.bind(agent_id.as_str())
|
||||
.bind(to_db_version(version))
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
for binding in bindings {
|
||||
sqlx::query(
|
||||
"insert into agent_operation_bindings (
|
||||
agent_id,
|
||||
agent_version,
|
||||
operation_id,
|
||||
operation_version,
|
||||
tool_name,
|
||||
tool_title,
|
||||
tool_description_override,
|
||||
enabled
|
||||
) values ($1, $2, $3, $4, $5, $6, $7, $8)",
|
||||
)
|
||||
.bind(agent_id.as_str())
|
||||
.bind(to_db_version(version))
|
||||
.bind(binding.operation_id.as_str())
|
||||
.bind(to_db_version(binding.operation_version))
|
||||
.bind(&binding.tool_name)
|
||||
.bind(&binding.tool_title)
|
||||
.bind(binding.tool_description_override.as_deref())
|
||||
.bind(binding.enabled)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn assert_immutable_fields(
|
||||
summary: &OperationSummary,
|
||||
snapshot: &RegistryOperation,
|
||||
@@ -1010,6 +1417,28 @@ fn map_workspace_record(row: &PgRow) -> Result<WorkspaceRecord, RegistryError> {
|
||||
})
|
||||
}
|
||||
|
||||
fn map_agent_summary(row: &PgRow) -> Result<AgentSummary, RegistryError> {
|
||||
Ok(AgentSummary {
|
||||
id: AgentId::new(row.try_get::<String, _>("id")?),
|
||||
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
|
||||
slug: row.try_get("slug")?,
|
||||
display_name: row.try_get("display_name")?,
|
||||
description: row.try_get("description")?,
|
||||
status: deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?,
|
||||
current_draft_version: from_db_version(
|
||||
row.try_get("current_draft_version")?,
|
||||
"current_draft_version",
|
||||
)?,
|
||||
latest_published_version: row
|
||||
.try_get::<Option<i32>, _>("latest_published_version")?
|
||||
.map(|value| from_db_version(value, "latest_published_version"))
|
||||
.transpose()?,
|
||||
created_at: row.try_get("created_at")?,
|
||||
updated_at: row.try_get("updated_at")?,
|
||||
published_at: row.try_get("published_at")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_operation_summary(row: &PgRow) -> Result<OperationSummary, RegistryError> {
|
||||
Ok(OperationSummary {
|
||||
id: OperationId::new(row.try_get::<String, _>("id")?),
|
||||
@@ -1102,6 +1531,60 @@ fn map_auth_profile(row: &PgRow) -> Result<AuthProfile, RegistryError> {
|
||||
})
|
||||
}
|
||||
|
||||
fn map_agent_binding(row: &PgRow) -> Result<AgentOperationBinding, RegistryError> {
|
||||
Ok(AgentOperationBinding {
|
||||
operation_id: OperationId::new(row.try_get::<String, _>("operation_id")?),
|
||||
operation_version: from_db_version(row.try_get("operation_version")?, "operation_version")?,
|
||||
tool_name: row.try_get("tool_name")?,
|
||||
tool_title: row.try_get("tool_title")?,
|
||||
tool_description_override: row.try_get("tool_description_override")?,
|
||||
enabled: row.try_get("enabled")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_agent_version_record(
|
||||
summary: &AgentSummary,
|
||||
row: &PgRow,
|
||||
bindings: Vec<AgentOperationBinding>,
|
||||
) -> Result<AgentVersionRecord, RegistryError> {
|
||||
let version = from_db_version(row.try_get("version")?, "version")?;
|
||||
let status = deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?;
|
||||
|
||||
Ok(AgentVersionRecord {
|
||||
agent_id: summary.id.clone(),
|
||||
workspace_id: summary.workspace_id.clone(),
|
||||
version,
|
||||
status,
|
||||
created_at: row.try_get("created_at")?,
|
||||
bindings,
|
||||
snapshot: AgentVersion {
|
||||
agent_id: summary.id.clone(),
|
||||
version,
|
||||
status,
|
||||
instructions: row.try_get::<Json<Value>, _>("instructions_json")?.0,
|
||||
tool_selection_policy: row
|
||||
.try_get::<Json<Value>, _>("tool_selection_policy_json")?
|
||||
.0,
|
||||
created_at: row.try_get("created_at")?,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn map_published_agent_tool(row: &PgRow) -> Result<PublishedAgentTool, RegistryError> {
|
||||
let record = map_operation_version_record(row)?;
|
||||
|
||||
Ok(PublishedAgentTool {
|
||||
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
|
||||
workspace_slug: row.try_get("workspace_slug")?,
|
||||
agent_id: AgentId::new(row.try_get::<String, _>("agent_id")?),
|
||||
agent_slug: row.try_get("agent_slug")?,
|
||||
tool_name: row.try_get("tool_name")?,
|
||||
tool_title: row.try_get("tool_title")?,
|
||||
tool_description: row.try_get("tool_description")?,
|
||||
operation: record.snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_sample_metadata(row: &PgRow) -> Result<OperationSampleMetadata, RegistryError> {
|
||||
Ok(OperationSampleMetadata {
|
||||
id: crank_core::SampleId::new(row.try_get::<String, _>("id")?),
|
||||
|
||||
Reference in New Issue
Block a user