agentroot_core/llm/
traits.rs

1//! LLM trait definitions
2
3use crate::error::Result;
4use async_trait::async_trait;
5
6/// Embedding generation trait
7#[async_trait]
8pub trait Embedder: Send + Sync {
9    /// Generate embedding for single text
10    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
11
12    /// Generate embeddings for batch of texts
13    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;
14
15    /// Get embedding dimensions
16    fn dimensions(&self) -> usize;
17
18    /// Get model name
19    fn model_name(&self) -> &str;
20}
21
22/// Document reranking trait
23#[async_trait]
24pub trait Reranker: Send + Sync {
25    /// Rerank documents for a query
26    async fn rerank(&self, query: &str, documents: &[RerankDocument]) -> Result<Vec<RerankResult>>;
27
28    /// Get model name
29    fn model_name(&self) -> &str;
30}
31
32/// Document for reranking
33#[derive(Debug, Clone)]
34pub struct RerankDocument {
35    pub id: String,
36    pub text: String,
37}
38
39/// Reranking result
40#[derive(Debug, Clone)]
41pub struct RerankResult {
42    pub id: String,
43    pub score: f64,
44}
45
46/// Query expansion trait
47#[async_trait]
48pub trait QueryExpander: Send + Sync {
49    /// Expand query into variants
50    async fn expand(&self, query: &str, context: Option<&str>) -> Result<ExpandedQuery>;
51
52    /// Get model name
53    fn model_name(&self) -> &str;
54}
55
56/// Expanded query variants
57#[derive(Debug, Clone, Default)]
58pub struct ExpandedQuery {
59    /// Lexical variations for BM25
60    pub lexical: Vec<String>,
61    /// Semantic variations for vector search
62    pub semantic: Vec<String>,
63    /// Hypothetical document (HyDE)
64    pub hyde: Option<String>,
65}
66
67/// Tokenization trait
68#[async_trait]
69pub trait Tokenizer: Send + Sync {
70    /// Tokenize text
71    async fn tokenize(&self, text: &str) -> Result<Vec<u32>>;
72
73    /// Detokenize tokens back to text
74    async fn detokenize(&self, tokens: &[u32]) -> Result<String>;
75
76    /// Count tokens
77    async fn count_tokens(&self, text: &str) -> Result<usize>;
78}
agentroot_core/llm/traits.rs

agentroot_core/llm/
traits.rs
