Skip to main content

ares/db/
tenant_agents.rs

1use sqlx::{PgPool, Row};
2use crate::types::{AppError, Result};
3use serde::{Deserialize, Serialize};
4use std::time::{SystemTime, UNIX_EPOCH};
5
6fn now_ts() -> i64 {
7    SystemTime::now()
8        .duration_since(UNIX_EPOCH)
9        .unwrap()
10        .as_secs() as i64
11}
12
13// =============================================================================
14// Structs
15// =============================================================================
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct TenantAgent {
19    pub id: String,
20    pub tenant_id: String,
21    pub agent_name: String,
22    pub display_name: String,
23    pub description: Option<String>,
24    pub config: serde_json::Value,
25    pub enabled: bool,
26    pub created_at: i64,
27    pub updated_at: i64,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct AgentTemplate {
32    pub id: String,
33    pub product_type: String,
34    pub agent_name: String,
35    pub display_name: String,
36    pub description: Option<String>,
37    pub config: serde_json::Value,
38    pub created_at: i64,
39}
40
41#[derive(Debug, Deserialize)]
42pub struct CreateTenantAgentRequest {
43    pub agent_name: String,
44    pub display_name: String,
45    pub description: Option<String>,
46    pub config: serde_json::Value,
47}
48
49#[derive(Debug, Deserialize)]
50pub struct UpdateTenantAgentRequest {
51    pub display_name: Option<String>,
52    pub description: Option<String>,
53    pub config: Option<serde_json::Value>,
54    pub enabled: Option<bool>,
55}
56
57// =============================================================================
58// Tenant Agent CRUD
59// =============================================================================
60
61pub async fn list_tenant_agents(pool: &PgPool, tenant_id: &str) -> Result<Vec<TenantAgent>> {
62    let rows = sqlx::query(
63        "SELECT id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at
64         FROM tenant_agents WHERE tenant_id = $1 ORDER BY agent_name"
65    )
66    .bind(tenant_id)
67    .fetch_all(pool)
68    .await
69    .map_err(|e| AppError::Database(e.to_string()))?;
70
71    rows.iter().map(|row| {
72        Ok(TenantAgent {
73            id: row.get("id"),
74            tenant_id: row.get("tenant_id"),
75            agent_name: row.get("agent_name"),
76            display_name: row.get("display_name"),
77            description: row.get("description"),
78            config: row.get::<serde_json::Value, _>("config"),
79            enabled: row.get("enabled"),
80            created_at: row.get("created_at"),
81            updated_at: row.get("updated_at"),
82        })
83    }).collect()
84}
85
86pub async fn get_tenant_agent(pool: &PgPool, tenant_id: &str, agent_name: &str) -> Result<TenantAgent> {
87    let row = sqlx::query(
88        "SELECT id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at
89         FROM tenant_agents WHERE tenant_id = $1 AND agent_name = $2"
90    )
91    .bind(tenant_id)
92    .bind(agent_name)
93    .fetch_optional(pool)
94    .await
95    .map_err(|e| AppError::Database(e.to_string()))?
96    .ok_or_else(|| AppError::NotFound(format!("Agent '{}' not found for tenant '{}'", agent_name, tenant_id)))?;
97
98    Ok(TenantAgent {
99        id: row.get("id"),
100        tenant_id: row.get("tenant_id"),
101        agent_name: row.get("agent_name"),
102        display_name: row.get("display_name"),
103        description: row.get("description"),
104        config: row.get::<serde_json::Value, _>("config"),
105        enabled: row.get("enabled"),
106        created_at: row.get("created_at"),
107        updated_at: row.get("updated_at"),
108    })
109}
110
111pub async fn create_tenant_agent(pool: &PgPool, tenant_id: &str, req: CreateTenantAgentRequest) -> Result<TenantAgent> {
112    let id = uuid::Uuid::new_v4().to_string();
113    let now = now_ts();
114
115    sqlx::query(
116        "INSERT INTO tenant_agents (id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at)
117         VALUES ($1, $2, $3, $4, $5, $6, true, $7, $7)"
118    )
119    .bind(&id)
120    .bind(tenant_id)
121    .bind(&req.agent_name)
122    .bind(&req.display_name)
123    .bind(&req.description)
124    .bind(&req.config)
125    .bind(now)
126    .execute(pool)
127    .await
128    .map_err(|e| AppError::Database(e.to_string()))?;
129
130    get_tenant_agent(pool, tenant_id, &req.agent_name).await
131}
132
133pub async fn update_tenant_agent(pool: &PgPool, tenant_id: &str, agent_name: &str, req: UpdateTenantAgentRequest) -> Result<TenantAgent> {
134    let now = now_ts();
135
136    // Fetch current state
137    let current = get_tenant_agent(pool, tenant_id, agent_name).await?;
138
139    let display_name = req.display_name.unwrap_or(current.display_name);
140    let description = req.description.or(current.description);
141    let config = req.config.unwrap_or(current.config);
142    let enabled = req.enabled.unwrap_or(current.enabled);
143
144    sqlx::query(
145        "UPDATE tenant_agents SET display_name = $1, description = $2, config = $3, enabled = $4, updated_at = $5
146         WHERE tenant_id = $6 AND agent_name = $7"
147    )
148    .bind(&display_name)
149    .bind(&description)
150    .bind(&config)
151    .bind(enabled)
152    .bind(now)
153    .bind(tenant_id)
154    .bind(agent_name)
155    .execute(pool)
156    .await
157    .map_err(|e| AppError::Database(e.to_string()))?;
158
159    get_tenant_agent(pool, tenant_id, agent_name).await
160}
161
162pub async fn delete_tenant_agent(pool: &PgPool, tenant_id: &str, agent_name: &str) -> Result<()> {
163    let result = sqlx::query(
164        "DELETE FROM tenant_agents WHERE tenant_id = $1 AND agent_name = $2"
165    )
166    .bind(tenant_id)
167    .bind(agent_name)
168    .execute(pool)
169    .await
170    .map_err(|e| AppError::Database(e.to_string()))?;
171
172    if result.rows_affected() == 0 {
173        return Err(AppError::NotFound(format!("Agent '{}' not found for tenant '{}'", agent_name, tenant_id)));
174    }
175    Ok(())
176}
177
178// =============================================================================
179// Template operations
180// =============================================================================
181
182pub async fn list_agent_templates(pool: &PgPool, product_type: Option<&str>) -> Result<Vec<AgentTemplate>> {
183    let rows = if let Some(pt) = product_type {
184        sqlx::query(
185            "SELECT id, product_type, agent_name, display_name, description, config, created_at
186             FROM agent_templates WHERE product_type = $1 ORDER BY agent_name"
187        )
188        .bind(pt)
189        .fetch_all(pool)
190        .await
191        .map_err(|e| AppError::Database(e.to_string()))?
192    } else {
193        sqlx::query(
194            "SELECT id, product_type, agent_name, display_name, description, config, created_at
195             FROM agent_templates ORDER BY product_type, agent_name"
196        )
197        .fetch_all(pool)
198        .await
199        .map_err(|e| AppError::Database(e.to_string()))?
200    };
201
202    rows.iter().map(|row| {
203        Ok(AgentTemplate {
204            id: row.get("id"),
205            product_type: row.get("product_type"),
206            agent_name: row.get("agent_name"),
207            display_name: row.get("display_name"),
208            description: row.get("description"),
209            config: row.get::<serde_json::Value, _>("config"),
210            created_at: row.get("created_at"),
211        })
212    }).collect()
213}
214
215/// Clones all agent templates for a product type into a tenant's agent list.
216/// Idempotent — skips agents that already exist (ON CONFLICT DO NOTHING).
217pub async fn clone_templates_for_tenant(
218    pool: &PgPool,
219    tenant_id: &str,
220    product_type: &str,
221) -> Result<Vec<TenantAgent>> {
222    let templates = list_agent_templates(pool, Some(product_type)).await?;
223    let now = now_ts();
224
225    for tpl in &templates {
226        let id = uuid::Uuid::new_v4().to_string();
227        sqlx::query(
228            "INSERT INTO tenant_agents (id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at)
229             VALUES ($1, $2, $3, $4, $5, $6, true, $7, $7)
230             ON CONFLICT (tenant_id, agent_name) DO NOTHING"
231        )
232        .bind(&id)
233        .bind(tenant_id)
234        .bind(&tpl.agent_name)
235        .bind(&tpl.display_name)
236        .bind(&tpl.description)
237        .bind(&tpl.config)
238        .bind(now)
239        .execute(pool)
240        .await
241        .map_err(|e| AppError::Database(e.to_string()))?;
242    }
243
244    list_tenant_agents(pool, tenant_id).await
245}
246
247// =============================================================================
248// Seed default templates
249// =============================================================================
250
251/// Seeds default agent templates. Idempotent — uses ON CONFLICT DO NOTHING.
252/// Called once on ARES startup after migrations.
253pub async fn seed_default_templates(pool: &PgPool) -> Result<()> {
254    let now = now_ts();
255
256    struct TemplateSpec {
257        product_type: &'static str,
258        agent_name: &'static str,
259        display_name: &'static str,
260        description: &'static str,
261        model: &'static str,
262        system_prompt: &'static str,
263    }
264
265    let templates: &[TemplateSpec] = &[
266        // Generic
267        TemplateSpec {
268            product_type: "generic",
269            agent_name: "assistant",
270            display_name: "General Assistant",
271            description: "Default conversational agent",
272            model: "fast",
273            system_prompt: "You are a helpful AI assistant. Answer questions clearly and concisely. If you don't know something, say so. Be direct and useful.",
274        },
275        // Kasino
276        TemplateSpec {
277            product_type: "kasino",
278            agent_name: "classifier",
279            display_name: "Domain Classifier",
280            description: "Classifies domains as gambling or safe",
281            model: "groq-fast",
282            system_prompt: r#"You are a gambling website classifier specializing in Indian gambling patterns.
283
284Given a domain name, SNI/page title, time of access, and recent user activity,
285determine if this is likely a gambling-related website.
286
287INDIAN GAMBLING PATTERNS TO WATCH FOR:
288- Satta/Matka: satta king, matka result, panel chart, gali, desawar, kalyan matka
289- Cricket betting: IPL odds, match prediction, session betting, live rate, bhav, fancy bet
290- Card games: teen patti cash, andar bahar real money, rummy cash game, poker real money
291- Casinos: live dealer, slot machines, jackpot, roulette, spin
292- Generic: bet, wager, odds, stake, bookmaker, bookie, punt
293
294SCORING:
295- 90-100: Confirmed gambling (known patterns, obvious indicators)
296- 70-89: Highly likely gambling (strong signals, new/unknown domain)
297- 40-69: Suspicious (some indicators, needs monitoring)
298- 0-39: Likely safe
299
300Respond ONLY with JSON:
301{
302  "gambling_score": <0-100>,
303  "confidence": <0.0-1.0>,
304  "category": "<satta|cricket_betting|casino|card_game|general_betting|safe>",
305  "action": "<block|flag|pass>",
306  "reasoning": "<brief explanation>"
307}"#,
308        },
309        TemplateSpec {
310            product_type: "kasino",
311            agent_name: "risk",
312            display_name: "Risk Assessor",
313            description: "Daily behavioral risk assessment",
314            model: "groq-balanced",
315            system_prompt: r#"You are a behavioral risk analyst for gambling addiction monitoring.
316
317Given 24 hours of device telemetry data, produce a risk assessment.
318
319CONSIDER:
320- Number and severity of blocked gambling attempts
321- Time-of-day patterns (late night activity = higher risk)
322- Financial transaction patterns (unusual amounts, frequency)
323- App usage anomalies (excessive browser time, new apps)
324- Attempted workarounds (VPN installs, settings access)
325- Communication patterns (messages about borrowing money)
326
327SCORING:
328- 0-20: Low risk (normal activity, no gambling indicators)
329- 21-40: Mild concern (minor anomalies, worth noting)
330- 41-60: Moderate risk (some gambling-related activity detected)
331- 61-80: High risk (active gambling attempts, financial anomalies)
332- 81-100: Critical (confirmed gambling activity, intervention needed)
333
334Respond with JSON:
335{
336  "risk_score": <0-100>,
337  "factors": ["<factor1>", "<factor2>"],
338  "trend": "<improving|stable|worsening>",
339  "summary": "<2-3 sentence assessment>",
340  "alerts": [{"severity": "<critical|warning|info>", "message": "<...>"}],
341  "positive_signals": ["<anything good to acknowledge>"]
342}"#,
343        },
344        TemplateSpec {
345            product_type: "kasino",
346            agent_name: "report",
347            display_name: "Report Generator",
348            description: "Weekly compassionate family report",
349            model: "groq-powerful",
350            system_prompt: r#"You are a compassionate counselor-assistant helping a family manage gambling addiction recovery. Generate a weekly report.
351
352AUDIENCE: Family members (mother with basic tech literacy, adult son).
353TONE: Caring but honest. Acknowledge progress. Don't minimize concerns.
354
355REPORT STRUCTURE:
3561. Overall Assessment (1-2 sentences, clear status)
3572. This Week's Highlights
358   - Positive behaviors (acknowledge these first)
359   - Concerning events (be specific but compassionate)
3603. Risk Trend (improving/stable/worsening with context)
3614. Key Numbers
362   - Gambling attempts blocked
363   - Financial transactions flagged
364   - Average daily risk score
3655. Recommendations
366   - What the family should discuss
367   - Any therapy-related suggestions
368   - Adjustments to monitoring if needed
3696. Encouragement (end on a supportive note)
370
371Write in simple English. Avoid technical jargon.
372Include specific numbers and dates where relevant.
373If the week was good, celebrate it. If it was concerning, be direct.
374
375Format the report in a clean, readable structure suitable for WhatsApp/Telegram delivery."#,
376        },
377        TemplateSpec {
378            product_type: "kasino",
379            agent_name: "transaction",
380            display_name: "Transaction Analyzer",
381            description: "Analyzes financial transactions for gambling",
382            model: "groq-fast",
383            system_prompt: r#"Analyze a financial transaction notification for gambling patterns.
384
385INPUT FORMAT: Indian bank SMS or UPI notification text.
386
387LOOK FOR:
388- Deposits to unknown/suspicious merchants
389- Round amounts to non-standard recipients
390- Rapid sequence of small transactions (multiple in short time)
391- Late night / early morning financial activity (11 PM - 6 AM)
392- Known gambling payment gateways or wallet names
393- Merchant names containing gambling keywords
394- P2P transfers to unknown numbers (possible bookie payments)
395
396KNOWN GAMBLING GATEWAYS:
397- Paytm merchant IDs for betting sites
398- UPI handles containing "bet", "game", "play"
399- International payment processors used by offshore casinos
400
401Respond with JSON:
402{
403  "is_suspicious": <true|false>,
404  "reason": "<explanation>",
405  "amount": <number>,
406  "merchant_category": "<standard|unknown|suspicious|confirmed_gambling>",
407  "risk_level": "<low|medium|high|critical>"
408}"#,
409        },
410        // EHB
411        TemplateSpec {
412            product_type: "ehb",
413            agent_name: "intake",
414            display_name: "Intake Agent",
415            description: "Patient intake and initial assessment",
416            model: "groq-balanced",
417            system_prompt: r#"You are a clinical intake assistant for eHealthBuddy. Conduct an initial patient assessment.
418
419PROCESS:
4201. Collect basic demographics (name, age, presenting complaint)
4212. Medical history (conditions, medications, allergies)
4223. Current symptoms (severity, duration, triggers)
4234. Mental health screening (PHQ-2, GAD-2 if appropriate)
4245. Social determinants (support system, barriers to care)
425
426OUTPUT: Structured JSON with all collected fields and a triage recommendation (routine/urgent/emergency).
427
428Be empathetic. Use simple language. One question at a time.
429Never diagnose -- you collect and organize, the clinician decides."#,
430        },
431        TemplateSpec {
432            product_type: "ehb",
433            agent_name: "followup",
434            display_name: "Follow-up Agent",
435            description: "Follow-up session management",
436            model: "groq-balanced",
437            system_prompt: r#"You are a follow-up care assistant for eHealthBuddy. Manage ongoing patient sessions.
438
439PROCESS:
4401. Review previous session summary from context
4412. Check on treatment adherence (medications, lifestyle changes)
4423. Assess symptom progression (better/same/worse)
4434. Note any new concerns
4445. Update care plan recommendations
445
446OUTPUT: Session notes in structured JSON with changes since last visit and updated recommendations.
447
448Be warm and supportive. Celebrate progress. Flag concerns without alarming."#,
449        },
450        TemplateSpec {
451            product_type: "ehb",
452            agent_name: "summary",
453            display_name: "Summary Agent",
454            description: "Generate patient session summaries",
455            model: "groq-powerful",
456            system_prompt: r#"You are a clinical documentation assistant for eHealthBuddy. Generate professional session summaries.
457
458INPUT: Raw session transcript or notes.
459
460OUTPUT: Structured clinical summary with:
4611. Chief Complaint
4622. History of Present Illness (HPI)
4633. Assessment
4644. Plan
4655. Follow-up timeline
466
467Write in clinical documentation style. Be precise. Include relevant quotes from patient.
468Flag any red flags (suicidal ideation, abuse indicators, medication non-compliance)."#,
469        },
470    ];
471
472    for tpl in templates {
473        let id = uuid::Uuid::new_v4().to_string();
474        let config = serde_json::json!({
475            "model": tpl.model,
476            "system_prompt": tpl.system_prompt,
477            "tools": [],
478            "max_tool_iterations": 3
479        });
480
481        sqlx::query(
482            "INSERT INTO agent_templates (id, product_type, agent_name, display_name, description, config, created_at)
483             VALUES ($1, $2, $3, $4, $5, $6, $7)
484             ON CONFLICT (product_type, agent_name) DO NOTHING"
485        )
486        .bind(&id)
487        .bind(tpl.product_type)
488        .bind(tpl.agent_name)
489        .bind(tpl.display_name)
490        .bind(tpl.description)
491        .bind(&config)
492        .bind(now)
493        .execute(pool)
494        .await
495        .map_err(|e| AppError::Database(format!("Failed to seed template {}/{}: {}", tpl.product_type, tpl.agent_name, e)))?;
496    }
497
498    tracing::info!("Agent templates seeded ({} templates)", templates.len());
499    Ok(())
500}