112 lines
3.6 KiB
Rust
112 lines
3.6 KiB
Rust
use anyhow::Result;
|
|
use mem_store::{MemoryL1, VectorStore, ChunkL0};
|
|
use mem_llm::EmbeddingsClient;
|
|
use sqlx::PgPool;
|
|
use uuid::Uuid;
|
|
use std::sync::Arc;
|
|
use pgvector::Vector;
|
|
|
|
/// Ingest worker — processes queued records through memory storage
|
|
pub struct IngestWorker {
|
|
pool: PgPool,
|
|
vector_store: Arc<VectorStore>,
|
|
embeddings: Arc<EmbeddingsClient>,
|
|
}
|
|
|
|
impl IngestWorker {
|
|
/// Create worker
|
|
pub fn new(
|
|
pool: PgPool,
|
|
embeddings: EmbeddingsClient,
|
|
) -> Self {
|
|
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
|
Self {
|
|
pool,
|
|
vector_store,
|
|
embeddings: Arc::new(embeddings),
|
|
}
|
|
}
|
|
|
|
/// Process ingest job: records -> chunks -> storage
|
|
pub async fn process_ingest(
|
|
&self,
|
|
project: &str,
|
|
ingest_id: &str,
|
|
records: Vec<(String, String)>, // (content, source)
|
|
) -> Result<()> {
|
|
tracing::info!("Processing ingest: project={}, id={}, records={}", project, ingest_id, records.len());
|
|
|
|
// Update job status to processing
|
|
sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2")
|
|
.bind("processing")
|
|
.bind(ingest_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
let mut total_chunks = 0;
|
|
let mut total_stored = 0;
|
|
|
|
// Process each record
|
|
for (content, source) in &records {
|
|
let chunk_id = Uuid::new_v4();
|
|
|
|
// Store L0 chunk
|
|
let l0_chunk = ChunkL0 {
|
|
id: chunk_id,
|
|
project: project.to_string(),
|
|
query_id: "ingest".to_string(),
|
|
source: source.clone(),
|
|
content: content.clone(),
|
|
tokens: (content.len() / 4) as i32,
|
|
};
|
|
self.vector_store.store_chunk_l0(&l0_chunk).await?;
|
|
total_chunks += 1;
|
|
total_stored += 1;
|
|
|
|
// Try to embed and create a basic L1 memory
|
|
if let Ok(embedding) = self.embeddings.embed(content).await {
|
|
let l1 = MemoryL1 {
|
|
id: Uuid::new_v4(),
|
|
project: project.to_string(),
|
|
query_id: "ingest".to_string(),
|
|
content: content.clone(),
|
|
tokens: (content.len() / 4) as i32,
|
|
embedding: Some(embedding.to_vec()),
|
|
chunks_seen: 1,
|
|
chunks_used: 1,
|
|
run_id: ingest_id.to_string(),
|
|
};
|
|
|
|
if let Err(e) = self.vector_store.store_memory_l1(&l1, &embedding).await {
|
|
tracing::warn!("Failed to store L1 memory: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Mark job complete
|
|
sqlx::query("UPDATE ingest_jobs SET status=$1, completed_at=NOW() WHERE ingest_id=$2")
|
|
.bind("done")
|
|
.bind(ingest_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
tracing::info!("Ingest completed: {} (stored {} chunks)", ingest_id, total_stored);
|
|
Ok(())
|
|
}
|
|
|
|
/// Process a single chunk
|
|
pub async fn process_chunk(&self, project: &str, query_id: &str, content: &str, source: &str) -> Result<()> {
|
|
let embedding = self.embeddings.embed(content).await?;
|
|
let chunk = ChunkL0 {
|
|
id: Uuid::new_v4(),
|
|
project: project.to_string(),
|
|
query_id: query_id.to_string(),
|
|
source: source.to_string(),
|
|
content: content.to_string(),
|
|
tokens: (content.len() / 4) as i32,
|
|
};
|
|
self.vector_store.store_chunk_l0(&chunk).await?;
|
|
Ok(())
|
|
}
|
|
}
|