Skip to main content

systemprompt_agent/
extension.rs

1use systemprompt_extension::prelude::*;
2
3#[derive(Debug, Clone, Copy, Default)]
4pub struct AgentExtension;
5
6impl Extension for AgentExtension {
7    fn metadata(&self) -> ExtensionMetadata {
8        ExtensionMetadata {
9            id: "agent",
10            name: "Agent",
11            version: env!("CARGO_PKG_VERSION"),
12        }
13    }
14
15    fn schemas(&self) -> Vec<SchemaDefinition> {
16        let mut schemas = conversation_schemas();
17        schemas.extend(artifact_schemas());
18        schemas.extend(context_schemas());
19        schemas.extend(task_tracking_schemas());
20        schemas.extend(service_schemas());
21        schemas
22    }
23
24    fn dependencies(&self) -> Vec<&'static str> {
25        vec!["users", "oauth", "mcp", "ai"]
26    }
27
28    fn migrations(&self) -> Vec<Migration> {
29        extension_migrations!()
30    }
31
32    fn cross_extension_tables(&self) -> Vec<&'static str> {
33        vec!["ai_requests"]
34    }
35}
36
37fn conversation_schemas() -> Vec<SchemaDefinition> {
38    vec![
39        SchemaDefinition::new("user_contexts", include_str!("../schema/user_contexts.sql"))
40            .with_required_columns(vec![
41                "context_id".into(),
42                "user_id".into(),
43                "created_at".into(),
44            ]),
45        SchemaDefinition::new("agent_tasks", include_str!("../schema/agent_tasks.sql"))
46            .with_required_columns(vec![
47                "task_id".into(),
48                "context_id".into(),
49                "status".into(),
50                "created_at".into(),
51            ]),
52        SchemaDefinition::new("task_messages", include_str!("../schema/task_messages.sql"))
53            .with_required_columns(vec![
54                "id".into(),
55                "task_id".into(),
56                "role".into(),
57                "created_at".into(),
58            ]),
59        SchemaDefinition::new("message_parts", include_str!("../schema/message_parts.sql"))
60            .with_required_columns(vec!["id".into(), "message_id".into(), "part_kind".into()]),
61    ]
62}
63
64fn context_schemas() -> Vec<SchemaDefinition> {
65    vec![
66        SchemaDefinition::new(
67            "context_agents",
68            include_str!("../schema/context_agents.sql"),
69        )
70        .with_required_columns(vec!["id".into(), "context_id".into(), "agent_name".into()]),
71        SchemaDefinition::new(
72            "context_notifications",
73            include_str!("../schema/context_notifications.sql"),
74        )
75        .with_required_columns(vec![
76            "id".into(),
77            "context_id".into(),
78            "notification_type".into(),
79        ]),
80    ]
81}
82
83fn task_tracking_schemas() -> Vec<SchemaDefinition> {
84    vec![
85        SchemaDefinition::new(
86            "task_push_notification_configs",
87            include_str!("../schema/task_push_notification_configs.sql"),
88        )
89        .with_required_columns(vec!["id".into(), "task_id".into()]),
90        SchemaDefinition::new(
91            "task_execution_steps",
92            include_str!("../schema/task_execution_steps.sql"),
93        )
94        .with_required_columns(vec![
95            "step_id".into(),
96            "task_id".into(),
97            "step_type".into(),
98        ]),
99    ]
100}
101
102fn artifact_schemas() -> Vec<SchemaDefinition> {
103    vec![
104        SchemaDefinition::new(
105            "task_artifacts",
106            include_str!("../schema/task_artifacts.sql"),
107        )
108        .with_required_columns(vec!["id".into(), "task_id".into(), "artifact_id".into()]),
109        SchemaDefinition::new(
110            "artifact_parts",
111            include_str!("../schema/artifact_parts.sql"),
112        )
113        .with_required_columns(vec!["id".into(), "artifact_id".into(), "part_kind".into()]),
114    ]
115}
116
117fn service_schemas() -> Vec<SchemaDefinition> {
118    vec![
119        SchemaDefinition::new("services", include_str!("../schema/services.sql"))
120            .with_required_columns(vec!["name".into(), "module_name".into(), "status".into()]),
121        SchemaDefinition::new(
122            "user_session_analytics",
123            include_str!("../schema/user_session_analytics.sql"),
124        ),
125    ]
126}
127
128register_extension!(AgentExtension);