use anyhow::{bail, Result};
use clap::Parser;
use qdrant_client::Qdrant; use std::sync::Arc;
use clap::{Subcommand};
use qdrant_client::qdrant::{
payload_index_params::IndexParams, FieldType, IntegerIndexParams, KeywordIndexParams, TextIndexParams, TokenizerType, UpdateStatus,
UpsertPointsBuilder,
UpsertPoints,
CreateFieldIndexCollectionBuilder,
GetCollectionInfoResponse,
PointStruct,
};
use indicatif::ProgressBar;
use log;
use crate::config::AppConfig;
pub(crate) const BATCH_SIZE: usize = 128;
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
pub struct CliArgs {
#[command(subcommand)]
pub command: Commands,
#[arg(short = 'm', long = "onnx-model", global = true, env = "VECTORDB_ONNX_MODEL")]
pub onnx_model_path_arg: Option<String>,
#[arg(short = 't', long = "onnx-tokenizer-dir", global = true, env = "VECTORDB_ONNX_TOKENIZER_DIR")]
pub onnx_tokenizer_dir_arg: Option<String>,
}
pub const FIELD_FILE_PATH: &str = "file_path";
pub const FIELD_START_LINE: &str = "start_line";
pub const FIELD_END_LINE: &str = "end_line";
pub const FIELD_FILE_EXTENSION: &str = "file_extension";
pub const FIELD_LANGUAGE: &str = "language";
pub const FIELD_ELEMENT_TYPE: &str = "element_type";
pub const FIELD_CHUNK_CONTENT: &str = "chunk_content";
pub const LEGACY_INDEX_COLLECTION: &str = "vectordb-code-search";
pub const FIELD_BRANCH: &str = "branch";
pub const FIELD_COMMIT_HASH: &str = "commit_hash";
pub const SIMPLE_INDEX_COLLECTION: &str = "vectordb-code-search";
#[derive(Subcommand, Debug)]
pub enum Commands {
#[command(subcommand_negates_reqs = true)]
Repo(super::repo_commands::RepoArgs),
#[command(subcommand_negates_reqs = true)]
Simple(super::simple::SimpleArgs), }
pub async fn handle_command(
args: CliArgs,
config: &mut AppConfig,
client: Arc<Qdrant>,
) -> Result<()> {
match args.command {
Commands::Repo(ref cmd_args) => super::repo_commands::handle_repo_command(cmd_args.clone(), &args, config, client, None).await,
Commands::Simple(ref cmd_args) => super::simple::handle_simple_command(cmd_args.clone(), &args, config.clone(), client).await,
}
}
pub async fn ensure_payload_index(
client: &Qdrant,
collection_name: &str,
field_name: &str,
field_type: FieldType,
_is_keyword: bool, tokenizer: Option<TokenizerType>,
) -> Result<()> {
let info: GetCollectionInfoResponse = client.collection_info(collection_name).await?;
if let Some(config) = info.result {
if config.payload_schema.get(field_name).is_some() {
log::debug!("Payload index for '{}' on field '{}' already exists.", collection_name, field_name);
return Ok(());
}
} else {
bail!("Could not retrieve collection info for {}", collection_name);
}
log::info!("Creating payload index for '{}' on field '{}'...", collection_name, field_name);
let index_params = match field_type {
FieldType::Keyword => Some(IndexParams::KeywordIndexParams(KeywordIndexParams {
on_disk: None,
is_tenant: Some(false),
})),
FieldType::Integer => Some(IndexParams::IntegerIndexParams(IntegerIndexParams {
lookup: Some(false),
range: Some(false),
is_principal: Some(false),
on_disk: None, })),
FieldType::Text => Some(IndexParams::TextIndexParams(TextIndexParams {
tokenizer: tokenizer.map(|t| t.into()).unwrap_or(TokenizerType::Word.into()),
lowercase: Some(true),
min_token_len: None,
max_token_len: None,
on_disk: None,
})),
_ => None,
};
let mut builder = CreateFieldIndexCollectionBuilder::new(collection_name, field_name, field_type);
if let Some(params) = index_params {
builder = builder.field_index_params(params);
}
match client.create_field_index(builder).await {
Ok(response) => {
if let Some(result) = response.result {
match UpdateStatus::try_from(result.status) {
Ok(UpdateStatus::Completed) => {
log::info!("Payload index created successfully for field '{}'.", field_name);
Ok(())
}
Ok(status) => {
log::warn!("Payload index creation for field '{}' resulted in status: {:?}", field_name, status);
Ok(())
}
Err(_) => {
log::warn!("Payload index creation for field '{}' returned unknown status: {}", field_name, result.status);
Ok(())
}
}
} else {
log::warn!("Payload index creation response for field '{}' did not contain a result.", field_name);
Ok(())
}
}
Err(e) => {
log::error!("Failed to create payload index for field '{}': {}. Ignoring error, assuming index might exist.", field_name, e);
Ok(())
}
}
}
pub(crate) async fn upsert_batch(
client: &Qdrant,
collection_name: &str,
points: Vec<PointStruct>,
_batch_num: usize,
_total_batches: usize,
progress_bar: &ProgressBar,
) -> Result<()> {
if points.is_empty() {
return Ok(());
}
let num_points = points.len();
let request: UpsertPoints = UpsertPointsBuilder::new(collection_name, points)
.wait(false)
.build();
match client.upsert_points(request).await {
Ok(response) => {
if let Some(result) = response.result {
match UpdateStatus::try_from(result.status) {
Ok(UpdateStatus::Completed) => {
progress_bar.inc(num_points as u64);
log::debug!("Upsert batch successful.");
Ok(())
},
Ok(status) => {
let msg = format!("Qdrant upsert batch completed with status: {:?}", status);
log::debug!("{}", msg);
Ok(())
},
Err(_) => {
let msg = format!("Qdrant upsert batch completed with unknown status code: {}", result.status);
progress_bar.println(format!("Error: {}", msg));
log::error!("{}", msg);
Err(anyhow::anyhow!(msg))
}
}
} else {
let msg = "Qdrant upsert response missing result status";
progress_bar.println(format!("Error: {}", msg));
log::error!("{}", msg);
Err(anyhow::anyhow!(msg))
}
},
Err(e) => {
let msg = format!("Failed to upsert batch to {}: {}", collection_name, e);
progress_bar.println(format!("Error: {}", msg));
log::error!("{}", msg);
Err(anyhow::anyhow!(msg).context(e))
}
}
}