dci_tool/

agent.rs

//! The Direct Corpus Interaction agent: a rig [`Agent`] pre-wired with the four
//! corpus tools and a preamble that teaches the search → narrow → read → cite
//! investigation loop.
//!
//! The agent is generic over any [`CompletionModel`], so callers bring their
//! own model/provider. Nothing here is tied to a specific vendor.

use rig_core::agent::{Agent, AgentBuilder};
use rig_core::completion::{CompletionModel, Prompt, PromptError};

use crate::sandbox::CorpusRoot;
use crate::tools::CorpusTools;

/// Default number of tool-calling turns allowed per investigation.
pub const DEFAULT_MAX_TURNS: usize = 24;

/// The default system preamble. It frames the corpus as something to be
/// *interrogated directly* with commands rather than retrieved from an index.
pub const DEFAULT_PREAMBLE: &str = "\
You are a Direct Corpus Interaction (DCI) analyst. You answer questions about a \
corpus of files (code, logs, documents) by issuing search commands directly \
against the raw text — there is no vector database and no pre-built index.

You have four tools:
- corpus_list: list a directory to orient yourself.
- corpus_find: locate files by a path glob (e.g. '**/*.log').
- corpus_search: search file contents with a regular expression; returns
  file:line:text evidence.
- corpus_read: read a bounded, line-numbered window from one file.

Method:
1. Start broad: search for the most specific term or pattern in the question.
2. Narrow using path globs and follow-up searches; pivot on identifiers,
   error codes, IPs, hashes, or usernames you discover.
3. Read the surrounding lines of promising hits to confirm before concluding.
4. Cite concrete evidence as `path:line` for every claim. If the corpus does
   not support a conclusion, say so plainly rather than guessing.

Prefer precise regular expressions over broad ones, and prefer reading a few
lines of real evidence over speculating.";

/// Builder for a [`DciAgent`].
pub struct DciAgentBuilder<M: CompletionModel> {
    model: M,
    corpus: CorpusRoot,
    preamble: Option<String>,
    appended: Vec<String>,
    max_turns: usize,
    temperature: Option<f64>,
    max_tokens: Option<u64>,
    model_label: Option<String>,
}

impl<M: CompletionModel + 'static> DciAgentBuilder<M> {
    /// Start building a DCI agent over `model` and `corpus`.
    pub fn new(model: M, corpus: CorpusRoot) -> Self {
        Self {
            model,
            corpus,
            preamble: None,
            appended: Vec::new(),
            max_turns: DEFAULT_MAX_TURNS,
            temperature: None,
            max_tokens: None,
            model_label: None,
        }
    }

    /// Replace the default preamble entirely.
    pub fn preamble(mut self, preamble: impl Into<String>) -> Self {
        self.preamble = Some(preamble.into());
        self
    }

    /// Append domain guidance (e.g. a cyber-investigation playbook) after the
    /// base preamble.
    pub fn append_preamble(mut self, extra: impl Into<String>) -> Self {
        self.appended.push(extra.into());
        self
    }

    /// Set the maximum number of tool-calling turns per investigation.
    pub fn max_turns(mut self, turns: usize) -> Self {
        self.max_turns = turns;
        self
    }

    /// Set the model sampling temperature.
    pub fn temperature(mut self, temperature: f64) -> Self {
        self.temperature = Some(temperature);
        self
    }

    /// Set the maximum number of output tokens per turn. Required by some
    /// providers (e.g. Anthropic).
    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    /// Set a human-readable model label recorded on telemetry events (e.g.
    /// `"gpt-4o"`). Defaults to `"unknown"`.
    pub fn model_label(mut self, label: impl Into<String>) -> Self {
        self.model_label = Some(label.into());
        self
    }

    /// Finish building, registering the four corpus tools on the agent.
    pub fn build(self) -> DciAgent<M> {
        let tools = CorpusTools::new(self.corpus);

        let mut preamble = self
            .preamble
            .unwrap_or_else(|| DEFAULT_PREAMBLE.to_string());
        for extra in &self.appended {
            preamble.push_str("\n\n");
            preamble.push_str(extra);
        }

        let mut builder = AgentBuilder::new(self.model).preamble(&preamble);
        if let Some(temp) = self.temperature {
            builder = builder.temperature(temp);
        }
        if let Some(max_tokens) = self.max_tokens {
            builder = builder.max_tokens(max_tokens);
        }
        let agent = builder
            .tool(tools.search)
            .tool(tools.find)
            .tool(tools.read)
            .tool(tools.list)
            .build();

        DciAgent {
            agent,
            max_turns: self.max_turns,
            model_label: self.model_label.unwrap_or_else(|| "unknown".to_string()),
        }
    }
}

/// A ready-to-run Direct Corpus Interaction agent.
pub struct DciAgent<M: CompletionModel, P: rig_core::agent::PromptHook<M> = ()> {
    agent: Agent<M, P>,
    max_turns: usize,
    model_label: String,
}

impl<M: CompletionModel + 'static> DciAgent<M, ()> {
    /// Start building a DCI agent.
    pub fn builder(model: M, corpus: CorpusRoot) -> DciAgentBuilder<M> {
        DciAgentBuilder::new(model, corpus)
    }
}

impl<M: CompletionModel + 'static, P: rig_core::agent::PromptHook<M> + 'static> DciAgent<M, P> {
    /// Borrow the underlying rig agent (e.g. to register additional tools or
    /// wrap it as a delegate in later phases).
    pub fn agent(&self) -> &Agent<M, P> {
        &self.agent
    }

    /// The configured per-investigation turn budget.
    pub fn max_turns(&self) -> usize {
        self.max_turns
    }

    /// Run an investigation: prompt the agent and let it interact with the
    /// corpus across multiple tool-calling turns until it produces an answer.
    ///
    /// Token usage for the whole run is emitted as a telemetry
    /// `prompt.completed` event (see [`crate::telemetry`]).
    pub async fn investigate(&self, question: &str) -> Result<String, PromptError> {
        let start = std::time::Instant::now();
        let response = self
            .agent
            .prompt(question)
            .max_turns(self.max_turns)
            .extended_details()
            .await?;
        let usage = &response.usage;
        crate::telemetry::record_prompt(
            &self.model_label,
            usage.input_tokens,
            usage.output_tokens,
            usage.cached_input_tokens,
            usage.reasoning_tokens,
            start.elapsed().as_millis() as u64,
        );
        Ok(response.output)
    }
}