Sharang Parnerkar
42cabf0582
Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com> Reviewed-on: #1
101 lines
3.0 KiB
Rust
101 lines
3.0 KiB
Rust
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))
|
|
}
|
|
}
|