core/rag/src/api/models.rs - main - alexefan136/flowstack

core/rag/src/api/models.rs
461 строка14 KB
Alexander Efanov
Обновление репозитория
15 июл 2026, 12:19
15 июл 2026, 12:1976704c6
//! Request и Response модели для API.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::domain::DocumentType;

// ============================================================================
// Ingest Requests
// ============================================================================

/// Запрос на ingestion текстового документа.
#[derive(Debug, Clone, Deserialize)]
pub struct IngestTextRequest {
    /// Содержимое документа.
    pub content: String,

    /// Источник (путь к файлу или URL).
    pub source: String,

    /// Тип документа.
    #[serde(default = "default_doc_type")]
    pub doc_type: DocumentType,

    /// Workspace ID (для multi-tenant).
    pub workspace_id: String,

    /// Заголовок документа (опционально).
    pub title: Option<String>,

    /// Автор (опционально).
    pub author: Option<String>,

    /// Язык (опционально, по умолчанию "ru").
    #[serde(default = "default_language")]
    pub language: String,

    /// Дополнительные метаданные.
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Запрос на batch ingestion.
#[derive(Debug, Clone, Deserialize)]
pub struct IngestBatchRequest {
    /// Список документов.
    pub documents: Vec<IngestTextRequest>,
}

fn default_doc_type() -> DocumentType {
    DocumentType::Markdown
}

fn default_language() -> String {
    "ru".to_string()
}

// ============================================================================
// Ingest Responses
// ============================================================================

/// Ответ на успешный ingestion.
#[derive(Debug, Clone, Serialize)]
pub struct IngestResponse {
    /// ID созданного документа.
    pub document_id: String,

    /// Количество созданных чанков.
    pub chunks_count: usize,

    /// Время обработки в миллисекундах.
    pub processing_time_ms: f64,

    /// Детальная разбивка времени.
    pub timing: IngestTiming,

    /// Использованная стратегия chunking.
    pub chunking_strategy: String,

    /// Размерность векторов.
    pub vector_dimensions: usize,
}

/// Разбивка времени по этапам.
#[derive(Debug, Clone, Serialize)]
pub struct IngestTiming {
    pub chunking_ms: f64,
    pub embedding_ms: f64,
    pub storage_ms: f64,
    pub total_ms: f64,
}

/// Ответ на batch ingestion.
#[derive(Debug, Clone, Serialize)]
pub struct IngestBatchResponse {
    /// Результаты для каждого документа.
    pub results: Vec<IngestBatchResult>,

    /// Общее время в миллисекундах.
    pub total_time_ms: f64,

    /// Количество успешных ingestions.
    pub success_count: usize,

    /// Количество ошибок.
    pub error_count: usize,
}

/// Результат для одного документа в batch.
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum IngestBatchResult {
    Success(IngestResponse),
    Error { error: String, source: String },
}

// ============================================================================
// Search Requests
// ============================================================================

/// Запрос на поиск.
#[derive(Debug, Clone, Deserialize)]
pub struct SearchRequest {
    /// Текст запроса.
    pub query: String,

    /// Workspace ID (опционально).
    pub workspace_id: Option<String>,

    /// Количество результатов (по умолчанию 10).
    #[serde(default = "default_top_k")]
    pub top_k: usize,

    /// Минимальный score (опционально).
    pub score_threshold: Option<f32>,

    /// Фильтр по типам документов.
    pub doc_types: Option<Vec<DocumentType>>,

    /// Фильтр по источникам.
    pub sources: Option<Vec<String>>,

    /// Дополнительные фильтры.
    #[serde(default)]
    pub filters: HashMap<String, serde_json::Value>,
}

fn default_top_k() -> usize {
    10
}

// ============================================================================
// Search Responses
// ============================================================================

/// Ответ на поиск.
#[derive(Debug, Clone, Serialize)]
pub struct SearchResponse {
    /// Оригинальный запрос.
    pub query: String,

    /// Результаты поиска.
    pub results: Vec<SearchResultItem>,

    /// Количество результатов.
    pub results_count: usize,

    /// Отформатированный контекст для LLM.
    pub context: String,

    /// Время выполнения.
    pub timing: SearchTiming,

    /// Использованные модели.
    pub models: SearchModels,
}

/// Один результат поиска.
#[derive(Debug, Clone, Serialize)]
pub struct SearchResultItem {
    /// ID чанка.
    pub chunk_id: String,

    /// ID документа.
    pub document_id: String,

    /// Содержимое чанка.
    pub content: String,

    /// Score релевантности (0.0 - 1.0).
    pub score: f32,

    /// Rerank score (если использовался reranker).
    pub rerank_score: Option<f32>,

    /// Эффективный score (rerank если есть, иначе обычный).
    pub effective_score: f32,

    /// Метаданные.
    pub metadata: SearchMetadata,
}

/// Метаданные результата.
#[derive(Debug, Clone, Serialize)]
pub struct SearchMetadata {
    pub source: String,
    pub title: String,
    pub workspace_id: String,
    pub doc_type: DocumentType,
    pub chunk_index: usize,
    pub char_count: usize,
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

/// Время выполнения поиска.
#[derive(Debug, Clone, Serialize)]
pub struct SearchTiming {
    pub retrieval_ms: f64,
    pub rerank_ms: Option<f64>,
    pub total_ms: f64,
}

/// Использованные модели.
#[derive(Debug, Clone, Serialize)]
pub struct SearchModels {
    pub embedding: Option<String>,
    pub reranker: Option<String>,
}

/// Ответ на batch search.
#[derive(Debug, Clone, Serialize)]
pub struct SearchBatchResponse {
    /// Результаты для каждого запроса.
    pub results: Vec<SearchBatchResult>,

    /// Общее время в миллисекундах.
    pub total_time_ms: f64,
}

/// Результат для одного запроса в batch.
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum SearchBatchResult {
    Success(SearchResponse),
    Error { query: String, error: String },
}

// ============================================================================
// Document Management
// ============================================================================

/// Информация о документе.
#[derive(Debug, Clone, Serialize)]
pub struct DocumentInfo {
    pub id: String,
    pub source: String,
    pub doc_type: DocumentType,
    pub title: String,
    pub workspace_id: String,
    pub chunks_count: usize,
    pub created_at: String,
}

/// Список документов.
#[derive(Debug, Clone, Serialize)]
pub struct DocumentListResponse {
    pub documents: Vec<DocumentInfo>,
    pub total: usize,
    pub workspace_id: Option<String>,
}

/// Ответ на удаление.
#[derive(Debug, Clone, Serialize)]
pub struct DeleteResponse {
    pub deleted_count: usize,
    pub target: String,
}

// ============================================================================
// System Endpoints
// ============================================================================

/// Health check response.
#[derive(Debug, Clone, Serialize)]
pub struct HealthResponse {
    pub status: String,
    pub version: String,
    pub components: ComponentsHealth,
}

/// Health отдельных компонентов.
#[derive(Debug, Clone, Serialize)]
pub struct ComponentsHealth {
    pub storage: bool,
    pub embedder: bool,
}

/// Readiness check response.
#[derive(Debug, Clone, Serialize)]
pub struct ReadyResponse {
    pub status: String,
    pub stats: crate::storage::StorageStats,
}

/// Stats response.
#[derive(Debug, Clone, Serialize)]
pub struct StatsResponse {
    pub storage: crate::storage::StorageStats,
    pub config: PipelineConfigResponse,
}

/// Конфигурация pipeline.
#[derive(Debug, Clone, Serialize)]
pub struct PipelineConfigResponse {
    pub chunking_strategy: String,
    pub chunk_size: usize,
    pub chunk_overlap: usize,
    pub retrieval_top_k: usize,
    pub use_reranker: bool,
    pub use_hybrid: bool,
}

// ============================================================================
// Conversions
// ============================================================================

impl From<crate::pipeline::IngestResult> for IngestResponse {
    fn from(result: crate::pipeline::IngestResult) -> Self {
        Self {
            document_id: result.document_id,
            chunks_count: result.chunks_count,
            processing_time_ms: result.processing_time_ms,
            timing: IngestTiming {
                chunking_ms: result.chunking_time_ms,
                embedding_ms: result.embedding_time_ms,
                storage_ms: result.storage_time_ms,
                total_ms: result.processing_time_ms,
            },
            chunking_strategy: result.chunking_strategy,
            vector_dimensions: result.vector_dimensions,
        }
    }
}

impl From<crate::pipeline::SearchResult> for SearchResponse {
    fn from(result: crate::pipeline::SearchResult) -> Self {
        let results: Vec<SearchResultItem> = result
            .retrieval
            .results
            .into_iter()
            .map(|r| {
                // ==========================================================
                // Вычисляем ВСЕ значения ДО перемещения полей из r.chunk
                // ==========================================================

                // Копируемые значения (score)
                let score = r.score;
                let rerank_score = r.rerank_score;
                let effective_score = r.effective_score();

                // Значения из chunk
                let char_count = r.chunk.char_count();
                let chunk_index = r.chunk.index;

                // Извлекаем метаданные ДО перемещения
                let source = r
                    .chunk
                    .metadata
                    .get("source")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let title = r
                    .chunk
                    .metadata
                    .get("title")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let workspace_id = r
                    .chunk
                    .metadata
                    .get("workspace_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let doc_type = r
                    .chunk
                    .metadata
                    .get("doc_type")
                    .and_then(|v| v.as_str())
                    .and_then(|s| serde_json::from_str(&format!("\"{}\"", s)).ok())
                    .unwrap_or(DocumentType::Markdown);

                // Extra metadata (всё кроме стандартных полей)
                let extra: HashMap<String, serde_json::Value> = r
                    .chunk
                    .metadata
                    .iter()
                    .filter(|(k, _)| {
                        !["source", "title", "workspace_id", "doc_type"].contains(&k.as_str())
                    })
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect();

                // ==========================================================
                // Теперь перемещаем поля из r.chunk
                // ==========================================================
                SearchResultItem {
                    chunk_id: r.chunk.id,
                    document_id: r.chunk.document_id,
                    content: r.chunk.content,
                    score,
                    rerank_score,
                    effective_score,
                    metadata: SearchMetadata {
                        source,
                        title,
                        workspace_id,
                        doc_type,
                        chunk_index,
                        char_count,
                        extra,
                    },
                }
            })
            .collect();

        let results_count = results.len();

        Self {
            query: result.query,
            results,
            results_count,
            context: result.context,
            timing: SearchTiming {
                retrieval_ms: result.retrieval.retrieval_time_ms,
                rerank_ms: result.retrieval.rerank_time_ms,
                total_ms: result.total_time_ms,
            },
            models: SearchModels {
                embedding: result.retrieval.embedding_model,
                reranker: result.retrieval.reranker_model,
            },
        }
    }
}

impl From<crate::pipeline::PipelineConfig> for PipelineConfigResponse {
    fn from(config: crate::pipeline::PipelineConfig) -> Self {
        Self {
            chunking_strategy: format!("{:?}", config.chunking_strategy),
            chunk_size: config.chunking.chunk_size,
            chunk_overlap: config.chunking.chunk_overlap,
            retrieval_top_k: config.retrieval.top_k,
            use_reranker: config.retrieval.use_reranker,
            use_hybrid: config.retrieval.use_hybrid,
        }
    }
}