359 lines
10 KiB
Rust
359 lines
10 KiB
Rust
use anyhow::Result;
|
|
use sqlx::{PgPool, FromRow};
|
|
use uuid::Uuid;
|
|
use serde::{Deserialize, Serialize};
|
|
use chrono::{DateTime, Utc};
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
|
pub struct AgentPrompt {
|
|
pub id: Uuid,
|
|
pub project_id: String,
|
|
pub name: String,
|
|
pub template: String,
|
|
pub target_model: Option<String>,
|
|
pub task_category: String,
|
|
pub usage_count: i64,
|
|
pub avg_quality: f32,
|
|
pub last_used: Option<DateTime<Utc>>,
|
|
pub active: bool,
|
|
pub version: i32,
|
|
pub tags: Vec<String>,
|
|
pub created_at: DateTime<Utc>,
|
|
pub updated_at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
|
pub struct AgentSkill {
|
|
pub id: Uuid,
|
|
pub project_id: String,
|
|
pub agent_id: String,
|
|
pub name: String,
|
|
pub description: String,
|
|
pub trigger_patterns: Vec<String>,
|
|
pub success_rate: f32,
|
|
pub invocation_count: i64,
|
|
pub avg_latency_ms: i64,
|
|
pub linked_prompts: Vec<Uuid>,
|
|
pub enabled: bool,
|
|
pub created_at: DateTime<Utc>,
|
|
pub updated_at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
|
pub struct AgentDecision {
|
|
pub id: Uuid,
|
|
pub project_id: String,
|
|
pub agent_id: String,
|
|
pub action: String,
|
|
pub reasoning: String,
|
|
pub alternatives: Vec<String>,
|
|
pub confidence: f32,
|
|
pub context_entities: Vec<Uuid>,
|
|
pub tool: Option<String>,
|
|
pub task: Option<String>,
|
|
pub outcome_success: Option<bool>,
|
|
pub outcome_quality: Option<f32>,
|
|
pub outcome_feedback: Option<String>,
|
|
pub outcome_recorded_at: Option<DateTime<Utc>>,
|
|
pub created_at: DateTime<Utc>,
|
|
pub updated_at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
|
pub struct RolePromptMapping {
|
|
pub id: Uuid,
|
|
pub project_id: String,
|
|
pub role_name: String,
|
|
pub prompt_id: Uuid,
|
|
pub priority: i32,
|
|
pub active: bool,
|
|
pub created_at: DateTime<Utc>,
|
|
pub updated_at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
|
pub struct AgentMetrics {
|
|
pub id: Uuid,
|
|
pub project_id: String,
|
|
pub agent_id: String,
|
|
pub requests_total: i64,
|
|
pub requests_success: i64,
|
|
pub requests_failed: i64,
|
|
pub average_latency_ms: f32,
|
|
pub p95_latency_ms: f32,
|
|
pub p99_latency_ms: f32,
|
|
pub error_rate: f32,
|
|
pub recorded_at: DateTime<Utc>,
|
|
}
|
|
|
|
pub struct AgentRepository {
|
|
pool: PgPool,
|
|
}
|
|
|
|
impl AgentRepository {
|
|
pub fn new(pool: PgPool) -> Self {
|
|
AgentRepository { pool }
|
|
}
|
|
|
|
pub async fn create_prompt(&self, prompt: AgentPrompt) -> Result<AgentPrompt> {
|
|
let result = sqlx::query_as::<_, AgentPrompt>(
|
|
r#"
|
|
INSERT INTO agent_prompt
|
|
(project_id, name, template, target_model, task_category, active, version, tags)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
|
RETURNING *
|
|
"#,
|
|
)
|
|
.bind(&prompt.project_id)
|
|
.bind(&prompt.name)
|
|
.bind(&prompt.template)
|
|
.bind(&prompt.target_model)
|
|
.bind(&prompt.task_category)
|
|
.bind(prompt.active)
|
|
.bind(prompt.version)
|
|
.bind(&prompt.tags)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn get_prompt(&self, id: Uuid) -> Result<Option<AgentPrompt>> {
|
|
let result = sqlx::query_as::<_, AgentPrompt>(
|
|
"SELECT * FROM agent_prompt WHERE id = $1"
|
|
)
|
|
.bind(id)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn list_prompts(&self, project_id: &str) -> Result<Vec<AgentPrompt>> {
|
|
let results = sqlx::query_as::<_, AgentPrompt>(
|
|
"SELECT * FROM agent_prompt WHERE project_id = $1 AND active = true ORDER BY created_at DESC"
|
|
)
|
|
.bind(project_id)
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
pub async fn update_prompt_usage(&self, id: Uuid, quality_score: f32) -> Result<()> {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE agent_prompt
|
|
SET usage_count = usage_count + 1,
|
|
avg_quality = (avg_quality * (usage_count) + $2) / (usage_count + 1),
|
|
last_used = NOW(),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
"#,
|
|
)
|
|
.bind(id)
|
|
.bind(quality_score)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn create_skill(&self, skill: AgentSkill) -> Result<AgentSkill> {
|
|
let result = sqlx::query_as::<_, AgentSkill>(
|
|
r#"
|
|
INSERT INTO agent_skill
|
|
(project_id, agent_id, name, description, enabled)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
RETURNING *
|
|
"#,
|
|
)
|
|
.bind(&skill.project_id)
|
|
.bind(&skill.agent_id)
|
|
.bind(&skill.name)
|
|
.bind(&skill.description)
|
|
.bind(skill.enabled)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn get_skill(&self, id: Uuid) -> Result<Option<AgentSkill>> {
|
|
let result = sqlx::query_as::<_, AgentSkill>(
|
|
"SELECT * FROM agent_skill WHERE id = $1"
|
|
)
|
|
.bind(id)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn list_skills(&self, project_id: &str, agent_id: &str) -> Result<Vec<AgentSkill>> {
|
|
let results = sqlx::query_as::<_, AgentSkill>(
|
|
"SELECT * FROM agent_skill WHERE project_id = $1 AND agent_id = $2 AND enabled = true ORDER BY created_at DESC"
|
|
)
|
|
.bind(project_id)
|
|
.bind(agent_id)
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
pub async fn create_decision(&self, decision: AgentDecision) -> Result<AgentDecision> {
|
|
let result = sqlx::query_as::<_, AgentDecision>(
|
|
r#"
|
|
INSERT INTO agent_decision
|
|
(project_id, agent_id, action, reasoning, confidence, tool, task)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
|
RETURNING *
|
|
"#,
|
|
)
|
|
.bind(&decision.project_id)
|
|
.bind(&decision.agent_id)
|
|
.bind(&decision.action)
|
|
.bind(&decision.reasoning)
|
|
.bind(decision.confidence)
|
|
.bind(&decision.tool)
|
|
.bind(&decision.task)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn record_decision_outcome(
|
|
&self,
|
|
id: Uuid,
|
|
success: bool,
|
|
quality: f32,
|
|
feedback: Option<&str>,
|
|
) -> Result<()> {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE agent_decision
|
|
SET outcome_success = $2,
|
|
outcome_quality = $3,
|
|
outcome_feedback = $4,
|
|
outcome_recorded_at = NOW(),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
"#,
|
|
)
|
|
.bind(id)
|
|
.bind(success)
|
|
.bind(quality)
|
|
.bind(feedback)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn create_role_mapping(&self, mapping: RolePromptMapping) -> Result<RolePromptMapping> {
|
|
let result = sqlx::query_as::<_, RolePromptMapping>(
|
|
r#"
|
|
INSERT INTO role_prompt_mapping
|
|
(project_id, role_name, prompt_id, priority, active)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
RETURNING *
|
|
"#,
|
|
)
|
|
.bind(&mapping.project_id)
|
|
.bind(&mapping.role_name)
|
|
.bind(mapping.prompt_id)
|
|
.bind(mapping.priority)
|
|
.bind(mapping.active)
|
|
.fetch_one(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn get_prompts_for_role(&self, project_id: &str, role_name: &str) -> Result<Vec<AgentPrompt>> {
|
|
let results = sqlx::query_as::<_, AgentPrompt>(
|
|
r#"
|
|
SELECT ap.* FROM agent_prompt ap
|
|
INNER JOIN role_prompt_mapping rpm ON ap.id = rpm.prompt_id
|
|
WHERE rpm.project_id = $1 AND rpm.role_name = $2 AND rpm.active = true
|
|
ORDER BY rpm.priority DESC, ap.created_at DESC
|
|
"#,
|
|
)
|
|
.bind(project_id)
|
|
.bind(role_name)
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
pub async fn save_metrics(&self, metrics: AgentMetrics) -> Result<()> {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO agent_metrics
|
|
(project_id, agent_id, requests_total, requests_success, requests_failed,
|
|
average_latency_ms, p95_latency_ms, p99_latency_ms, error_rate)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
|
|
ON CONFLICT (project_id, agent_id, DATE(recorded_at)) DO UPDATE SET
|
|
requests_total = EXCLUDED.requests_total,
|
|
requests_success = EXCLUDED.requests_success,
|
|
requests_failed = EXCLUDED.requests_failed,
|
|
average_latency_ms = EXCLUDED.average_latency_ms,
|
|
p95_latency_ms = EXCLUDED.p95_latency_ms,
|
|
p99_latency_ms = EXCLUDED.p99_latency_ms,
|
|
error_rate = EXCLUDED.error_rate
|
|
"#,
|
|
)
|
|
.bind(&metrics.project_id)
|
|
.bind(&metrics.agent_id)
|
|
.bind(metrics.requests_total)
|
|
.bind(metrics.requests_success)
|
|
.bind(metrics.requests_failed)
|
|
.bind(metrics.average_latency_ms)
|
|
.bind(metrics.p95_latency_ms)
|
|
.bind(metrics.p99_latency_ms)
|
|
.bind(metrics.error_rate)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn log_prompt_usage(
|
|
&self,
|
|
project_id: &str,
|
|
prompt_id: Uuid,
|
|
agent_id: Option<&str>,
|
|
model: Option<&str>,
|
|
input_tokens: Option<i32>,
|
|
output_tokens: Option<i32>,
|
|
quality_score: Option<f32>,
|
|
duration_ms: i64,
|
|
error_message: Option<&str>,
|
|
) -> Result<()> {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO prompt_usage_log
|
|
(project_id, prompt_id, agent_id, model_used, input_tokens, output_tokens,
|
|
quality_score, duration_ms, error_message)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
|
|
"#,
|
|
)
|
|
.bind(project_id)
|
|
.bind(prompt_id)
|
|
.bind(agent_id)
|
|
.bind(model)
|
|
.bind(input_tokens)
|
|
.bind(output_tokens)
|
|
.bind(quality_score)
|
|
.bind(duration_ms)
|
|
.bind(error_message)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(())
|
|
}
|
|
}
|