rig_experimental/
routing.rs

1//! This module provides an abstraction for semantic routing.
2//!
3//! Example usage can be found in the `routing` example on the repository: <https://github.com/joshua-mo-143/rig-extra/blob/main/examples/routing.rs>
4use std::collections::HashMap;
5
6use serde::{Deserialize, Serialize};
7
8use rig::{
9    agent::Agent,
10    completion::{CompletionModel, Prompt},
11    vector_store::VectorStoreIndex,
12};
13
14/// The core semantic router abstraction.
15/// Contains a vector store index and a cosine similarity score threshold.
16pub struct SemanticRouter<V> {
17    store: V,
18    threshold: f64,
19}
20
21/// An abstraction over [`SemanticRouter`] that additionally contains Rig agents.
22/// Currently, each agent must be of the same completion model.
23pub struct SemanticRouterWithAgents<V, M: CompletionModel> {
24    store: V,
25    threshold: f64,
26    agents: HashMap<String, Agent<M>>,
27}
28
29impl<V> SemanticRouter<V> {
30    /// Create an instance of [`SemanticRouterBuilder`].
31    pub fn builder() -> SemanticRouterBuilder<V> {
32        SemanticRouterBuilder::new()
33    }
34}
35
36impl<V> SemanticRouter<V>
37where
38    V: VectorStoreIndex,
39{
40    pub async fn prompt(&self, query: &str) -> Option<String> {
41        let res = self.store.top_n(query, 1).await.ok()?;
42        let (score, _, SemanticRoute { tag }) = res.first()?;
43
44        tracing::info!("Retrieved route: {tag}, {score}");
45
46        if *score < self.threshold {
47            return None;
48        }
49
50        Some(tag.to_owned())
51    }
52
53    pub fn agent<M: CompletionModel>(
54        self,
55        route: &str,
56        agent: Agent<M>,
57    ) -> SemanticRouterWithAgents<V, M> {
58        let mut agents = HashMap::new();
59        agents.insert(route.to_string(), agent);
60        SemanticRouterWithAgents {
61            store: self.store,
62            threshold: self.threshold,
63            agents,
64        }
65    }
66}
67
68impl<V, M> SemanticRouterWithAgents<V, M>
69where
70    V: VectorStoreIndex,
71    M: CompletionModel,
72{
73    /// P
74    pub async fn prompt<R>(&self, query: R) -> Result<Option<String>, Box<dyn std::error::Error>>
75    where
76        R: Into<RouterRequest>,
77    {
78        let RouterRequest { query, turns } = query.into();
79        let res = self.store.top_n(&query, 1).await?;
80        let (score, _, SemanticRoute { tag }) = if let Some(result) = res.first() {
81            result
82        } else {
83            return Ok(None);
84        };
85
86        if *score < self.threshold {
87            return Ok(None);
88        }
89
90        let Some(agent) = self.agents.get(tag) else {
91            panic!("Couldn't find an agent that exists at tag: {tag}");
92        };
93
94        let res = if turns > 0 {
95            agent
96                .prompt(query)
97                .multi_turn(turns as usize)
98                .await
99                .unwrap()
100        } else {
101            agent.prompt(query).await.unwrap()
102        };
103
104        Ok(Some(res))
105    }
106
107    pub fn agent(mut self, route: &str, agent: Agent<M>) -> Self {
108        self.agents.insert(route.to_string(), agent);
109        self
110    }
111}
112
113pub struct RouterRequest {
114    query: String,
115    turns: u64,
116}
117
118impl RouterRequest {
119    pub fn new(query: String) -> Self {
120        Self::from(query)
121    }
122
123    pub fn with_turns(mut self, turns: u64) -> Self {
124        self.turns = turns;
125        self
126    }
127}
128
129impl From<String> for RouterRequest {
130    fn from(value: String) -> Self {
131        Self {
132            query: value,
133            turns: 0,
134        }
135    }
136}
137
138impl From<&str> for RouterRequest {
139    fn from(value: &str) -> Self {
140        Self {
141            query: value.to_string(),
142            turns: 0,
143        }
144    }
145}
146
147impl From<(String, u64)> for RouterRequest {
148    fn from((query, turns): (String, u64)) -> Self {
149        Self { query, turns }
150    }
151}
152
153impl From<(&str, u64)> for RouterRequest {
154    fn from((query, turns): (&str, u64)) -> Self {
155        Self {
156            query: query.to_string(),
157            turns,
158        }
159    }
160}
161
162#[derive(Serialize, Deserialize)]
163pub struct SemanticRoute {
164    tag: String,
165}
166
167pub trait Router: VectorStoreIndex {
168    fn retrieve_route() -> impl std::future::Future<Output = Option<String>> + Send;
169}
170
171pub struct SemanticRouterBuilder<V> {
172    store: Option<V>,
173    threshold: Option<f64>,
174}
175
176impl<V> Default for SemanticRouterBuilder<V> {
177    fn default() -> Self {
178        Self::new()
179    }
180}
181
182impl<V> SemanticRouterBuilder<V> {
183    pub fn new() -> Self {
184        Self {
185            store: None,
186            threshold: None,
187        }
188    }
189
190    pub fn store(mut self, router: V) -> Self {
191        self.store = Some(router);
192
193        self
194    }
195
196    pub fn threshold(mut self, threshold: f64) -> Self {
197        self.threshold = Some(threshold);
198
199        self
200    }
201
202    pub fn build(self) -> Result<SemanticRouter<V>, SemanticRouterError> {
203        let Some(store) = self.store else {
204            return Err(SemanticRouterError::StoreNotFound);
205        };
206
207        let threshold = self.threshold.unwrap_or(0.8);
208
209        Ok(SemanticRouter { store, threshold })
210    }
211}
212
213#[derive(thiserror::Error, Debug)]
214pub enum SemanticRouterError {
215    #[error("Vector store not found")]
216    StoreNotFound,
217}
rig_experimental/routing.rs

rig_experimental/
routing.rs
