d45e9e28f4
Extract agent, compact, embed, task, and modes modules from single service.rs files into focused sub-modules. Add orao module for O1-like reasoning loop. Move RigAgentService to rig_tool.rs.
91 lines
2.4 KiB
Rust
91 lines
2.4 KiB
Rust
pub mod chunk;
|
|
pub mod client;
|
|
pub mod embeddable;
|
|
pub mod entity_embed;
|
|
pub mod qdrant;
|
|
pub mod search;
|
|
|
|
pub use client::{EmbedClient, EmbedPayload, EmbedVector, SearchResult};
|
|
pub use embeddable::{EmbedMemoryInput, Embeddable, TagEmbedInput};
|
|
pub use qdrant::QdrantClient;
|
|
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Clone)]
|
|
pub struct EmbedService {
|
|
client: Arc<EmbedClient>,
|
|
db: sea_orm::DatabaseConnection,
|
|
model_name: String,
|
|
dimensions: u64,
|
|
}
|
|
|
|
impl EmbedService {
|
|
pub fn new(
|
|
client: EmbedClient,
|
|
db: sea_orm::DatabaseConnection,
|
|
model_name: String,
|
|
dimensions: u64,
|
|
) -> Self {
|
|
Self {
|
|
client: Arc::new(client),
|
|
db,
|
|
model_name,
|
|
dimensions,
|
|
}
|
|
}
|
|
|
|
pub async fn ensure_collections(&self) -> crate::Result<()> {
|
|
self.client
|
|
.ensure_collection("issue", self.dimensions)
|
|
.await?;
|
|
self.client
|
|
.ensure_collection("repo", self.dimensions)
|
|
.await?;
|
|
self.client.ensure_skill_collection(self.dimensions).await?;
|
|
self.client
|
|
.ensure_collection("repo_tag", self.dimensions)
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn db(&self) -> &sea_orm::DatabaseConnection {
|
|
&self.db
|
|
}
|
|
|
|
pub fn client(&self) -> &Arc<EmbedClient> {
|
|
&self.client
|
|
}
|
|
|
|
pub fn model_name(&self) -> &str {
|
|
&self.model_name
|
|
}
|
|
|
|
pub fn dimensions(&self) -> u64 {
|
|
self.dimensions
|
|
}
|
|
}
|
|
|
|
pub async fn new_embed_client(config: &config::AppConfig) -> crate::Result<EmbedClient> {
|
|
let base_url = config
|
|
.get_embed_model_base_url()
|
|
.map_err(|e| crate::AgentError::Internal(e.to_string()))?;
|
|
let api_key = config
|
|
.get_embed_model_api_key()
|
|
.map_err(|e| crate::AgentError::Internal(e.to_string()))?;
|
|
let qdrant_url = config
|
|
.get_qdrant_url()
|
|
.map_err(|e| crate::AgentError::Internal(e.to_string()))?;
|
|
let qdrant_api_key = config.get_qdrant_api_key();
|
|
|
|
let openai = rig::providers::openai::Client::builder()
|
|
.api_key(&api_key)
|
|
.base_url(&base_url)
|
|
.build()
|
|
.map_err(|e| {
|
|
crate::AgentError::Internal(format!("failed to build rig openai client: {}", e))
|
|
})?;
|
|
|
|
let qdrant = QdrantClient::new(&qdrant_url, qdrant_api_key.as_deref()).await?;
|
|
Ok(EmbedClient::new(openai, qdrant))
|
|
}
|