use stateset_core::{
Customer, EmbeddingService, EmbeddingStats, EntityType, InventoryItem, Order, Product, Result,
VectorRepository, VectorSearchResult,
};
use stateset_db::sqlite::SqliteVectorRepository;
use std::sync::Arc;
#[derive(Debug)]
pub struct Vector {
repo: SqliteVectorRepository,
embedding_service: Arc<EmbeddingService>,
}
impl Vector {
pub(crate) fn new(repo: SqliteVectorRepository, api_key: String) -> Self {
Self { repo, embedding_service: Arc::new(EmbeddingService::new(api_key)) }
}
pub fn with_model(
repo: SqliteVectorRepository,
api_key: String,
model: String,
dimensions: usize,
) -> Self {
Self {
repo,
embedding_service: Arc::new(EmbeddingService::with_model(api_key, model, dimensions)),
}
}
pub fn index_product(&self, product: &Product) -> Result<()> {
let text = EmbeddingService::product_text(product);
let result = self.embedding_service.embed(&text)?;
self.repo.store_embedding(
EntityType::Product,
&product.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)
}
pub fn search_products(
&self,
query: &str,
limit: usize,
) -> Result<Vec<VectorSearchResult<Product>>> {
let result = self.embedding_service.embed(query)?;
self.repo.search_products_hybrid(&result.embedding, query, limit)
}
pub fn search_products_by_embedding(
&self,
embedding: &[f32],
limit: usize,
) -> Result<Vec<VectorSearchResult<Product>>> {
self.repo.search_products(embedding, limit)
}
pub fn unindex_product(&self, product_id: &str) -> Result<()> {
self.repo.delete_embedding(EntityType::Product, product_id)
}
pub fn index_products(&self, products: &[Product]) -> Result<usize> {
let mut indexed = 0;
for chunk in products.chunks(100) {
let texts: Vec<String> = chunk.iter().map(EmbeddingService::product_text).collect();
let results = self.embedding_service.embed_batch(&texts)?;
for (product, result) in chunk.iter().zip(results.iter()) {
self.repo.store_embedding(
EntityType::Product,
&product.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)?;
indexed += 1;
}
}
Ok(indexed)
}
pub fn index_customer(&self, customer: &Customer) -> Result<()> {
let text = EmbeddingService::customer_text(customer);
let result = self.embedding_service.embed(&text)?;
self.repo.store_embedding(
EntityType::Customer,
&customer.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)
}
pub fn search_customers(
&self,
query: &str,
limit: usize,
) -> Result<Vec<VectorSearchResult<Customer>>> {
let result = self.embedding_service.embed(query)?;
self.repo.search_customers_hybrid(&result.embedding, query, limit)
}
pub fn unindex_customer(&self, customer_id: &str) -> Result<()> {
self.repo.delete_embedding(EntityType::Customer, customer_id)
}
pub fn index_customers(&self, customers: &[Customer]) -> Result<usize> {
let mut indexed = 0;
for chunk in customers.chunks(100) {
let texts: Vec<String> = chunk.iter().map(EmbeddingService::customer_text).collect();
let results = self.embedding_service.embed_batch(&texts)?;
for (customer, result) in chunk.iter().zip(results.iter()) {
self.repo.store_embedding(
EntityType::Customer,
&customer.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)?;
indexed += 1;
}
}
Ok(indexed)
}
pub fn index_order(&self, order: &Order) -> Result<()> {
let text = EmbeddingService::order_text(order);
let result = self.embedding_service.embed(&text)?;
self.repo.store_embedding(
EntityType::Order,
&order.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)
}
pub fn search_orders(
&self,
query: &str,
limit: usize,
) -> Result<Vec<VectorSearchResult<Order>>> {
let result = self.embedding_service.embed(query)?;
self.repo.search_orders_hybrid(&result.embedding, query, limit)
}
pub fn unindex_order(&self, order_id: &str) -> Result<()> {
self.repo.delete_embedding(EntityType::Order, order_id)
}
pub fn index_orders(&self, orders: &[Order]) -> Result<usize> {
let mut indexed = 0;
for chunk in orders.chunks(100) {
let texts: Vec<String> = chunk.iter().map(EmbeddingService::order_text).collect();
let results = self.embedding_service.embed_batch(&texts)?;
for (order, result) in chunk.iter().zip(results.iter()) {
self.repo.store_embedding(
EntityType::Order,
&order.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)?;
indexed += 1;
}
}
Ok(indexed)
}
pub fn index_inventory_item(&self, item: &InventoryItem) -> Result<()> {
let text = EmbeddingService::inventory_item_text(item);
let result = self.embedding_service.embed(&text)?;
self.repo.store_embedding(
EntityType::InventoryItem,
&item.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)
}
pub fn search_inventory(
&self,
query: &str,
limit: usize,
) -> Result<Vec<VectorSearchResult<InventoryItem>>> {
let result = self.embedding_service.embed(query)?;
self.repo.search_inventory_hybrid(&result.embedding, query, limit)
}
pub fn unindex_inventory_item(&self, item_id: &str) -> Result<()> {
self.repo.delete_embedding(EntityType::InventoryItem, item_id)
}
pub fn index_inventory_items(&self, items: &[InventoryItem]) -> Result<usize> {
let mut indexed = 0;
for chunk in items.chunks(100) {
let texts: Vec<String> =
chunk.iter().map(EmbeddingService::inventory_item_text).collect();
let results = self.embedding_service.embed_batch(&texts)?;
for (item, result) in chunk.iter().zip(results.iter()) {
self.repo.store_embedding(
EntityType::InventoryItem,
&item.id.to_string(),
&result.embedding,
&result.text_hash,
self.embedding_service.model(),
)?;
indexed += 1;
}
}
Ok(indexed)
}
pub fn stats(&self) -> Result<EmbeddingStats> {
self.repo.get_stats()
}
pub fn is_indexed(&self, entity_type: EntityType, entity_id: &str) -> Result<bool> {
self.repo.has_embedding(entity_type, entity_id)
}
pub fn clear(&self, entity_type: EntityType) -> Result<u64> {
self.repo.clear_embeddings(entity_type)
}
pub fn clear_all(&self) -> Result<u64> {
let mut total = 0;
total += self.repo.clear_embeddings(EntityType::Product)?;
total += self.repo.clear_embeddings(EntityType::Customer)?;
total += self.repo.clear_embeddings(EntityType::Order)?;
total += self.repo.clear_embeddings(EntityType::InventoryItem)?;
Ok(total)
}
pub fn embed(&self, text: &str) -> Result<Vec<f32>> {
let result = self.embedding_service.embed(text)?;
Ok(result.embedding)
}
}