feat(core): implement full memory pipeline (#11)
This commit was merged in pull request #11.
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
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(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user