src/main/java/com/example/cmisagent/VectorSearchService.java - master - akadnikov/CMIS_agent_Java

src/main/java/com/example/cmisagent/VectorSearchService.java
223 строки8 KB
akadnikov
Initial commit
25 фев 2026, 23:17
25 фев 2026, 23:17882add4
package com.example.cmisagent;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Сервис для семантического поиска по документам
 * Использует векторные эмбеддинги и similarity search
 */
@Slf4j
@Service
public class VectorSearchService {

    private final EmbeddingModel embeddingModel;
    private final EmbeddingStore<TextSegment> embeddingStore;
    private final DocumentSplitter documentSplitter;
    private final CMISService cmisService;

    // Кэш метаданных документов
    private final Map<String, DocumentMetadata> documentCache = new HashMap<>();

    public VectorSearchService(CMISService cmisService) {
        this.cmisService = cmisService;
        
        // Инициализация модели эмбеддингов
        // Для русского языка можно использовать multilingual модели
        this.embeddingModel = new AllMiniLmL6V2EmbeddingModel();
        
        // In-memory хранилище эмбеддингов
        // В продакшене лучше использовать Pinecone, Weaviate, или Chroma
        this.embeddingStore = new InMemoryEmbeddingStore<>();
        
        // Сплиттер для разбиения документов на чанки
        this.documentSplitter = DocumentSplitters.recursive(500, 50);
        
        log.info("VectorSearchService initialized");
    }

    /**
     * Инициализация - индексация документов при старте
     */
    @PostConstruct
    public void initialize() {
        log.info("Starting document indexing...");
        // В реальном проекте здесь можно запустить фоновую задачу
        // для индексации всех документов из CMIS
    }

    /**
     * Семантический поиск документов
     */
    public List<SemanticSearchResult> search(String query, int topK) {
        log.info("Semantic search: query='{}', topK={}", query, topK);

        try {
            // Создаем эмбеддинг для запроса
            Embedding queryEmbedding = embeddingModel.embed(query).content();

            // Ищем наиболее похожие сегменты
            List<EmbeddingMatch<TextSegment>> matches = embeddingStore.findRelevant(
                    queryEmbedding, 
                    topK
            );

            // Конвертируем в результаты
            List<SemanticSearchResult> results = new ArrayList<>();
            for (EmbeddingMatch<TextSegment> match : matches) {
                TextSegment segment = match.embedded();
                
                // Извлекаем метаданные из сегмента
                Map<String, String> metadata = segment.metadata().toMap();
                String documentId = metadata.get("documentId");
                String documentName = metadata.get("documentName");

                results.add(SemanticSearchResult.builder()
                        .documentId(documentId)
                        .documentName(documentName)
                        .content(segment.text())
                        .relevance(match.score())
                        .build());
            }

            log.info("Found {} relevant documents", results.size());
            return results;

        } catch (Exception e) {
            log.error("Error during semantic search", e);
            throw new CMISException("Error during semantic search", e);
        }
    }

    /**
     * Индексация документа для семантического поиска
     */
    public void indexDocument(String documentId) {
        log.info("Indexing document: {}", documentId);

        try {
            // Получаем содержимое документа
            DocumentContent content = cmisService.getContent(documentId);
            if (content == null || content.getText() == null) {
                log.warn("Cannot index document without content: {}", documentId);
                return;
            }

            // Получаем метаданные
            DocumentMetadata metadata = cmisService.getMetadata(documentId);
            if (metadata != null) {
                documentCache.put(documentId, metadata);
            }

            // Создаем документ для индексации
            Document document = Document.from(
                    content.getText(),
                    dev.langchain4j.data.document.Metadata.from(Map.of(
                            "documentId", documentId,
                            "documentName", content.getName(),
                            "mimeType", content.getMimeType()
                    ))
            );

            // Разбиваем на сегменты
            List<TextSegment> segments = documentSplitter.split(document);

            // Создаем эмбеддинги для каждого сегмента
            for (TextSegment segment : segments) {
                Embedding embedding = embeddingModel.embed(segment).content();
                embeddingStore.add(embedding, segment);
            }

            log.info("Document indexed successfully: {} ({} segments)", 
                    documentId, segments.size());

        } catch (Exception e) {
            log.error("Error indexing document: {}", documentId, e);
            throw new CMISException("Error indexing document", e);
        }
    }

    /**
     * Пакетная индексация документов
     */
    public void indexDocuments(List<String> documentIds) {
        log.info("Batch indexing {} documents", documentIds.size());

        int successCount = 0;
        int failureCount = 0;

        for (String documentId : documentIds) {
            try {
                indexDocument(documentId);
                successCount++;
            } catch (Exception e) {
                log.error("Failed to index document: {}", documentId, e);
                failureCount++;
            }
        }

        log.info("Batch indexing complete: {} successful, {} failed", 
                successCount, failureCount);
    }

    /**
     * Индексация всех документов из CMIS
     */
    public void indexAllDocuments() {
        log.info("Starting full document indexing...");

        try {
            // Поиск всех документов
            SearchCriteria criteria = SearchCriteria.builder()
                    .maxResults(1000)  // Ограничение для безопасности
                    .build();

            List<DocumentInfo> documents = cmisService.search(criteria);
            
            List<String> documentIds = documents.stream()
                    .map(DocumentInfo::getId)
                    .collect(Collectors.toList());

            indexDocuments(documentIds);

            log.info("Full indexing complete: {} documents", documents.size());

        } catch (Exception e) {
            log.error("Error during full indexing", e);
            throw new CMISException("Error during full indexing", e);
        }
    }

    /**
     * Очистка индекса
     */
    public void clearIndex() {
        embeddingStore.removeAll();
        documentCache.clear();
        log.info("Embedding store cleared");
    }

    /**
     * Получение статистики индекса
     */
    public Map<String, Object> getIndexStats() {
        return Map.of(
                "totalSegments", embeddingStore.findAll().size(),
                "cachedDocuments", documentCache.size()
        );
    }
}