feat: bump version to 0.1.6 and add new features
- Updated Cargo.toml to version 0.1.6 and added new features: `http`, `sessions-db`, `ai`, and `broker-nats`. - Introduced `FastEmbedAdapter` for generating text embeddings using the `fastembed` crate. - Added `QdrantAdapter` for interacting with Qdrant vector database. - Implemented a message broker trait and a NATS broker for asynchronous message publishing and subscribing. - Created HTTP server middleware for CORS and session management. - Added session store implementation with support for both PostgreSQL and SQLite. - Organized module structure by creating new modules for AI, broker, and HTTP functionalities.
This commit is contained in:
2841
Cargo.lock
generated
2841
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
29
Cargo.toml
29
Cargo.toml
@@ -1,15 +1,20 @@
|
||||
[package]
|
||||
name = "k-core"
|
||||
version = "0.1.5"
|
||||
version = "0.1.6"
|
||||
edition = "2024"
|
||||
|
||||
[features]
|
||||
default = ["logging", "db-sqlx", "auth"]
|
||||
default = ["logging", "db-sqlx", "auth", "http"]
|
||||
logging = ["dep:tracing", "dep:tracing-subscriber"]
|
||||
db-sqlx = ["dep:sqlx"]
|
||||
postgres = ["db-sqlx", "sqlx/postgres"]
|
||||
sqlite = ["db-sqlx", "sqlx/sqlite"]
|
||||
postgres = ["db-sqlx", "sqlx/postgres", "tower-sessions-sqlx-store?/postgres"]
|
||||
sqlite = ["db-sqlx", "sqlx/sqlite", "tower-sessions-sqlx-store?/sqlite"]
|
||||
auth = ["dep:tower-sessions"]
|
||||
sessions-db = ["auth", "dep:tower-sessions-sqlx-store"]
|
||||
ai = ["dep:fastembed", "dep:qdrant-client"]
|
||||
broker = []
|
||||
broker-nats = ["broker", "dep:async-nats"]
|
||||
http = ["dep:axum", "dep:tower", "dep:tower-http", "dep:time", "logging"]
|
||||
|
||||
[dependencies]
|
||||
# Error handling
|
||||
@@ -26,7 +31,6 @@ sqlx = { version = "0.8.6", features = [
|
||||
"chrono",
|
||||
"uuid",
|
||||
], optional = true }
|
||||
sqlx-core = { version = "0.8.6", optional = true }
|
||||
|
||||
# Logging
|
||||
tracing = { version = "0.1", optional = true }
|
||||
@@ -37,8 +41,23 @@ tracing-subscriber = { version = "0.3.22", features = [
|
||||
|
||||
# Auth
|
||||
tower-sessions = { version = "0.14.0", optional = true }
|
||||
tower-sessions-sqlx-store = { version = "0.15", optional = true}
|
||||
|
||||
# Utils
|
||||
chrono = { version = "0.4.42", features = ["serde"] }
|
||||
uuid = { version = "1.19.0", features = ["v4", "serde"] }
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
async-trait = "0.1.89"
|
||||
|
||||
# AI
|
||||
fastembed = { version = "5.4", optional = true }
|
||||
qdrant-client = { version = "1.16", optional = true }
|
||||
|
||||
# Broker
|
||||
async-nats = { version = "0.45", optional = true }
|
||||
|
||||
# HTTP
|
||||
axum = { version = "0.8.8", features = ["macros"], optional = true }
|
||||
tower = { version = "0.5.2", optional = true }
|
||||
tower-http = { version = "0.6.2", features = ["cors", "trace"], optional = true}
|
||||
time = { version = "0.3", optional = true }
|
||||
69
src/ai/embeddings.rs
Normal file
69
src/ai/embeddings.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FastEmbedAdapter {
|
||||
model: Arc<Mutex<TextEmbedding>>,
|
||||
}
|
||||
|
||||
impl FastEmbedAdapter {
|
||||
pub fn new() -> DomainResult<Self> {
|
||||
let mut options = InitOptions::default();
|
||||
options.model_name = EmbeddingModel::AllMiniLML6V2;
|
||||
options.show_download_progress = false;
|
||||
|
||||
let model = TextEmbedding::try_new(options)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to init fastembed: {}", e))?;
|
||||
|
||||
Ok(Self {
|
||||
model: Arc::new(Mutex::new(model)),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn generate_embedding(&self, text: &str) -> anyhow::Result<Vec<f32>> {
|
||||
let model = self.model.clone();
|
||||
let text = text.to_string();
|
||||
|
||||
// FastEmbed is blocking, so we run it in a blocking task if we are in an async context,
|
||||
// but since this method signature doesn't force async, we wrap the internal logic.
|
||||
// For strictly async usage in k-core:
|
||||
let embeddings = std::thread::scope(|s| {
|
||||
s.spawn(|| {
|
||||
let mut model = model
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {}", e))?;
|
||||
model
|
||||
.embed(vec![text], None)
|
||||
.map_err(|e| anyhow::anyhow!("Embed error: {}", e))
|
||||
})
|
||||
.join()
|
||||
.map_err(|_| anyhow::anyhow!("Thread join error"))?
|
||||
})?;
|
||||
|
||||
embeddings
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("No embedding generated"))
|
||||
}
|
||||
|
||||
/// Async wrapper for use in async contexts (like Axum handlers)
|
||||
pub async fn generate_embedding_async(&self, text: &str) -> anyhow::Result<Vec<f32>> {
|
||||
let model = self.model.clone();
|
||||
let text = text.to_string();
|
||||
|
||||
let embeddings = tokio::task::spawn_blocking(move || {
|
||||
let mut model = model
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("Lock error: {}", e))?;
|
||||
model
|
||||
.embed(vec![text], None)
|
||||
.map_err(|e| anyhow::anyhow!("Embed error: {}", e))
|
||||
})
|
||||
.await??;
|
||||
|
||||
embeddings
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("No embedding generated"))
|
||||
}
|
||||
}
|
||||
2
src/ai/mod.rs
Normal file
2
src/ai/mod.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
pub mod embeddings;
|
||||
pub mod qdrant;
|
||||
72
src/ai/qdrant.rs
Normal file
72
src/ai/qdrant.rs
Normal file
@@ -0,0 +1,72 @@
|
||||
use qdrant_client::Qdrant;
|
||||
use qdrant_client::qdrant::{
|
||||
CreateCollectionBuilder, Distance, PointStruct, SearchPointsBuilder, UpsertPointsBuilder,
|
||||
Value, VectorParamsBuilder,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QdrantAdapter {
|
||||
client: Arc<Qdrant>,
|
||||
collection_name: String,
|
||||
}
|
||||
|
||||
impl QdrantAdapter {
|
||||
pub fn new(url: &str, collection_name: &str) -> anyhow::Result<Self> {
|
||||
let client = Qdrant::from_url(url).build()?;
|
||||
Ok(Self {
|
||||
client: Arc::new(client),
|
||||
collection_name: collection_name.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn create_collection_if_not_exists(&self, vector_size: u64) -> anyhow::Result<()> {
|
||||
if !self.client.collection_exists(&self.collection_name).await? {
|
||||
self.client
|
||||
.create_collection(
|
||||
CreateCollectionBuilder::new(self.collection_name.clone())
|
||||
.vectors_config(VectorParamsBuilder::new(vector_size, Distance::Cosine)),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
id: Uuid,
|
||||
vector: Vec<f32>,
|
||||
payload: HashMap<String, Value>,
|
||||
) -> anyhow::Result<()> {
|
||||
let point = PointStruct::new(id.to_string(), vector, payload);
|
||||
let upsert_points = UpsertPointsBuilder::new(self.collection_name.clone(), vec![point]);
|
||||
|
||||
self.client.upsert_points(upsert_points).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn search(&self, vector: Vec<f32>, limit: u64) -> anyhow::Result<Vec<(Uuid, f32)>> {
|
||||
let search_points = SearchPointsBuilder::new(self.collection_name.clone(), vector, limit)
|
||||
.with_payload(true);
|
||||
|
||||
let search_result = self.client.search_points(search_points).await?;
|
||||
|
||||
let results = search_result
|
||||
.result
|
||||
.into_iter()
|
||||
.filter_map(|point| {
|
||||
let id = point.id?;
|
||||
let uuid_str = match id.point_id_options? {
|
||||
qdrant_client::qdrant::point_id::PointIdOptions::Uuid(u) => u,
|
||||
_ => return None,
|
||||
};
|
||||
let uuid = Uuid::parse_str(&uuid_str).ok()?;
|
||||
Some((uuid, point.score))
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
16
src/broker/mod.rs
Normal file
16
src/broker/mod.rs
Normal file
@@ -0,0 +1,16 @@
|
||||
use async_trait::async_trait;
|
||||
use futures_core::Stream;
|
||||
use std::pin::Pin;
|
||||
|
||||
#[cfg(feature = "broker-nats")]
|
||||
pub mod nats;
|
||||
|
||||
#[async_trait]
|
||||
pub trait MessageBroker: Send + Sync {
|
||||
async fn publish(&self, topic: &str, payload: Vec<u8>) -> anyhow::Result<()>;
|
||||
|
||||
async fn subscribe(
|
||||
&self,
|
||||
topic: &str,
|
||||
) -> anyhow::Result<Pin<Box<dyn Stream<Item = Vec<u8>> + Send>>>;
|
||||
}
|
||||
43
src/broker/nats.rs
Normal file
43
src/broker/nats.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use super::MessageBroker;
|
||||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use std::pin::Pin;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct NatsBroker {
|
||||
client: async_nats::Client,
|
||||
}
|
||||
|
||||
impl NatsBroker {
|
||||
pub async fn connect(url: &str) -> Result<Self, async_nats::ConnectError> {
|
||||
let client = async_nats::connect(url).await?;
|
||||
Ok(Self { client })
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MessageBroker for NatsBroker {
|
||||
async fn publish(&self, topic: &str, payload: Vec<u8>) -> anyhow::Result<()> {
|
||||
self.client
|
||||
.publish(topic.to_string(), payload.into())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("NATS publish error: {}", e))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn subscribe(
|
||||
&self,
|
||||
topic: &str,
|
||||
) -> anyhow::Result<Pin<Box<dyn futures_core::Stream<Item = Vec<u8>> + Send>>> {
|
||||
let subscriber = self
|
||||
.client
|
||||
.subscribe(topic.to_string())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("NATS subscribe error: {}", e))?;
|
||||
|
||||
// Map NATS Message to generic Vec<u8>
|
||||
let stream = subscriber.map(|msg| msg.payload.to_vec());
|
||||
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,3 @@
|
||||
use sqlx::Pool;
|
||||
use std::time::Duration;
|
||||
|
||||
#[cfg(feature = "sqlite")]
|
||||
|
||||
1
src/http/mod.rs
Normal file
1
src/http/mod.rs
Normal file
@@ -0,0 +1 @@
|
||||
pub mod server;
|
||||
52
src/http/server.rs
Normal file
52
src/http/server.rs
Normal file
@@ -0,0 +1,52 @@
|
||||
use axum::Router;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::trace::TraceLayer;
|
||||
use tower_sessions::{Expiry, SessionManagerLayer, SessionStore};
|
||||
|
||||
pub struct ServerConfig {
|
||||
pub cors_origins: Vec<String>,
|
||||
pub session_secret: Option<String>,
|
||||
}
|
||||
|
||||
pub fn apply_standard_middleware(app: Router, config: &ServerConfig) -> Router {
|
||||
let mut cors = CorsLayer::new()
|
||||
.allow_methods([
|
||||
axum::http::Method::GET,
|
||||
axum::http::Method::POST,
|
||||
axum::http::Method::PATCH,
|
||||
axum::http::Method::DELETE,
|
||||
axum::http::Method::OPTIONS,
|
||||
])
|
||||
.allow_headers([
|
||||
axum::http::header::AUTHORIZATION,
|
||||
axum::http::header::ACCEPT,
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
])
|
||||
.allow_credentials(true);
|
||||
|
||||
// Configure CORS origins
|
||||
let mut allowed_origins = Vec::new();
|
||||
for origin in &config.cors_origins {
|
||||
if let Ok(value) = origin.parse::<axum::http::HeaderValue>() {
|
||||
allowed_origins.push(value);
|
||||
}
|
||||
}
|
||||
|
||||
if !allowed_origins.is_empty() {
|
||||
cors = cors.allow_origin(allowed_origins);
|
||||
}
|
||||
|
||||
app.layer(cors).layer(TraceLayer::new_for_http())
|
||||
}
|
||||
|
||||
/// Helper to attach a session layer with standard K-Suite configuration
|
||||
pub fn attach_session_layer<S>(app: Router, store: S) -> Router
|
||||
where
|
||||
S: SessionStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
let session_layer = SessionManagerLayer::new(store)
|
||||
.with_secure(false) // Set to true if you handle HTTPS termination in the app
|
||||
.with_expiry(Expiry::OnInactivity(time::Duration::days(7)));
|
||||
|
||||
app.layer(session_layer)
|
||||
}
|
||||
11
src/lib.rs
11
src/lib.rs
@@ -5,3 +5,14 @@ pub mod logging;
|
||||
pub mod db;
|
||||
|
||||
pub mod error;
|
||||
|
||||
#[cfg(feature = "ai")]
|
||||
pub mod ai;
|
||||
|
||||
#[cfg(feature = "broker")]
|
||||
pub mod broker;
|
||||
#[cfg(feature = "auth")]
|
||||
pub mod session;
|
||||
|
||||
#[cfg(feature = "http")]
|
||||
pub mod http;
|
||||
|
||||
79
src/session.rs
Normal file
79
src/session.rs
Normal file
@@ -0,0 +1,79 @@
|
||||
// Only compile this module if we have the store dependency AND sqlx enabled
|
||||
#[cfg(all(feature = "sessions-db", feature = "db-sqlx"))]
|
||||
pub mod store {
|
||||
use async_trait::async_trait;
|
||||
use tower_sessions::{
|
||||
SessionStore,
|
||||
session::{Id, Record},
|
||||
session_store,
|
||||
};
|
||||
|
||||
// We only import what is actually enabled by features
|
||||
#[cfg(feature = "postgres")]
|
||||
use tower_sessions_sqlx_store::PostgresStore;
|
||||
#[cfg(feature = "sqlite")]
|
||||
use tower_sessions_sqlx_store::SqliteStore;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum InfraSessionStore {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Sqlite(SqliteStore),
|
||||
#[cfg(feature = "postgres")]
|
||||
Postgres(PostgresStore),
|
||||
// Fallback variant to allow compilation if sessions-db is enabled but no generic DB feature is selected
|
||||
// (Though in practice you should ensure your config validates this)
|
||||
#[allow(dead_code)]
|
||||
Unused,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SessionStore for InfraSessionStore {
|
||||
async fn save(&self, session_record: &Record) -> session_store::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(store) => store.save(session_record).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(store) => store.save(session_record).await,
|
||||
_ => Err(session_store::Error::Backend(
|
||||
"No database feature enabled".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn load(&self, session_id: &Id) -> session_store::Result<Option<Record>> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(store) => store.load(session_id).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(store) => store.load(session_id).await,
|
||||
_ => Err(session_store::Error::Backend(
|
||||
"No database feature enabled".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete(&self, session_id: &Id) -> session_store::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(store) => store.delete(session_id).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(store) => store.delete(session_id).await,
|
||||
_ => Err(session_store::Error::Backend(
|
||||
"No database feature enabled".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl InfraSessionStore {
|
||||
pub async fn migrate(&self) -> anyhow::Result<()> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(store) => store.migrate().await.map_err(|e| anyhow::anyhow!(e)),
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(store) => store.migrate().await.map_err(|e| anyhow::anyhow!(e)),
|
||||
_ => Ok(()), // No migration needed for no-op
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user