use tracing::info;
use crate::agent::{self, config::AgentConfig, events::EventEmitter as AgentEventEmitter};
use crate::client::types::QueryResult;
use crate::document::{DocumentTree, NavigationIndex, ReasoningIndex};
use crate::error::Result;
use crate::events::{EventEmitter, QueryEvent};
use crate::llm::LlmClient;
use crate::retrieval::{dispatcher, postprocessor};
pub(crate) struct RetrieverClient {
llm: LlmClient,
config: AgentConfig,
events: EventEmitter,
}
impl RetrieverClient {
pub fn new(llm: LlmClient) -> Self {
Self {
llm,
config: AgentConfig::default(),
events: EventEmitter::new(),
}
}
pub fn with_events(mut self, events: EventEmitter) -> Self {
self.events = events;
self
}
pub fn with_config(mut self, config: AgentConfig) -> Self {
self.config = config;
self
}
pub fn config(&self) -> &AgentConfig {
&self.config
}
pub fn llm(&self) -> &LlmClient {
&self.llm
}
#[tracing::instrument(skip_all, fields(question = %question, docs = documents.len()))]
pub async fn query(
&self,
documents: &[(DocumentTree, NavigationIndex, ReasoningIndex, String)],
question: &str,
skip_analysis: bool,
) -> Result<QueryResult> {
self.events.emit_query(QueryEvent::Started {
query: question.to_string(),
});
info!(
docs = documents.len(),
skip_analysis, "Querying: {:?}", question
);
let doc_contexts: Vec<agent::DocContext> = documents
.iter()
.map(|(tree, nav, ridx, id)| agent::DocContext {
tree,
nav_index: nav,
reasoning_index: ridx,
doc_name: id.as_str(),
})
.collect();
let scope = if skip_analysis {
agent::Scope::Specified(doc_contexts)
} else {
agent::Scope::Workspace(agent::WorkspaceContext::new(doc_contexts))
};
let emitter = AgentEventEmitter::noop();
let output =
dispatcher::dispatch(question, scope, &self.config, &self.llm, &emitter).await?;
let fallback_id = documents
.first()
.map(|(_, _, _, id)| id.as_str())
.unwrap_or("");
let items = postprocessor::to_results(&output, fallback_id);
let result = QueryResult::new_with_items(items);
self.events.emit_query(QueryEvent::Complete {
total_results: result.len(),
confidence: result.single().map(|i| i.confidence).unwrap_or(0.0),
});
Ok(result)
}
}
impl Clone for RetrieverClient {
fn clone(&self) -> Self {
Self {
llm: self.llm.clone(),
config: self.config.clone(),
events: self.events.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_retriever_client_creation() {
let _client =
RetrieverClient::new(LlmClient::new(crate::llm::config::LlmConfig::default()));
}
}