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

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

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))
+    }
+}