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

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

compliance-agent/src/database.rs

@@ -127,11 +127,7 @@ impl Database {
 
         // dast_targets: index on repo_id
         self.dast_targets()
-            .create_index(
-                IndexModel::builder()
-                    .keys(doc! { "repo_id": 1 })
-                    .build(),
-            )
+            .create_index(IndexModel::builder().keys(doc! { "repo_id": 1 }).build())
             .await?;
 
         // dast_scan_runs: compound (target_id, started_at DESC)
@@ -152,6 +148,24 @@ impl Database {
             )
             .await?;
 
+        // code_embeddings: compound (repo_id, graph_build_id)
+        self.code_embeddings()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "repo_id": 1, "graph_build_id": 1 })
+                    .build(),
+            )
+            .await?;
+
+        // embedding_builds: compound (repo_id, started_at DESC)
+        self.embedding_builds()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "repo_id": 1, "started_at": -1 })
+                    .build(),
+            )
+            .await?;
+
         tracing::info!("Database indexes ensured");
         Ok(())
     }
@@ -210,6 +224,17 @@ impl Database {
         self.inner.collection("dast_findings")
     }
 
+    // Embedding collections
+    pub fn code_embeddings(&self) -> Collection<compliance_core::models::embedding::CodeEmbedding> {
+        self.inner.collection("code_embeddings")
+    }
+
+    pub fn embedding_builds(
+        &self,
+    ) -> Collection<compliance_core::models::embedding::EmbeddingBuildRun> {
+        self.inner.collection("embedding_builds")
+    }
+
     #[allow(dead_code)]
     pub fn raw_collection(&self, name: &str) -> Collection<mongodb::bson::Document> {
         self.inner.collection(name)