From 520a8b6be689377c37ca5e43bd8743eba03eb1a3 Mon Sep 17 00:00:00 2001 From: Hiram Chirino Date: Thu, 30 Jan 2025 17:55:21 -0500 Subject: [PATCH] fix: add an ai feature so that ai features are not compiled by default. Fixes #1219 Signed-off-by: Hiram Chirino --- .cargo/config.toml | 1 + .github/workflows/ci.yaml | 2 +- common/Cargo.toml | 3 + common/src/db/mod.rs | 4 + migration/Cargo.toml | 3 + migration/src/ai.rs | 19 + migration/src/lib.rs | 4 +- migration/tests/tests.rs | 20 ++ modules/fundamental/Cargo.toml | 6 +- modules/fundamental/src/ai/endpoints/mod.rs | 1 + modules/fundamental/src/ai/endpoints/test.rs | 25 +- modules/fundamental/src/endpoints.rs | 1 + modules/fundamental/src/lib.rs | 1 + openapi.yaml | 349 ------------------- server/Cargo.toml | 3 + trustd/Cargo.toml | 1 + xtask/Cargo.toml | 4 + xtask/src/main.rs | 3 + 18 files changed, 88 insertions(+), 362 deletions(-) create mode 100644 migration/src/ai.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index 950ddcc6c..7d37bdd19 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,5 +1,6 @@ [alias] xtask = "run --package xtask --" +xtask-ai = "run --features ai --package xtask --" [env] CARGO_WORKSPACE_ROOT = { value = "", relative = true } \ No newline at end of file diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index e118572c2..00f31fbbe 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -53,7 +53,7 @@ jobs: key: ${{ runner.os }}-theseus-postgresql-${{ hashFiles('**/Cargo.lock') }} - name: Test - run: cargo test -- --nocapture + run: cargo test --all-features -- --nocapture env: RUST_LOG: warn,sqlx=error,sea_orm=error - name: Export and Validate Generated Openapi Spec diff --git a/common/Cargo.toml b/common/Cargo.toml index ce818b675..a12ff4a8d 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true publish.workspace = true license.workspace = true +[features] +ai = [ "trustify-migration/ai" ] + [dependencies] trustify-migration = { workspace = true } diff --git a/common/src/db/mod.rs b/common/src/db/mod.rs index c125a386e..d823cfc04 100644 --- a/common/src/db/mod.rs +++ b/common/src/db/mod.rs @@ -47,6 +47,8 @@ impl Database { pub async fn migrate(&self) -> Result<(), anyhow::Error> { log::debug!("applying migrations"); Migrator::up(&self.db, None).await?; + #[cfg(feature = "ai")] + migration::ai::Migrator::up(&self.db, None).await?; log::debug!("applied migrations"); Ok(()) @@ -56,6 +58,8 @@ impl Database { pub async fn refresh(&self) -> Result<(), anyhow::Error> { log::warn!("refreshing database schema..."); Migrator::refresh(&self.db).await?; + #[cfg(feature = "ai")] + migration::ai::Migrator::refresh(&self.db).await?; log::warn!("refreshing database schema... done!"); Ok(()) diff --git a/migration/Cargo.toml b/migration/Cargo.toml index 01dcd350f..90b94090c 100644 --- a/migration/Cargo.toml +++ b/migration/Cargo.toml @@ -9,6 +9,9 @@ license.workspace = true name = "migration" path = "src/lib.rs" +[features] +ai = ["trustify-common/ai"] + [dependencies] sea-orm-migration = { workspace = true, features = ["runtime-tokio-rustls", "sqlx-postgres", "with-uuid"] } serde_json = { workspace = true } diff --git a/migration/src/ai.rs b/migration/src/ai.rs new file mode 100644 index 000000000..72f1ec0fb --- /dev/null +++ b/migration/src/ai.rs @@ -0,0 +1,19 @@ +pub use sea_orm_migration::prelude::*; + +pub struct Migrator; + +#[async_trait::async_trait] +impl MigratorTrait for Migrator { + fn migration_table_name() -> DynIden { + #[derive(DeriveIden)] + enum AiTables { + AiMigrations, + } + + AiTables::AiMigrations.into_iden() + } + + fn migrations() -> Vec> { + vec![Box::new(crate::m0000820_create_conversation::Migration)] + } +} diff --git a/migration/src/lib.rs b/migration/src/lib.rs index 1ddb7a8c6..de028e150 100644 --- a/migration/src/lib.rs +++ b/migration/src/lib.rs @@ -103,6 +103,9 @@ mod m0000830_perf_indexes; mod m0000840_add_relationship_14_15; mod m0000850_python_version; +#[cfg(feature = "ai")] +pub mod ai; + pub struct Migrator; #[async_trait::async_trait] @@ -207,7 +210,6 @@ impl MigratorTrait for Migrator { Box::new(m0000790_alter_sbom_alter_document_id::Migration), Box::new(m0000800_alter_product_version_range_scheme::Migration), Box::new(m0000810_fix_get_purl::Migration), - Box::new(m0000820_create_conversation::Migration), Box::new(m0000830_perf_indexes::Migration), Box::new(m0000840_add_relationship_14_15::Migration), Box::new(m0000850_python_version::Migration), diff --git a/migration/tests/tests.rs b/migration/tests/tests.rs index 052022b36..688606f81 100644 --- a/migration/tests/tests.rs +++ b/migration/tests/tests.rs @@ -21,6 +21,26 @@ async fn test_migrations(ctx: TrustifyContext) -> Result<(), anyhow::Error> { Ok(()) } +#[cfg(feature = "ai")] +#[test_context(TrustifyContext, skip_teardown)] +#[test(tokio::test)] +async fn test_ai_migrations(ctx: TrustifyContext) -> Result<(), anyhow::Error> { + let db = ctx.db; + + let migrations = migration::ai::Migrator::get_applied_migrations(&db).await?; + // 'Migrator.up' was called in bootstrap function when using TrustifyContext. + // At this point we already have migrations. + assert!(!migrations.is_empty()); + + db.refresh().await?; + + let rolled_back_and_reapplied_migrations = + migration::ai::Migrator::get_applied_migrations(&db).await?; + assert!(!rolled_back_and_reapplied_migrations.is_empty()); + + Ok(()) +} + #[test_context(TrustifyContext, skip_teardown)] #[test(tokio::test)] async fn only_up_migration(_ctx: TrustifyContext) -> Result<(), anyhow::Error> { diff --git a/modules/fundamental/Cargo.toml b/modules/fundamental/Cargo.toml index ca1038fcc..804967104 100644 --- a/modules/fundamental/Cargo.toml +++ b/modules/fundamental/Cargo.toml @@ -5,9 +5,13 @@ edition.workspace = true publish.workspace = true license.workspace = true +[features] +default = [] +ai = [ "trustify-common/ai" ] + [dependencies] trustify-auth = { workspace = true } -trustify-common = { workspace = true } +trustify-common = { workspace = true} trustify-cvss = { workspace = true } trustify-entity = { workspace = true } trustify-module-analysis = { workspace = true } diff --git a/modules/fundamental/src/ai/endpoints/mod.rs b/modules/fundamental/src/ai/endpoints/mod.rs index 1506db8d9..42b1c93a1 100644 --- a/modules/fundamental/src/ai/endpoints/mod.rs +++ b/modules/fundamental/src/ai/endpoints/mod.rs @@ -1,3 +1,4 @@ +#[cfg(feature = "ai")] #[cfg(test)] mod test; diff --git a/modules/fundamental/src/ai/endpoints/test.rs b/modules/fundamental/src/ai/endpoints/test.rs index 9b3c1e726..3f193e7ea 100644 --- a/modules/fundamental/src/ai/endpoints/test.rs +++ b/modules/fundamental/src/ai/endpoints/test.rs @@ -173,12 +173,17 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that there are no conversations let request = TestRequest::get() - .uri("/api/v1/ai/conversations") + .uri("/api/v2/ai/conversations") .to_request() .test_auth("user-a"); let response = app.call_service(request).await; - assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.status(), + StatusCode::OK, + "response: {:?}", + read_text(response).await + ); let result: PaginatedResults = read_body_json(response).await; assert_eq!(result.total, 0); @@ -186,7 +191,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Create a conversation let request = TestRequest::post() - .uri("/api/v1/ai/conversations") + .uri("/api/v2/ai/conversations") .to_request() .test_auth("user-a"); @@ -208,7 +213,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { )); let request = TestRequest::put() - .uri(format!("/api/v1/ai/conversations/{}", conversation_v1.id).as_str()) + .uri(format!("/api/v2/ai/conversations/{}", conversation_v1.id).as_str()) .set_json(update1.clone()) .to_request() .test_auth("user-a"); @@ -225,7 +230,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that the conversation can be listed let request = TestRequest::get() - .uri("/api/v1/ai/conversations") + .uri("/api/v2/ai/conversations") .to_request() .test_auth("user-a"); @@ -239,7 +244,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that we can retrieve the conversation by ID let request = TestRequest::get() - .uri(format!("/api/v1/ai/conversations/{}", conversation_v1.id).as_str()) + .uri(format!("/api/v2/ai/conversations/{}", conversation_v1.id).as_str()) .to_request() .test_auth("user-a"); @@ -256,7 +261,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { )); let request = TestRequest::put() - .uri(format!("/api/v1/ai/conversations/{}", conversation_v1.id).as_str()) + .uri(format!("/api/v2/ai/conversations/{}", conversation_v1.id).as_str()) .append_header(("if-match", format!("\"{}\"", conversation_v2.seq).as_str())) .set_json(update2.clone()) .to_request() @@ -274,7 +279,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that we can retrieve the updated conversation by ID let request = TestRequest::get() - .uri(format!("/api/v1/ai/conversations/{}", conversation_v1.id).as_str()) + .uri(format!("/api/v2/ai/conversations/{}", conversation_v1.id).as_str()) .to_request() .test_auth("user-a"); @@ -286,7 +291,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that we can delete the conversation let request = TestRequest::delete() - .uri(format!("/api/v1/ai/conversations/{}", conversation_v1.id).as_str()) + .uri(format!("/api/v2/ai/conversations/{}", conversation_v1.id).as_str()) .to_request() .test_auth("user-a"); @@ -295,7 +300,7 @@ async fn conversation_crud(ctx: &TrustifyContext) -> anyhow::Result<()> { // Verify that the conversation is deleted let request = TestRequest::get() - .uri("/api/v1/ai/conversations") + .uri("/api/v2/ai/conversations") .to_request() .test_auth("user-a"); diff --git a/modules/fundamental/src/endpoints.rs b/modules/fundamental/src/endpoints.rs index 02efc5db8..fa57f3d5d 100644 --- a/modules/fundamental/src/endpoints.rs +++ b/modules/fundamental/src/endpoints.rs @@ -24,6 +24,7 @@ pub fn configure( crate::advisory::endpoints::configure(svc, db.clone(), config.advisory_upload_limit); crate::license::endpoints::configure(svc, db.clone()); + #[cfg(feature = "ai")] crate::ai::endpoints::configure(svc, db.clone()); crate::organization::endpoints::configure(svc, db.clone()); crate::purl::endpoints::configure(svc, db.clone()); diff --git a/modules/fundamental/src/lib.rs b/modules/fundamental/src/lib.rs index 601726606..3ee3576dd 100644 --- a/modules/fundamental/src/lib.rs +++ b/modules/fundamental/src/lib.rs @@ -1,4 +1,5 @@ pub mod advisory; +#[cfg(feature = "ai")] pub mod ai; pub mod endpoints; pub mod error; diff --git a/openapi.yaml b/openapi.yaml index 417aae42b..1dc633403 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -226,239 +226,6 @@ paths: format: binary '404': description: The document could not be found - /api/v2/ai/completions: - post: - tags: - - ai - operationId: completions - requestBody: - content: - application/json: - schema: - $ref: '#/components/schemas/ChatState' - required: true - responses: - '200': - description: The resulting completion - content: - application/json: - schema: - $ref: '#/components/schemas/ChatState' - '400': - description: The request was invalid - '404': - description: The AI service is not enabled - /api/v2/ai/conversations: - get: - tags: - - ai - operationId: listConversations - parameters: - - name: q - in: query - required: false - schema: - type: string - - name: sort - in: query - required: false - schema: - type: string - - name: offset - in: query - description: |- - The first item to return, skipping all that come before it. - - NOTE: The order of items is defined by the API being called. - required: false - schema: - type: integer - format: int64 - minimum: 0 - - name: limit - in: query - description: |- - The maximum number of entries to return. - - Zero means: no limit - required: false - schema: - type: integer - format: int64 - minimum: 0 - responses: - '200': - description: The resulting list of conversation summaries - content: - application/json: - schema: - $ref: '#/components/schemas/PaginatedResults_ConversationSummary' - '404': - description: The AI service is not enabled - post: - tags: - - ai - operationId: createConversation - responses: - '200': - description: The resulting conversation - content: - application/json: - schema: - $ref: '#/components/schemas/Conversation' - '400': - description: The request was invalid - '404': - description: The AI service is not enabled - /api/v2/ai/conversations/{id}: - get: - tags: - - ai - operationId: getConversation - parameters: - - name: id - in: path - description: Opaque ID of the conversation - required: true - schema: - type: string - format: uuid - responses: - '200': - description: The resulting conversation - headers: - etag: - schema: - type: string - description: Sequence ID - content: - application/json: - schema: - $ref: '#/components/schemas/Conversation' - '400': - description: The request was invalid - '404': - description: The AI service is not enabled - put: - tags: - - ai - operationId: updateConversation - parameters: - - name: id - in: path - description: Opaque ID of the conversation - required: true - schema: - type: string - format: uuid - - name: if-match - in: header - description: The revision to update - required: false - schema: - type: - - string - - 'null' - requestBody: - content: - application/json: - schema: - type: array - items: - $ref: '#/components/schemas/ChatMessage' - required: true - responses: - '200': - description: The resulting conversation - content: - application/json: - schema: - $ref: '#/components/schemas/Conversation' - '400': - description: The request was invalid - '404': - description: The AI service is not enabled or the conversation was not found - delete: - tags: - - ai - operationId: deleteConversation - parameters: - - name: id - in: path - description: Opaque ID of the conversation - required: true - schema: - type: string - format: uuid - responses: - '200': - description: The resulting conversation - content: - application/json: - schema: - $ref: '#/components/schemas/Conversation' - '400': - description: The request was invalid - '404': - description: The AI service is not enabled or the conversation was not found - /api/v2/ai/flags: - get: - tags: - - ai - operationId: aiFlags - responses: - '200': - description: The resulting Flags - content: - application/json: - schema: - $ref: '#/components/schemas/AiFlags' - '404': - description: The AI service is not enabled - /api/v2/ai/tools: - get: - tags: - - ai - operationId: aiTools - responses: - '200': - description: The resulting list of tools - content: - application/json: - schema: - type: array - items: - $ref: '#/components/schemas/AiTool' - '404': - description: The AI service is not enabled - /api/v2/ai/tools/{name}: - post: - tags: - - ai - operationId: aiToolCall - parameters: - - name: name - in: path - description: Name of the tool to call - required: true - schema: - type: string - requestBody: - content: - application/json: - schema: {} - required: true - responses: - '200': - description: The result of the tool call - content: - text/plain: - schema: - type: string - '400': - description: The tool request was invalid - '404': - description: The tool was not found /api/v2/analysis/component/{key}: get: tags: @@ -2367,25 +2134,6 @@ components: All CVSS3 scores from the advisory for the given vulnerability. May include several, varying by minor version of the CVSS3 vector. description: Summary of information from this advisory regarding a single specific vulnerability. - AiFlags: - type: object - required: - - completions - properties: - completions: - type: boolean - AiTool: - type: object - required: - - name - - description - - parameters - properties: - description: - type: string - name: - type: string - parameters: {} AnalysisStatus: type: object required: @@ -2509,33 +2257,6 @@ components: type: string BinaryByteSize: type: string - ChatMessage: - type: object - required: - - message_type - - content - - timestamp - properties: - content: - type: string - message_type: - $ref: '#/components/schemas/MessageType' - timestamp: - type: string - format: date-time - ChatState: - type: object - required: - - messages - properties: - internal_state: - type: - - string - - 'null' - messages: - type: array - items: - $ref: '#/components/schemas/ChatMessage' ClearlyDefinedCurationImporter: allOf: - $ref: '#/components/schemas/CommonImporter' @@ -2593,42 +2314,6 @@ components: period: type: string description: The period the importer should be run. - Conversation: - type: object - required: - - id - - messages - - updated_at - - seq - properties: - id: - type: string - format: uuid - messages: - type: array - items: - $ref: '#/components/schemas/ChatMessage' - seq: - type: integer - format: int32 - updated_at: - type: string - format: date-time - ConversationSummary: - type: object - required: - - id - - updated_at - - summary - properties: - id: - type: string - format: uuid - summary: - type: string - updated_at: - type: string - format: date-time Cpe: type: string format: uri @@ -2921,13 +2606,6 @@ components: severity: $ref: '#/components/schemas/Severity' description: The severity of the message - MessageType: - type: string - enum: - - human - - system - - ai - - tool OrganizationDetails: allOf: - $ref: '#/components/schemas/OrganizationHead' @@ -3133,33 +2811,6 @@ components: type: integer format: int64 minimum: 0 - PaginatedResults_ConversationSummary: - type: object - required: - - items - - total - properties: - items: - type: array - items: - type: object - required: - - id - - updated_at - - summary - properties: - id: - type: string - format: uuid - summary: - type: string - updated_at: - type: string - format: date-time - total: - type: integer - format: int64 - minimum: 0 PaginatedResults_DepSummary: type: object required: diff --git a/server/Cargo.toml b/server/Cargo.toml index ebf84027c..51f60b695 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true publish.workspace = true license.workspace = true +[features] +ai = ["trustify-module-fundamental/ai"] + [dependencies] trustify-auth = { workspace = true } trustify-common = { workspace = true } diff --git a/trustd/Cargo.toml b/trustd/Cargo.toml index 2e60164d7..9257e27b3 100644 --- a/trustd/Cargo.toml +++ b/trustd/Cargo.toml @@ -33,6 +33,7 @@ default = ["pm"] bundled = ["postgresql_embedded/bundled"] garage-door = ["trustify-server/garage-door"] +ai = ["trustify-server/ai"] vendored = [ "openssl/vendored", diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index ab18ec27c..e002260bc 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -5,6 +5,10 @@ edition.workspace = true publish.workspace = true license.workspace = true +[features] +default = [] +ai = [ "trustify-common/ai", "trustify-module-fundamental/ai" ] + [dependencies] anyhow = { workspace = true } async-trait = { workspace = true } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 2beb730c0..9e2199c8b 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -4,6 +4,7 @@ use crate::log::init_log; use clap::{Parser, Subcommand}; +#[cfg(feature = "ai")] mod ai; mod dataset; mod log; @@ -24,6 +25,7 @@ impl Xtask { Command::GenerateDump(command) => command.run().await, Command::GenerateSchemas(command) => command.run(), Command::Precommit(command) => command.run().await, + #[cfg(feature = "ai")] Command::Ai(command) => command.run().await, } } @@ -40,6 +42,7 @@ pub enum Command { /// Run precommit checks Precommit(precommit::Precommit), /// Run ai tool + #[cfg(feature = "ai")] Ai(ai::Ai), }