feat: rag-embedding-ai-chat (#1)

Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #1

compliance-core/src/models/chat.rs

@@ -0,0 +1,35 @@
+use serde::{Deserialize, Serialize};
+
+/// A message in the chat history
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatMessage {
+    pub role: String,
+    pub content: String,
+}
+
+/// Request body for the chat endpoint
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatRequest {
+    pub message: String,
+    #[serde(default)]
+    pub history: Vec<ChatMessage>,
+}
+
+/// A source reference from the RAG retrieval
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SourceReference {
+    pub file_path: String,
+    pub qualified_name: String,
+    pub start_line: u32,
+    pub end_line: u32,
+    pub language: String,
+    pub snippet: String,
+    pub score: f64,
+}
+
+/// Response from the chat endpoint
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatResponse {
+    pub message: String,
+    pub sources: Vec<SourceReference>,
+}

compliance-core/src/models/dast.rs

@@ -244,6 +244,7 @@ pub struct DastFinding {
 }
 
 impl DastFinding {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         scan_run_id: String,
         target_id: String,

compliance-core/src/models/embedding.rs

@@ -0,0 +1,100 @@
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+
+/// Status of an embedding build operation
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum EmbeddingBuildStatus {
+    Running,
+    Completed,
+    Failed,
+}
+
+/// A code embedding stored in MongoDB Atlas Vector Search
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CodeEmbedding {
+    #[serde(rename = "_id", skip_serializing_if = "Option::is_none")]
+    pub id: Option<bson::oid::ObjectId>,
+    pub repo_id: String,
+    pub graph_build_id: String,
+    pub qualified_name: String,
+    pub kind: String,
+    pub file_path: String,
+    pub start_line: u32,
+    pub end_line: u32,
+    pub language: String,
+    pub content: String,
+    pub context_header: String,
+    pub embedding: Vec<f64>,
+    pub token_estimate: u32,
+    #[serde(with = "bson::serde_helpers::chrono_datetime_as_bson_datetime")]
+    pub created_at: DateTime<Utc>,
+}
+
+/// Tracks an embedding build operation for a repository
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EmbeddingBuildRun {
+    #[serde(rename = "_id", skip_serializing_if = "Option::is_none")]
+    pub id: Option<bson::oid::ObjectId>,
+    pub repo_id: String,
+    pub graph_build_id: String,
+    pub status: EmbeddingBuildStatus,
+    pub total_chunks: u32,
+    pub embedded_chunks: u32,
+    pub embedding_model: String,
+    pub error_message: Option<String>,
+    #[serde(with = "bson::serde_helpers::chrono_datetime_as_bson_datetime")]
+    pub started_at: DateTime<Utc>,
+    #[serde(
+        default,
+        skip_serializing_if = "Option::is_none",
+        with = "opt_chrono_as_bson"
+    )]
+    pub completed_at: Option<DateTime<Utc>>,
+}
+
+impl EmbeddingBuildRun {
+    pub fn new(repo_id: String, graph_build_id: String, embedding_model: String) -> Self {
+        Self {
+            id: None,
+            repo_id,
+            graph_build_id,
+            status: EmbeddingBuildStatus::Running,
+            total_chunks: 0,
+            embedded_chunks: 0,
+            embedding_model,
+            error_message: None,
+            started_at: Utc::now(),
+            completed_at: None,
+        }
+    }
+}
+
+/// Serde helper for Option<DateTime<Utc>> as BSON DateTime
+mod opt_chrono_as_bson {
+    use chrono::{DateTime, Utc};
+    use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+    #[derive(Serialize, Deserialize)]
+    struct BsonDt(
+        #[serde(with = "bson::serde_helpers::chrono_datetime_as_bson_datetime")] DateTime<Utc>,
+    );
+
+    pub fn serialize<S>(value: &Option<DateTime<Utc>>, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match value {
+            Some(dt) => BsonDt(*dt).serialize(serializer),
+            None => serializer.serialize_none(),
+        }
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<DateTime<Utc>>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let opt: Option<BsonDt> = Option::deserialize(deserializer)?;
+        Ok(opt.map(|d| d.0))
+    }
+}

compliance-core/src/models/mod.rs

@@ -1,5 +1,7 @@
+pub mod chat;
 pub mod cve;
 pub mod dast;
+pub mod embedding;
 pub mod finding;
 pub mod graph;
 pub mod issue;
@@ -7,15 +9,16 @@ pub mod repository;
 pub mod sbom;
 pub mod scan;
 
+pub use chat::{ChatMessage, ChatRequest, ChatResponse, SourceReference};
 pub use cve::{CveAlert, CveSource};
 pub use dast::{
     DastAuthConfig, DastEvidence, DastFinding, DastScanPhase, DastScanRun, DastScanStatus,
     DastTarget, DastTargetType, DastVulnType,
 };
+pub use embedding::{CodeEmbedding, EmbeddingBuildRun, EmbeddingBuildStatus};
 pub use finding::{Finding, FindingStatus, Severity};
 pub use graph::{
-    CodeEdge, CodeEdgeKind, CodeNode, CodeNodeKind, GraphBuildRun, GraphBuildStatus,
-    ImpactAnalysis,
+    CodeEdge, CodeEdgeKind, CodeNode, CodeNodeKind, GraphBuildRun, GraphBuildStatus, ImpactAnalysis,
 };
 pub use issue::{IssueStatus, TrackerIssue, TrackerType};
 pub use repository::{ScanTrigger, TrackedRepository};

compliance-core/src/models/repository.rs

@@ -31,9 +31,15 @@ pub struct TrackedRepository {
     pub last_scanned_commit: Option<String>,
     #[serde(default, deserialize_with = "deserialize_findings_count")]
     pub findings_count: u32,
-    #[serde(default = "chrono::Utc::now", deserialize_with = "deserialize_datetime")]
+    #[serde(
+        default = "chrono::Utc::now",
+        deserialize_with = "deserialize_datetime"
+    )]
     pub created_at: DateTime<Utc>,
-    #[serde(default = "chrono::Utc::now", deserialize_with = "deserialize_datetime")]
+    #[serde(
+        default = "chrono::Utc::now",
+        deserialize_with = "deserialize_datetime"
+    )]
     pub updated_at: DateTime<Utc>,
 }
 
@@ -51,9 +57,7 @@ where
     let bson = bson::Bson::deserialize(deserializer)?;
     match bson {
         bson::Bson::DateTime(dt) => Ok(dt.into()),
-        bson::Bson::String(s) => s
-            .parse::<DateTime<Utc>>()
-            .map_err(serde::de::Error::custom),
+        bson::Bson::String(s) => s.parse::<DateTime<Utc>>().map_err(serde::de::Error::custom),
         other => Err(serde::de::Error::custom(format!(
             "expected DateTime or string, got: {other:?}"
         ))),