diff --git a/iped-api/src/main/java/iped/properties/ExtraProperties.java b/iped-api/src/main/java/iped/properties/ExtraProperties.java index 252e840e6b..8f79dd70f7 100644 --- a/iped-api/src/main/java/iped/properties/ExtraProperties.java +++ b/iped-api/src/main/java/iped/properties/ExtraProperties.java @@ -203,4 +203,11 @@ public class ExtraProperties { public static final List COMMUNICATION_BASIC_PROPS = Arrays.asList(MESSAGE_SUBJECT, MESSAGE_BODY, Message.MESSAGE_CC, Message.MESSAGE_BCC, Message.MESSAGE_RECIPIENT_ADDRESS, MESSAGE_IS_ATTACHMENT, MESSAGE_ATTACHMENT_COUNT.getName()); + + /** + * Extra attribute set to {@code "true"} on a parent evidence item when at least + * one of its text fragments has had an embedding generated and stored in the + * RAG vector index. Allows filtering/searching for embeddable items in the UI. + */ + public static final String HAS_EMBEDDING = "hasEmbedding"; //$NON-NLS-1$ } diff --git a/iped-app/resources/config/IPEDConfig.txt b/iped-app/resources/config/IPEDConfig.txt index b385a8632f..d2a443f4a9 100644 --- a/iped-app/resources/config/IPEDConfig.txt +++ b/iped-app/resources/config/IPEDConfig.txt @@ -1,4 +1,4 @@ -######################################################################## +######################################################################## # Processing Settings ######################################################################## @@ -104,6 +104,10 @@ indexFileContents = true # You must configure at least host/port options in conf/ElasticSearchConfig.txt enableIndexToElasticSearch = false +# Enables Retrieval-Augmented Generation (RAG) capabilities. +# Requires an embedding model server (e.g. Ollama). Advanced settings can be configured in conf/RAGConfig.txt. +enableRAG = false + # Enables exporting files to MinIO object storage cluster. # You must configure at least host/port options in conf/MinIOConfig.txt enableMinIO = false diff --git a/iped-app/resources/config/conf/IndexTaskConfig.txt b/iped-app/resources/config/conf/IndexTaskConfig.txt index df58abfba0..7f1609a1ee 100644 --- a/iped-app/resources/config/conf/IndexTaskConfig.txt +++ b/iped-app/resources/config/conf/IndexTaskConfig.txt @@ -17,8 +17,13 @@ filterNonLatinChars = false # Size (bytes) of the text segments extracted from items before indexing. Includes all items, # not just the ones indexed via strings. This avoids OutOfMemory errors while indexing items with large chunks of extracted text. +# Reduce to a smaller value (e.g. 4000) when using local embeddings (RAG) to fit context window limits. textSplitSize = 10485760 +# Overlap size (characters) between consecutive text segments. +# Reduce to a smaller value (e.g. 1000) if textSplitSize is reduced for local embeddings. +# textOverlapSize = 10000 + # Uses NIOFSDirectory instead of MMAPDirectory to open index (https://lucene.apache.org/core/4_9_0/core/org/apache/lucene/store/FSDirectory.html) # It is a bit slower, but prevents JVM crashes when reading index through network shares. useNIOFSDirectory = false diff --git a/iped-app/resources/config/conf/RAGConfig.txt b/iped-app/resources/config/conf/RAGConfig.txt new file mode 100644 index 0000000000..09838a3592 --- /dev/null +++ b/iped-app/resources/config/conf/RAGConfig.txt @@ -0,0 +1,171 @@ +############################################################################### +# IPED RAG & AI ASSISTANT CONFIGURATION # +# # +# Native Retrieval-Augmented Generation & Vector Search Engine for IPED # +# Default Opt-In: Set enableRAG = true to activate vector embedding indexing. # +############################################################################### + +############################################################################### +# 1. EMBEDDING GENERATION & MODEL PROVIDER +############################################################################### + +# Provider for text embeddings. Options: local, remote, gemini +# - local : Ollama server (use 'local' for Ollama EVEN IF running on a remote GPU server over IP/network). +# - remote : OpenAI-compatible API endpoint (e.g. OpenAI text-embedding-3-small, Azure, Cohere). +# - gemini : Google Generative Language Embeddings API. +embeddingProvider = local + +# Model name for text embeddings: +# - Ollama (local) : bge-m3 (1024 dims), nomic-embed-text (768 dims), mxbai-embed-large (1024 dims) +# - Gemini (remote): gemini-embedding-001 (768 dims) +# - OpenAI (remote): text-embedding-3-small (1536 dims), text-embedding-3-large (3072 dims) +embeddingModel = bge-m3 + +# HTTP Endpoint for generating embeddings: +# - Ollama (local) : http://localhost:11434/api/embeddings or http://:11434/api/embeddings +# * Concurrency Tip: For dedicated GPU setups, environment variables like OLLAMA_NUM_PARALLEL (e.g. 2-4) and OLLAMA_MAX_QUEUE (e.g. 512) can improve throughput depending on available VRAM/GPU hardware (requires further testing for optimal lab setup). +# - OpenAI (remote): https://api.openai.com/v1/embeddings +# - Gemini : leave empty to use official Google API endpoint automatically +embeddingEndpoint = http://localhost:11434/api/embeddings + +# Vector dimension size of the embedding model (must match embeddingModel): +# - bge-m3: 1024 | gemini-embedding-001: 768 | nomic-embed-text: 768 | text-embedding-3-small: 1536 +embeddingDimensions = 1024 + + +############################################################################### +# 2. CHUNKING & LLM CONTEXT WINDOW TUNING +# Controls text fragmenting during indexing and context size allocation for LLMs. +# +# MATHEMATICAL RELATIONSHIP & FORMULAS: +# 1. Character-to-Token Ratio: 1 token ??? 4 characters in Portuguese (3500 chars ??? 875 tokens). +# 2. Retrieved Context Formula: +# Retrieved Context Tokens ??? maxRetrievedChunks * (chunkSize / 4) +# 3. Overlap Guideline: +# chunkOverlap should be 10% to 25% of chunkSize (e.g. 700 for 3500). +# +# RECOMMENDED SETTINGS BY GPU VRAM / HARDWARE (For Local Ollama Models): +# - 6 GB VRAM GPU (RTX 2060 / 3050): +# llmContextWindow = 4096 | chunkSize = 3000 | maxRetrievedChunks = 4 (4B models) +# - 8 GB VRAM GPU (RTX 3060 / 4060) [DEFAULT BASELINE]: +# llmContextWindow = 9216 | chunkSize = 3500 | maxRetrievedChunks = 5 (4B models ~6GB VRAM) +# - 16 GB VRAM GPU (RTX A4000 / 4060 Ti 16GB / 4080): +# llmContextWindow = 9216 | chunkSize = 3500 | maxRetrievedChunks = 5 (8B models ~16GB VRAM, 100% GPU) +# * Note: Setting 16384 on 8B models balloons VRAM allocation to 25GB, forcing 43% CPU offload! +# - 24 GB VRAM GPU (RTX 3090 / 4090): +# llmContextWindow = 16384 | chunkSize = 4000 | maxRetrievedChunks = 10 (8B models ~24GB VRAM) +# +# NOTE FOR COMMERCIAL CLOUD MODELS (Google Gemini / Anthropic Claude / OpenAI): +# - Commercial APIs have massive native context windows (128k to 2M tokens). +# - Therefore, llmContextWindow is IGNORED for commercial providers (llmProvider != local). +# - For Cloud Models, simply increase maxRetrievedChunks = 15-20 and chunkSize = 4000. +############################################################################### + +# Size (in characters) of each text chunk/fragment created during indexing. +# Recommendation: 3000-5000 chars. Default: 3500 +chunkSize = 3500 + +# Overlap (in characters) between consecutive chunks to prevent boundary context loss. +# Recommendation: 10%-25% of chunkSize. Default: 700 +chunkOverlap = 700 + +# Maximum number of fragment chunks to retrieve and include in LLM context. +# Recommended for 8GB VRAM GPU: 5 | Cloud (Gemini/Claude/OpenAI): 15-20. Default: 5 +maxRetrievedChunks = 5 + +# Context window size (in tokens) STRICTLY FOR LOCAL OLLAMA MODELS (llmProvider = local). +# Dynamically configures 'num_ctx' on Ollama to prevent truncation and VRAM OOM. +# NOTE: Ignored for commercial APIs (Gemini, Claude, OpenAI) which natively accept huge contexts. +# Range: 2048 to 16384. Default: 9216 (optimized for 8GB VRAM GPUs). +llmContextWindow = 9216 + + +############################################################################### +# 3. VECTOR STORE BACKEND & HYBRID SEARCH RANKING (BM25 + KNN) +############################################################################### + +# Vector store backend for semantic (KNN) search: +# - lucene : (default) Vectors stored directly in local Lucene index. Cases are 100% portable. +# - opensearch : Vectors stored in external OpenSearch cluster via REST API. +# NOTE: Requires 'enableIndexToElasticSearch = true' in IPEDConfig.txt for cluster indexing. +# FALLBACK: Automatically falls back to 'lucene' if OpenSearch cluster is unreachable. +vectorStoreMode = lucene + +# Minimum raw vector similarity threshold applied BEFORE Reciprocal Rank Fusion (RRF). +# Discards weak vector matches below 0.62. Range: 0.0 to 1.0. Default: 0.62 +chunkSimilarityThreshold = 0.62 + +# Weight of vector (semantic) search in Reciprocal Rank Fusion (RRF). Range: 0.0 to 1.0. Default: 0.4 +# Note: Lexical (BM25) weight is automatically (1.0 - vectorSearchWeight) = 0.6. +vectorSearchWeight = 0.4 + + +############################################################################### +# 4. EVIDENCE FILTERING RULES (BLACKLIST & WHITELIST) +# Controls which evidence items generate embeddings during IPED indexing. +############################################################################### + +# Filtering mode for generating embeddings during indexing. Options: blacklist, whitelist +# - blacklist : excludes categories and MIME types matching the blacklist patterns below +# - whitelist : ONLY includes categories and MIME types matching the whitelist patterns below (Default) +embeddingFilterMode = whitelist + +# ----------------------------------------------------------------------------- +# BLACKLIST RULES (Active when embeddingFilterMode = blacklist) +# Excludes non-semantic binary data, system registries, archives, executables, and unallocated space. +# ----------------------------------------------------------------------------- + +# Blacklisted IPED Categories: +# Excludes: Programs/Libraries, Compressed Archives, Image Disks, Unallocated Space, File Slacks, Registry Files, Fonts, Geographic Data, Databases, XML Files. +embeddingCategoryBlacklist = ^(Programs and Libraries|Compressed Archives|Image Disks|Unallocated|File Slacks|Main Registry Files|Other Registry Files|Fonts|Geographic Data|Databases|XML Files|Other)$ + +# Blacklisted MIME Types: +# Excludes: Binaries (.exe, .dll, .so, .elf), Android binaries (.dex, .apk), Java archives (.jar), Databases (.sqlite, .edb), XML dumps/reports, Archives (.zip, .rar, .7z, .tar, .gz), Fonts, Shortcuts (.lnk), Debug Symbols (.pdb). +embeddingMimeTypeBlacklist = ^(application/octet-stream|application/x-dosexec|application/x-sharedlib|application/x-elf|application/x-executable|application/x-object|application/x-msdownload|application/x-ms-installer|application/x-msi|application/x-windows-registry-main|application/x-windows-registry|application/x-sqlite3|application/x-edb|application/x-msaccess|text/xml|application/xml|application/x-dex|application/vnd\.android\.dex|application/vnd\.android\.package-archive|application/java-archive|application/x-java-archive|font/.*|application/zip|application/x-7z-compressed|application/x-rar-compressed|application/x-tar|application/x-gzip|application/x-bzip2|application/x-ms-pdb|application/x-win-lnk)$ + +# ----------------------------------------------------------------------------- +# WHITELIST RULES (Active when embeddingFilterMode = whitelist) +# Restricts vector embeddings strictly to high-value textual evidence, documents, emails, and chats. +# ----------------------------------------------------------------------------- + +# Whitelisted IPED Categories: +# Includes: Documents, Emails, Chats, SMS, Instant Messages, Notes, Plain Texts, HTML and Web pages, PDFs, Office Documents, Audio. +embeddingCategoryWhitelist = ^(Documents|Emails and Mailboxes|Chats|SMS Messages|Instant Messages Artifacts|Notes|Plain Texts|HTML and Web pages|PDF Documents|Office Documents|Audio)$ + +# Whitelisted MIME Types: +# Includes: Plain Text, HTML, XHTML, CSV, TSV, Emails, PDF, Word, Excel, PowerPoint, OpenDocument, RTF, JSON, Audio files (transcriptions). +embeddingMimeTypeWhitelist = ^(text/plain|text/html|application/xhtml\+xml|text/csv|text/tab-separated-values|message/rfc822|application/pdf|application/msword|application/vnd\.openxmlformats-officedocument\..*|application/vnd\.ms-excel|application/vnd\.ms-powerpoint|application/vnd\.oasis\.opendocument\..*|application/rtf|application/json|audio/.*)$ + + +############################################################################### +# 5. LLM ASSISTANT INTEGRATION & QUERY SETTINGS +############################################################################### + +# Enables natural language questions through LLM integration (AI Assistant Chat Tab in IPED-SearchApp) +enableLLMQueries = true + +# Provider for LLM chat/queries. Options: local, remote, gemini, claude +# - local : Ollama server (use 'local' for Ollama EVEN IF running on a remote GPU server over IP/network! +# 'local' dynamically manages num_ctx context window tuning for Ollama without OOM/truncation). +# - remote : Official OpenAI API (gpt-4o, gpt-4o-mini) or OpenAI-compatible APIs requiring API Key (vLLM, LMStudio, Azure). +# - gemini : Google Gemini API (use free API Key from Google AI Studio). +# - claude : Anthropic Claude API. +llmProvider = local + +# Model name for the LLM: +# - Ollama (local) : qwen3:4b, gemma3:4b, llama3.1:8b, phi4 +# - OpenAI (remote): gpt-4o, gpt-4o-mini, o3-mini, gpt-4-turbo +# - Gemini (remote): gemini-2.5-flash, gemini-1.5-pro +# - Claude (remote): claude-sonnet-4-6 +llmModel = qwen3:8b + +# HTTP Endpoint for LLM requests: +# - Ollama (local) : http://localhost:11434/v1 or http://:11434/v1 (e.g. http://10.89.42.211:11434/v1) +# - OpenAI (remote): https://api.openai.com/v1 (or custom enterprise endpoint, e.g. http://:8000/v1) +# - Gemini/Claude : leave empty to use official API endpoint automatically +llmEndpoint = http://localhost:11434/v1 + +# API Key for remote cloud providers (remote, gemini, claude): +# - Required when using OpenAI (gpt-4o / gpt-4o-mini), Gemini, Claude, or authenticated APIs. +# - Leave blank for Ollama (llmProvider = local). +llmApiKey = diff --git a/iped-app/resources/localization/iped-desktop-messages.properties b/iped-app/resources/localization/iped-desktop-messages.properties index dfa6776be1..13338b941e 100644 --- a/iped-app/resources/localization/iped-desktop-messages.properties +++ b/iped-app/resources/localization/iped-desktop-messages.properties @@ -464,4 +464,35 @@ FieldValuePopupMenu.LessThan=Filter less than... FieldValuePopupMenu.Equals=Equals to... FieldValuePopupMenu.StartsWith=Starts with... FieldValuePopupMenu.Clear=Clear -FieldValuePopupMenu.Filter=Filter \ No newline at end of file +FieldValuePopupMenu.Filter=Filter +# RAG AI Assistant panel +App.RAGAssistant=AI Assistant +RAGAssistant.Question=Question... +RAGAssistant.Ask=Ask +RAGAssistant.Cancel=Cancel +RAGAssistant.History=History +RAGAssistant.NewChat=New Chat +RAGAssistant.DeleteChat=Delete Chat +RAGAssistant.RenameChat=Rename Chat +RAGAssistant.RenameChatPrompt=Enter new chat name: +RAGAssistant.RenameChatTitle=Rename Chat +RAGAssistant.ScopeSelected=Restrict to selection +RAGAssistant.ScopeSelectedTooltip=Restricts RAG search exclusively to checked, selected or filtered items in the table. +RAGAssistant.ConfirmDeleteChat=Are you sure you want to delete this chat thread from the history? +RAGAssistant.Sources=Sources (double-click to open) +RAGAssistant.DeleteQuestion=Delete Question +RAGAssistant.ConfirmDelete=Are you sure you want to delete this question from the history? +RAGAssistant.ConfirmDeleteTitle=Confirm Deletion +RAGAssistant.Thinking=Thinking... +RAGAssistant.ErrorTitle=AI Assistant Error +RAGAssistant.NotEnabled=The AI Assistant (RAG) is not enabled for this case.\nEnable 'enableRAG=true' in IPEDConfig.txt and re-index the case. +RAGAssistant.NoContextFound=No relevant documents were found in the material to answer this question. +RAGAssistant.Header.Id=ID +RAGAssistant.Header.Name=Name +RAGAssistant.Header.Score=Score +RAGAssistant.Header.Path=Path +iped.app.ui.SemanticSimilarityFilterer.filtererName=Semantic Similarity Filter (AI) +FilterValue.SemanticSimilarity=Semantic Similarity (AI) +RAGAssistant.YouLabel=You +RAGAssistant.AssistantLabel=Assistant +RAGAssistant.ContextWindowWarning=Warning: context window limit reached ({0} chunks omitted). diff --git a/iped-app/resources/localization/iped-desktop-messages_pt_BR.properties b/iped-app/resources/localization/iped-desktop-messages_pt_BR.properties index 6eb2c25725..37b37c4d27 100644 --- a/iped-app/resources/localization/iped-desktop-messages_pt_BR.properties +++ b/iped-app/resources/localization/iped-desktop-messages_pt_BR.properties @@ -465,3 +465,35 @@ FieldValuePopupMenu.Equals=Igual a... FieldValuePopupMenu.StartsWith=Iniciam com... FieldValuePopupMenu.Clear=Limpar FieldValuePopupMenu.Filter=Filtrar + +# Painel do Assistente IA (RAG) +App.RAGAssistant=Assistente IA +RAGAssistant.Question=Pergunta... +RAGAssistant.Ask=Perguntar +RAGAssistant.Cancel=Cancelar +RAGAssistant.History=Hist\u00f3rico +RAGAssistant.NewChat=Nova Conversa +RAGAssistant.DeleteChat=Excluir Conversa +RAGAssistant.RenameChat=Renomear Conversa +RAGAssistant.RenameChatPrompt=Digite o novo nome para a conversa: +RAGAssistant.RenameChatTitle=Renomear Conversa +RAGAssistant.ScopeSelected=Restringir a selecao +RAGAssistant.ScopeSelectedTooltip=Restringe a busca RAG apenas aos itens marcados, selecionados na tabela ou visiveis nos filtros atuais. +RAGAssistant.ConfirmDeleteChat=Tem certeza de que deseja excluir esta conversa do historico? +RAGAssistant.Sources=Fontes (duplo clique para abrir) +RAGAssistant.DeleteQuestion=Excluir Pergunta +RAGAssistant.ConfirmDelete=Tem certeza de que deseja excluir esta pergunta do historico? +RAGAssistant.ConfirmDeleteTitle=Confirmar Exclusao +RAGAssistant.Thinking=Consultando IA... +RAGAssistant.ErrorTitle=Erro no Assistente IA +RAGAssistant.NotEnabled=O Assistente IA (RAG) nao esta habilitado para este caso.\nHabilite 'enableRAG=true' no IPEDConfig.txt e reindexe o caso. +RAGAssistant.NoContextFound=Nenhum documento relevante foi encontrado no material para responder a esta pergunta. +RAGAssistant.Header.Id=ID +RAGAssistant.Header.Name=Nome +RAGAssistant.Header.Score=Score +RAGAssistant.Header.Path=Caminho +iped.app.ui.SemanticSimilarityFilterer.filtererName=Filtro de Similaridade Semantica (IA) +FilterValue.SemanticSimilarity=Similaridade Semantica (IA) +RAGAssistant.YouLabel=Voce +RAGAssistant.AssistantLabel=Assistente +RAGAssistant.ContextWindowWarning=Aviso: limite da janela de contexto atingido ({0} trechos omitidos). diff --git a/iped-app/src/main/java/iped/app/ui/App.java b/iped-app/src/main/java/iped/app/ui/App.java index 89e7442087..8ec089c2af 100644 --- a/iped-app/src/main/java/iped/app/ui/App.java +++ b/iped-app/src/main/java/iped/app/ui/App.java @@ -149,6 +149,8 @@ import iped.engine.data.IPEDSource; import iped.engine.data.ItemId; import iped.engine.search.IPEDSearcher; +import iped.engine.rag.RAGService; +import iped.app.ui.rag.RAGAssistantPanel; import iped.engine.search.ImageSimilarityLowScoreFilter; import iped.engine.search.ImageSimilarityScorer; import iped.engine.search.MultiSearchResult; @@ -230,7 +232,9 @@ public class App extends JFrame implements WindowListener, IMultiSearchResultPro BookmarksTreeListener bookmarksListener; TimelineListener timelineListener; public CControl dockingControl; - private DefaultSingleCDockable categoriesTabDock, aiFiltersTabDock, metadataTabDock, bookmarksTabDock, evidenceTabDock; + private DefaultSingleCDockable categoriesTabDock, aiFiltersTabDock, metadataTabDock, bookmarksTabDock, + evidenceTabDock; + private DefaultSingleCDockable ragAssistantTabDock; private List rsTabDock = new ArrayList(); private DefaultSingleCDockable tableTabDock, galleryTabDock, graphDock; @@ -246,7 +250,8 @@ public class App extends JFrame implements WindowListener, IMultiSearchResultPro Color defaultColor; Color defaultFocusedColor; Color defaultSelectedColor; - private JScrollPane hitsScroll, subItemScroll, parentItemScroll, duplicatesScroll, referencesScroll, referencedByScroll; + private JScrollPane hitsScroll, subItemScroll, parentItemScroll, duplicatesScroll, referencesScroll, + referencedByScroll; JScrollPane viewerScroll, resultsScroll, galleryScroll; JPanel topPanel; ClearFilterButton clearAllFilters; @@ -275,6 +280,7 @@ public class App extends JFrame implements WindowListener, IMultiSearchResultPro public IItem similarFacesRefItem; public List similarFacesPrevSortKeys; SimilarDocumentFilterer similarDocumentFilterer; + public SemanticSimilarityFilterer semanticSimilarityFilterer; public File casesPathFile; boolean isMultiCase; @@ -356,7 +362,7 @@ public Manager getProcessingManager() { public MenuClass getContextMenu() { IItemId id = resultTableListener.getSelectedItemId(); - IItem item = id == null ? null : appCase.getItemByItemId(id); + IItem item = id == null ? null : appCase.getItemByItemId(id); return new MenuClass(item); } @@ -368,7 +374,8 @@ public String getFontStartTag() { return this.fontStartTag; } - public void init(LogConfiguration logConfiguration, boolean isMultiCase, File casesPathFile, Manager processingManager, String codePath) { + public void init(LogConfiguration logConfiguration, boolean isMultiCase, File casesPathFile, + Manager processingManager, String codePath) { this.logConfiguration = logConfiguration; this.isMultiCase = isMultiCase; @@ -547,6 +554,7 @@ public void createGUI() { filterManager = new FilterManager(filterComboBox); similarDocumentFilterer = new SimilarDocumentFilterer(); + semanticSimilarityFilterer = new SemanticSimilarityFilterer(); similarImagesFilterer = new SimilarImagesQueryFilterer(); @@ -731,7 +739,8 @@ public TabComponent createTabComponent(EclipseTabPane pane, Dockable dockable) { private static final long serialVersionUID = -9020339124009415001L; public void setLabelInsets(Insets labelInsets) { - labelInsets = new Insets(labelInsets.top - 1, labelInsets.left - 3, labelInsets.bottom - 1, labelInsets.right - 3); + labelInsets = new Insets(labelInsets.top - 1, labelInsets.left - 3, labelInsets.bottom - 1, + labelInsets.right - 3); super.setLabelInsets(labelInsets); } }; @@ -749,30 +758,36 @@ public TabPanePainter createDecorationPainter(EclipseTabPane pane) { // Customize appearance of buttons and check boxes shown in docking frames title // bar, so focus is not painted (avoiding intersection with buttons icons) and // more clear indication when a CCheckBox is selected. - dockingControl.getController().getActionViewConverter().putClient(ActionType.BUTTON, ViewTarget.TITLE, new ViewGenerator>() { - public BasicTitleViewItem create(ActionViewConverter converter, ButtonDockAction action, Dockable dockable) { - BasicButtonHandler handler = new BasicButtonHandler(action, dockable); - CustomButton button = new CustomButton(handler, handler); - handler.setModel(button.getModel()); - return handler; - } - }); - dockingControl.getController().getActionViewConverter().putTheme(ActionType.CHECK, ViewTarget.TITLE, new ViewGenerator>() { - public BasicTitleViewItem create(ActionViewConverter converter, SelectableDockAction action, Dockable dockable) { - BasicSelectableHandler.Check handler = new BasicSelectableHandler.Check(action, dockable); - CustomButton button = new CustomButton(handler, handler); - handler.setModel(button.getModel()); - return handler; - } - }); + dockingControl.getController().getActionViewConverter().putClient(ActionType.BUTTON, ViewTarget.TITLE, + new ViewGenerator>() { + public BasicTitleViewItem create(ActionViewConverter converter, ButtonDockAction action, + Dockable dockable) { + BasicButtonHandler handler = new BasicButtonHandler(action, dockable); + CustomButton button = new CustomButton(handler, handler); + handler.setModel(button.getModel()); + return handler; + } + }); + dockingControl.getController().getActionViewConverter().putTheme(ActionType.CHECK, ViewTarget.TITLE, + new ViewGenerator>() { + public BasicTitleViewItem create(ActionViewConverter converter, + SelectableDockAction action, Dockable dockable) { + BasicSelectableHandler.Check handler = new BasicSelectableHandler.Check(action, dockable); + CustomButton button = new CustomButton(handler, handler); + handler.setModel(button.getModel()); + return handler; + } + }); dockingControl.putProperty(StackDockStation.TAB_PLACEMENT, TabPlacement.TOP_OF_DOCKABLE); this.getContentPane().add(dockingControl.getContentArea(), BorderLayout.CENTER); defaultColor = dockingControl.getController().getColors().get(ColorMap.COLOR_KEY_TAB_BACKGROUND); defaultFocusedColor = dockingControl.getController().getColors().get(ColorMap.COLOR_KEY_TAB_BACKGROUND_FOCUSED); - defaultSelectedColor = dockingControl.getController().getColors().get(ColorMap.COLOR_KEY_TAB_BACKGROUND_SELECTED); + defaultSelectedColor = dockingControl.getController().getColors() + .get(ColorMap.COLOR_KEY_TAB_BACKGROUND_SELECTED); - timelineButton = new CCheckBox(Messages.get("App.ToggleTimelineView"), IconUtil.getToolbarIcon("time", resPath)) { + timelineButton = new CCheckBox(Messages.get("App.ToggleTimelineView"), + IconUtil.getToolbarIcon("time", resPath)) { protected void changed() { if (timelineListener != null) timelineListener.setTimelineTableView(isSelected()); @@ -829,6 +844,7 @@ protected void changed() { filterManager.addQueryFilterer(treeListener); filterManager.addQueryFilterer(similarImagesFilterer); filterManager.addQueryFilterer(similarDocumentFilterer); + filterManager.addQueryFilterer(semanticSimilarityFilterer); filterManager.addQueryFilterer(TableHeaderFilterManager.get()); filterManager.addResultSetFilterer(bookmarksListener); filterManager.addResultSetFilterer(FilterSelectedEdges.getInstance()); @@ -980,7 +996,8 @@ private void createAllDockables() { graphDock = createDockable("graphtab", Messages.getString("App.Links"), appGraphAnalytics); - CCheckBox butToggleVideoFramesMode = new CCheckBox(Messages.getString("Gallery.ToggleVideoFrames"), IconUtil.getToolbarIcon("video", resPath)) { + CCheckBox butToggleVideoFramesMode = new CCheckBox(Messages.getString("Gallery.ToggleVideoFrames"), + IconUtil.getToolbarIcon("video", resPath)) { protected void changed() { galleryModel.clearVideoThumbsInCache(); useVideoThumbsInGallery = isSelected(); @@ -990,7 +1007,8 @@ protected void changed() { galleryTabDock.addAction(butToggleVideoFramesMode); galleryTabDock.addSeparator(); - galleryGrayButton = new CCheckBox(Messages.getString("Gallery.GalleryGrayFilter"), IconUtil.getToolbarIcon("gray-scale", resPath)) { + galleryGrayButton = new CCheckBox(Messages.getString("Gallery.GalleryGrayFilter"), + IconUtil.getToolbarIcon("gray-scale", resPath)) { protected void changed() { galleryModel.setGrayFilter(isSelected()); gallery.repaint(); @@ -998,7 +1016,8 @@ protected void changed() { }; galleryTabDock.addAction(galleryGrayButton); - galleryBlurButton = new CCheckBox(Messages.getString("Gallery.GalleryBlurFilter"), IconUtil.getToolbarIcon("blur-image", resPath)) { + galleryBlurButton = new CCheckBox(Messages.getString("Gallery.GalleryBlurFilter"), + IconUtil.getToolbarIcon("blur-image", resPath)) { protected void changed() { galleryModel.setBlurFilter(isSelected()); gallery.repaint(); @@ -1007,7 +1026,8 @@ protected void changed() { galleryTabDock.addAction(galleryBlurButton); galleryTabDock.addSeparator(); - butSimSearch = new CButton(Messages.getString("MenuClass.FindSimilarImages"), IconUtil.getToolbarIcon("find", resPath)); + butSimSearch = new CButton(Messages.getString("MenuClass.FindSimilarImages"), + IconUtil.getToolbarIcon("find", resPath)); galleryTabDock.addAction(butSimSearch); butSimSearch.addActionListener(new ActionListener() { public void actionPerformed(ActionEvent e) { @@ -1016,7 +1036,8 @@ public void actionPerformed(ActionEvent e) { }); butSimSearch.setEnabled(false); - butFaceSearch = new CButton(Messages.getString("MenuClass.FindSimilarFaces"), IconUtil.getToolbarIcon("face", resPath)); + butFaceSearch = new CButton(Messages.getString("MenuClass.FindSimilarFaces"), + IconUtil.getToolbarIcon("face", resPath)); galleryTabDock.addAction(butFaceSearch); butFaceSearch.addActionListener(new ActionListener() { public void actionPerformed(ActionEvent e) { @@ -1024,13 +1045,15 @@ public void actionPerformed(ActionEvent e) { } }); butFaceSearch.setEnabled(false); - + galleryTabDock.addSeparator(); // Add buttons to control the thumbnails size / number of columns in the gallery - CButton butDec = new CButton(Messages.getString("Gallery.DecreaseThumbsSize"), IconUtil.getToolbarIcon("minus", resPath)); + CButton butDec = new CButton(Messages.getString("Gallery.DecreaseThumbsSize"), + IconUtil.getToolbarIcon("minus", resPath)); galleryTabDock.addAction(butDec); - CButton butInc = new CButton(Messages.getString("Gallery.IncreaseThumbsSize"), IconUtil.getToolbarIcon("plus", resPath)); + CButton butInc = new CButton(Messages.getString("Gallery.IncreaseThumbsSize"), + IconUtil.getToolbarIcon("plus", resPath)); galleryTabDock.addAction(butInc); galleryTabDock.addSeparator(); butDec.addActionListener(new ActionListener() { @@ -1049,7 +1072,8 @@ public void actionPerformed(ActionEvent e) { final ResultSetViewer resultSetViewer = iterator.next(); resultSetViewer.init(resultsTable, this, this); - DefaultSingleCDockable tabDock = createDockable(resultSetViewer.getID(), resultSetViewer.getTitle(), resultSetViewer.getPanel()); + DefaultSingleCDockable tabDock = createDockable(resultSetViewer.getID(), resultSetViewer.getTitle(), + resultSetViewer.getPanel()); resultSetViewer.setDockableContainer(tabDock); @@ -1079,7 +1103,8 @@ public void changed(CDockableLocationEvent event) { duplicateDock = createDockable("duplicatestab", Messages.getString("DuplicatesTableModel.Duplicates"), //$NON-NLS-1$ //$NON-NLS-2$ duplicatesScroll); referencesDock = createDockable("referencestab", Messages.getString("ReferencesTab.Title"), referencesScroll); - referencedByDock = createDockable("referencedbytab", Messages.getString("ReferencedByTab.Title"), referencedByScroll); + referencedByDock = createDockable("referencedbytab", Messages.getString("ReferencedByTab.Title"), + referencedByScroll); dockingControl.addDockable(categoriesTabDock); dockingControl.addDockable(aiFiltersTabDock); @@ -1095,6 +1120,26 @@ public void changed(CDockableLocationEvent event) { dockingControl.addDockable(graphDock); } + // Add AI Assistant tab if RAG is enabled for this case + RAGService ragService = RAGService.getInstance(); + if (ragService != null && ragService.getConfig().isEnabled()) { + ragAssistantTabDock = createDockable("ragassistanttab", Messages.getString("App.RAGAssistant"), + new RAGAssistantPanel()); + dockingControl.addDockable(ragAssistantTabDock); + + // Inject the Lucene IndexSearcher so RAGService can perform local KNN searches. + // This is required when vectorStoreMode=lucene (the default portable mode). + if ("lucene".equalsIgnoreCase(ragService.getConfig().getVectorStoreMode()) + && appCase != null) { + try { + org.apache.lucene.search.IndexSearcher luceneSearcher = appCase.getSearcher(); + ragService.setLuceneSearcher(luceneSearcher); + } catch (Exception e) { + LOGGER.warn("Could not inject Lucene IndexSearcher into RAGService.", e); + } + } + } + for (Iterator iterator = rsTabDock.iterator(); iterator.hasNext();) { DefaultSingleCDockable tabDock = iterator.next(); dockingControl.addDockable(tabDock); @@ -1110,7 +1155,8 @@ public void changed(CDockableLocationEvent event) { List viewers = viewerController.getViewers(); viewerDocks = new ArrayList(); for (AbstractViewer viewer : viewers) { - DefaultSingleCDockable viewerDock = createDockable(viewer.getClass().getName(), viewer.getName(), viewer.getPanel()); + DefaultSingleCDockable viewerDock = createDockable(viewer.getClass().getName(), viewer.getName(), + viewer.getPanel()); viewerDocks.add(viewerDock); dockingControl.addDockable(viewerDock); viewerController.put(viewer, viewerDock); @@ -1120,7 +1166,8 @@ public void changed(CDockableLocationEvent event) { } private void setupViewerDocks() { - CCheckBox chkFixed = new CCheckBox(Messages.getString("ViewerController.FixViewer"), IconUtil.getToolbarIcon("pin", resPath)) { + CCheckBox chkFixed = new CCheckBox(Messages.getString("ViewerController.FixViewer"), + IconUtil.getToolbarIcon("pin", resPath)) { protected void changed() { viewerController.setFixed(isSelected()); } @@ -1139,7 +1186,8 @@ public void changed(CDockableLocationEvent event) { if (newLocation != null) { newLocation = newLocation.getParent(); } - if ((oldLocation == null && newLocation != null) || (oldLocation != null && !oldLocation.equals(newLocation))) { + if ((oldLocation == null && newLocation != null) + || (oldLocation != null && !oldLocation.equals(newLocation))) { validated = viewerController.validateViewer(viewer); } } @@ -1157,7 +1205,7 @@ public void changed(CDockableLocationEvent event) { public void actionPerformed(ActionEvent e) { ((CButton) viewerDock.getAction("prevHit")).setEnabled(false); ((CButton) viewerDock.getAction("nextHit")).setEnabled(false); - ((MultiViewer)viewer).searchInViewer(); + ((MultiViewer) viewer).searchInViewer(); } }); viewerDock.addAction(searchInViewer); @@ -1174,8 +1222,10 @@ public void actionPerformed(ActionEvent e) { viewerDock.addAction(copyViewerImage); } - CButton prevHit = new CButton(Messages.getString("ViewerController.PrevHit"), IconUtil.getToolbarIcon("prev", resPath)); - CButton nextHit = new CButton(Messages.getString("ViewerController.NextHit"), IconUtil.getToolbarIcon("next", resPath)); + CButton prevHit = new CButton(Messages.getString("ViewerController.PrevHit"), + IconUtil.getToolbarIcon("prev", resPath)); + CButton nextHit = new CButton(Messages.getString("ViewerController.NextHit"), + IconUtil.getToolbarIcon("next", resPath)); prevHit.addActionListener(new ActionListener() { public void actionPerformed(ActionEvent e) { viewer.scrollToNextHit(false); @@ -1196,7 +1246,8 @@ public void actionPerformed(ActionEvent e) { if (toolbarSupport >= 0) { Icon downIcon = IconUtil.getToolbarIcon("down", resPath); Icon upIcon = IconUtil.getToolbarIcon("up", resPath); - CSelButton butToolbar = new CSelButton(Messages.getString("ViewerController.ShowToolBar"), upIcon, downIcon); + CSelButton butToolbar = new CSelButton(Messages.getString("ViewerController.ShowToolBar"), upIcon, + downIcon); butToolbar.addActionListener(new ActionListener() { public void actionPerformed(ActionEvent e) { butToolbar.toggle(); @@ -1266,7 +1317,7 @@ private void removeAllDockables() { List docks = new ArrayList<>(); docks.addAll(Arrays.asList(hitsDock, subitemDock, duplicateDock, parentDock, tableTabDock, galleryTabDock, bookmarksTabDock, evidenceTabDock, metadataTabDock, categoriesTabDock, aiFiltersTabDock, graphDock, - referencesDock, referencedByDock, filtersTabDock)); + referencesDock, referencedByDock, filtersTabDock, ragAssistantTabDock)); docks.addAll(viewerDocks); docks.addAll(rsTabDock); rsTabDock.clear(); @@ -1372,7 +1423,8 @@ private void setTabColor(DefaultSingleCDockable dock, boolean isDefault) { } } - private void setTabColor(DefaultSingleCDockable dock, Color colorBackground, Color colorFocused, Color colorSelected) { + private void setTabColor(DefaultSingleCDockable dock, Color colorBackground, Color colorFocused, + Color colorSelected) { dock.getColors().setColor(ColorMap.COLOR_KEY_TAB_BACKGROUND, colorBackground); dock.getColors().setColor(ColorMap.COLOR_KEY_TAB_BACKGROUND_FOCUSED, colorFocused); dock.getColors().setColor(ColorMap.COLOR_KEY_TAB_BACKGROUND_SELECTED, colorSelected); @@ -1422,7 +1474,8 @@ public void adjustLayout(boolean isReset) { removeAllDockables(); createAllDockables(); - tableTabDock.setLocation(verticalLayout ? CLocation.base().normalNorth(0.7) : CLocation.base().normalNorth(0.5)); + tableTabDock + .setLocation(verticalLayout ? CLocation.base().normalNorth(0.7) : CLocation.base().normalNorth(0.5)); tableTabDock.setVisible(true); galleryTabDock.setLocationsAside(tableTabDock); @@ -1441,7 +1494,8 @@ public void adjustLayout(boolean isReset) { graphDock.setVisible(true); } - hitsDock.setLocation(verticalLayout ? CLocation.base().normalSouth(0.3) : CLocation.base().normalSouth(0.5).west(0.4)); + hitsDock.setLocation( + verticalLayout ? CLocation.base().normalSouth(0.3) : CLocation.base().normalSouth(0.5).west(0.4)); hitsDock.setVisible(true); subitemDock.setLocationsAside(hitsDock); @@ -1484,6 +1538,11 @@ public void adjustLayout(boolean isReset) { aiFiltersTabDock.setLocationsAside(metadataTabDock); aiFiltersTabDock.setVisible(true); + if (ragAssistantTabDock != null) { + ragAssistantTabDock.setLocationsAside(aiFiltersTabDock); + ragAssistantTabDock.setVisible(true); + } + if (verticalLayout) adjustViewerLayout(); @@ -1501,7 +1560,8 @@ public void adjustLayout(boolean isReset) { private void adjustViewerLayout() { DefaultSingleCDockable prevDock = viewerDocks.get(0); - prevDock.setLocation(verticalLayout ? CLocation.base().normalEast(0.35) : CLocation.base().normalSouth(0.5).east(0.6)); + prevDock.setLocation( + verticalLayout ? CLocation.base().normalEast(0.35) : CLocation.base().normalSouth(0.5).east(0.6)); prevDock.setVisible(true); for (int i = 1; i < viewerDocks.size(); i++) { DefaultSingleCDockable dock = viewerDocks.get(i); @@ -1516,6 +1576,20 @@ public void toggleHorizontalVerticalLayout() { adjustLayout(true); } + public void initRAGFeatures() { + RAGService ragService = RAGService.getInstance(); + if (ragService != null && ragService.getConfig().isEnabled()) { + if (ragAssistantTabDock == null) { + ragAssistantTabDock = createDockable("ragassistanttab", Messages.getString("App.RAGAssistant"), + new RAGAssistantPanel()); + dockingControl.addDockable(ragAssistantTabDock); + } + ragAssistantTabDock.setLocationsAside(aiFiltersTabDock); + ragAssistantTabDock.setVisible(true); + selectDockableTab(ragAssistantTabDock); + } + } + @Override public void windowActivated(WindowEvent e) { } @@ -1767,7 +1841,8 @@ public SimilarImageFilter(IItemId itemId, IItem similarImagesQueryRefItem) { } @Override - public IMultiSearchResult filterResult(IMultiSearchResult src) throws ParseException, QueryNodeException, IOException { + public IMultiSearchResult filterResult(IMultiSearchResult src) + throws ParseException, QueryNodeException, IOException { IMultiSearchResult result = src; if (refItem != null) { LOGGER.info("Starting similar image search..."); @@ -1862,7 +1937,8 @@ public String toString() { } @Override - public IMultiSearchResult filterResult(IMultiSearchResult src) throws ParseException, QueryNodeException, IOException { + public IMultiSearchResult filterResult(IMultiSearchResult src) + throws ParseException, QueryNodeException, IOException { IMultiSearchResult result = src; DynamicDuplicateFilter duplicateFilter = new DynamicDuplicateFilter(App.get().appCase); result = duplicateFilter.filter(src); @@ -1934,7 +2010,8 @@ public SimilarFacesSearchFilter(IItemId itemIdRef, IItem similarFacesRefItem) { } @Override - public IMultiSearchResult filterResult(IMultiSearchResult src) throws ParseException, QueryNodeException, IOException { + public IMultiSearchResult filterResult(IMultiSearchResult src) + throws ParseException, QueryNodeException, IOException { if (itemRef != null) { return sfs.filter((MultiSearchResult) src); } @@ -2136,7 +2213,8 @@ public void setPercent(int percent) { } public String toString() { - return Messages.get("FilterValue.SimilarDocument") + " (" + Integer.toString(filterPercent) + "%): " + itemRef.getName(); + return Messages.get("FilterValue.SimilarDocument") + " (" + Integer.toString(filterPercent) + "%): " + + itemRef.getName(); } @Override diff --git a/iped-app/src/main/java/iped/app/ui/CategoryTreeModel.java b/iped-app/src/main/java/iped/app/ui/CategoryTreeModel.java index f872f61c63..24c970bdd4 100644 --- a/iped-app/src/main/java/iped/app/ui/CategoryTreeModel.java +++ b/iped-app/src/main/java/iped/app/ui/CategoryTreeModel.java @@ -32,7 +32,15 @@ public static void install() { } private CategoryTreeModel() { - this.root = App.get().appCase.getCategoryTree(); + App appInstance = App.get(); + if (appInstance == null) { + System.err.println("DIAGNOSTIC: App.get() is NULL!"); + } else if (appInstance.appCase == null) { + System.err.println("DIAGNOSTIC: App.get().appCase is NULL!"); + } else { + System.err.println("DIAGNOSTIC: App.get().appCase is NOT null, categoryTree = " + appInstance.appCase.getCategoryTree()); + } + this.root = appInstance.appCase.getCategoryTree(); this.root.setName(rootName); } diff --git a/iped-app/src/main/java/iped/app/ui/MenuClass.java b/iped-app/src/main/java/iped/app/ui/MenuClass.java index 4cea5aa811..ae9eb3f663 100644 --- a/iped-app/src/main/java/iped/app/ui/MenuClass.java +++ b/iped-app/src/main/java/iped/app/ui/MenuClass.java @@ -48,7 +48,7 @@ public class MenuClass extends JPopupMenu { JMenuItem exportHighlighted, copyHighlighted, checkHighlighted, uncheckHighlighted, readHighlighted, unreadHighlighted, exportChecked, copyChecked, saveBookmarks, loadBookmarks, checkHighlightedAndSubItems, uncheckHighlightedAndSubItems, checkHighlightedAndParent, uncheckHighlightedAndParent, checkHighlightedAndReferences, uncheckHighlightedAndReferences, checkHighlightedAndReferencedBy, uncheckHighlightedAndReferencedBy, changeGalleryColCount, defaultLayout, changeLayout, previewScreenshot, manageBookmarks, clearSearchHistory, importKeywords, navigateToParent, exportTerms, manageFilters, manageColumns, exportCheckedToZip, exportCheckedTreeToZip, - exportTree, exportTreeChecked, similarDocs, openViewfile, createReport, resetColLayout, lastColLayout, saveColLayout, addToGraph, navigateToParentChat, pinFirstColumns, similarImagesCurrent, similarImagesExternal, + exportTree, exportTreeChecked, similarDocs, similarDocsAI, openViewfile, createReport, resetColLayout, lastColLayout, saveColLayout, addToGraph, navigateToParentChat, pinFirstColumns, similarImagesCurrent, similarImagesExternal, similarFacesCurrent, similarFacesExternal, toggleTimelineView, uiZoom, catIconSize, savePanelsLayout, loadPanelsLayout; MenuListener menuListener = new MenuListener(this); @@ -290,6 +290,11 @@ public void actionPerformed(ActionEvent e) { similarDocs.setEnabled(indexConfig.isStoreTermVectors()); this.add(similarDocs); + similarDocsAI = new JMenuItem(Messages.getString("MenuClass.FindSimilarDocsAI")); //$NON-NLS-1$ + similarDocsAI.addActionListener(menuListener); + similarDocsAI.setEnabled(SemanticSimilarityFilterActions.isFeatureEnabled()); + this.add(similarDocsAI); + submenu = new JMenu(Messages.getString("MenuClass.FindSimilarImages")); //$NON-NLS-1$ submenu.setEnabled(SimilarImagesFilterActions.isFeatureEnabled()); this.add(submenu); diff --git a/iped-app/src/main/java/iped/app/ui/MenuListener.java b/iped-app/src/main/java/iped/app/ui/MenuListener.java index 03d4560613..95e9b4a345 100644 --- a/iped-app/src/main/java/iped/app/ui/MenuListener.java +++ b/iped-app/src/main/java/iped/app/ui/MenuListener.java @@ -422,6 +422,9 @@ public void actionPerformed(ActionEvent e) { App.get().appletListener.updateFileListing(); } + } else if (e.getSource() == menu.similarDocsAI) { + SemanticSimilarityFilterActions.searchSimilarItems(); + } else if (e.getSource() == menu.openViewfile) { int selIdx = App.get().resultsTable.getSelectedRow(); IItemId itemId = App.get().ipedResult.getItem(App.get().resultsTable.convertRowIndexToModel(selIdx)); diff --git a/iped-app/src/main/java/iped/app/ui/MetadataPanel.java b/iped-app/src/main/java/iped/app/ui/MetadataPanel.java index c062534505..65d82136c6 100644 --- a/iped-app/src/main/java/iped/app/ui/MetadataPanel.java +++ b/iped-app/src/main/java/iped/app/ui/MetadataPanel.java @@ -57,6 +57,7 @@ import iped.engine.search.SimilarFacesSearch; import iped.engine.task.NamedEntityTask; import iped.engine.task.index.IndexItem; +import iped.engine.task.index.IndexTask; import iped.engine.task.regex.RegexTask; import iped.engine.task.similarity.ImageSimilarityTask; import iped.exception.ParseException; @@ -208,7 +209,7 @@ private void updateProps() { fields = fields.stream().map(f -> LocalizedProperties.getLocalizedField(f)).collect(Collectors.toList()); Collections.sort(fields, StringUtil.getIgnoreCaseComparator()); for (String f : fields) { - if (f.equals(ResultTableModel.BOOKMARK_COL) || f.equals(ResultTableModel.SCORE_COL) || f.startsWith(ImageSimilarityTask.IMAGE_FEATURES) || f.startsWith(SimilarFacesSearch.FACE_FEATURES)) + if (f.equals(ResultTableModel.BOOKMARK_COL) || f.equals(ResultTableModel.SCORE_COL) || f.startsWith(ImageSimilarityTask.IMAGE_FEATURES) || f.startsWith(SimilarFacesSearch.FACE_FEATURES) || f.equalsIgnoreCase(IndexTask.CONTENT_EMBEDDING) || f.equalsIgnoreCase(IndexTask.HAS_EMBEDDING_FRAG)) continue; if (filterStr.isEmpty() || f.toLowerCase().contains(filterStr)) props.addItem(f); diff --git a/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilter.java b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilter.java new file mode 100644 index 0000000000..e4f0730b57 --- /dev/null +++ b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilter.java @@ -0,0 +1,55 @@ +package iped.app.ui; + +import java.util.Set; + +import org.apache.lucene.document.IntPoint; +import org.apache.lucene.search.BooleanClause.Occur; +import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.Query; + +import iped.viewers.api.IQueryFilter; + +/** + * A Lucene query filter that matches items whose IPED ID is in the provided + * set of item-IDs (as returned by the semantic-similarity RAG search). + * + * Item IDs are stored in the index as integers in the {@code BasicProps.ID} + * IntPoint field. This filter builds a BooleanQuery of exact IntPoint matches. + */ +public class SemanticSimilarityFilter implements IQueryFilter { + + private final Set itemIds; + private Query cachedQuery; + + public SemanticSimilarityFilter(Set itemIds) { + this.itemIds = itemIds; + } + + @Override + public Query getQuery() { + if (cachedQuery != null) { + return cachedQuery; + } + if (itemIds == null || itemIds.isEmpty()) { + return null; + } + + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + for (String idStr : itemIds) { + try { + int id = Integer.parseInt(idStr); + builder.add(IntPoint.newExactQuery(iped.properties.BasicProps.ID, id), Occur.SHOULD); + } catch (NumberFormatException ignore) { + // skip malformed IDs + } + } + builder.setMinimumNumberShouldMatch(1); + cachedQuery = builder.build(); + return cachedQuery; + } + + @Override + public String toString() { + return Messages.get("FilterValue.SemanticSimilarity"); + } +} diff --git a/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterActions.java b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterActions.java new file mode 100644 index 0000000000..1a9f076722 --- /dev/null +++ b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterActions.java @@ -0,0 +1,120 @@ +package iped.app.ui; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import javax.swing.JOptionPane; +import javax.swing.SwingWorker; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import iped.data.IItem; +import iped.data.IItemId; +import iped.engine.rag.RAGService; + +/** + * Static action handler for the "Find Similar Evidence (AI)" context menu + * item. This class: + *
    + *
  1. Retrieves the average embedding vector of all text fragments belonging + * to the currently selected item from OpenSearch.
    2. + *
  2. Performs a KNN vector search to find semantically similar items.
    3. + *
  3. Activates {@link SemanticSimilarityFilterer} with the returned item IDs, + * which causes the result table to be refreshed.
    4. + *
+ */ +public class SemanticSimilarityFilterActions { + + private static final Logger LOGGER = LoggerFactory.getLogger(SemanticSimilarityFilterActions.class); + + /** Clears the current semantic-similarity filter and refreshes the table. */ + public static void clear() { + App app = App.get(); + app.semanticSimilarityFilterer.clearFilter(); + app.appletListener.updateFileListing(); + } + + /** + * Starts an async search for items semantically similar to the currently + * highlighted item and updates the result table when done. + */ + public static void searchSimilarItems() { + App app = App.get(); + + int selIdx = app.resultsTable.getSelectedRow(); + if (selIdx == -1) { + return; + } + + IItemId itemId = app.ipedResult.getItem(app.resultsTable.convertRowIndexToModel(selIdx)); + if (itemId == null) { + return; + } + + IItem item = app.appCase.getItemByItemId(itemId); + if (item == null) { + return; + } + + RAGService ragService = RAGService.getInstance(); + if (ragService == null || !ragService.getConfig().isEnabled()) { + JOptionPane.showMessageDialog( + app, + Messages.getString("RAGAssistant.NotEnabled"), + Messages.getString("RAGAssistant.ErrorTitle"), + JOptionPane.WARNING_MESSAGE); + return; + } + + // Use the string representation of the item's IPED integer ID as the parent ID. + // The RAGService.getAverageEmbeddingForItem expects the parent ID stored in + // OpenSearch's document_content.parent field (which is the IPED item integer ID). + String parentId = String.valueOf(item.getId()); + + new SwingWorker, Void>() { + @Override + protected Set doInBackground() throws Exception { + float[] avgVector = ragService.getAverageEmbeddingForItem(parentId); + if (avgVector == null) { + return null; + } + List similarIds = ragService.findSimilarDocumentIds(avgVector, 200); + return new HashSet<>(similarIds); + } + + @Override + protected void done() { + try { + Set similarIds = get(); + if (similarIds == null || similarIds.isEmpty()) { + JOptionPane.showMessageDialog( + app, + "No semantically similar items were found for the selected evidence.", + Messages.getString("RAGAssistant.ErrorTitle"), + JOptionPane.INFORMATION_MESSAGE); + return; + } + app.semanticSimilarityFilterer.setSimilarItemIds(similarIds); + app.appletListener.updateFileListing(); + } catch (Exception e) { + LOGGER.error("Error performing semantic similarity search", e); + JOptionPane.showMessageDialog( + app, + "Error: " + e.getMessage(), + Messages.getString("RAGAssistant.ErrorTitle"), + JOptionPane.ERROR_MESSAGE); + } + } + }.execute(); + } + + /** Returns {@code true} if the RAG module is initialized for this case. */ + public static boolean isFeatureEnabled() { + RAGService rag = RAGService.getInstance(); + return rag != null && rag.getConfig().isEnabled(); + } +} diff --git a/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterer.java b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterer.java new file mode 100644 index 0000000000..8b6d57329c --- /dev/null +++ b/iped-app/src/main/java/iped/app/ui/SemanticSimilarityFilterer.java @@ -0,0 +1,88 @@ +package iped.app.ui; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import org.apache.lucene.search.Query; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import iped.engine.rag.RAGService; +import iped.viewers.api.IFilter; +import iped.viewers.api.IQueryFilterer; + +/** + * Filterer that restricts the result table to items that are semantically + * similar to a reference document, as determined by the RAG module. + * + * The set of matching IPED item-IDs (source-relative integer IDs stored as + * strings) is populated externally by {@code SemanticSimilarityFilterActions} + * and translated into a Lucene id: query here. + */ +public class SemanticSimilarityFilterer implements IQueryFilterer { + + private static final Logger LOGGER = LoggerFactory.getLogger(SemanticSimilarityFilterer.class); + + /** IPED item-IDs (string form) that are semantically similar to the reference. */ + private Set similarItemIds; + private SemanticSimilarityFilter currentFilter; + + // ----------------------------------------------------- public API + + /** + * Activates the semantic-similarity filter with the given set of item-IDs. + * Pass {@code null} to clear. + */ + public void setSimilarItemIds(Set ids) { + this.similarItemIds = ids; + this.currentFilter = null; // force recreation on next getQuery() + } + + public Set getSimilarItemIds() { + return similarItemIds; + } + + public boolean isActive() { + return similarItemIds != null && !similarItemIds.isEmpty(); + } + + // ----------------------------------------------------- IQueryFilterer + + @Override + public List getDefinedFilters() { + if (currentFilter != null) { + List list = new ArrayList<>(); + list.add(currentFilter); + return list; + } + return null; + } + + @Override + public boolean hasFilters() { + return isActive(); + } + + @Override + public boolean hasFiltersApplied() { + return false; // colour-coding handled externally + } + + @Override + public void clearFilter() { + similarItemIds = null; + currentFilter = null; + } + + @Override + public Query getQuery() { + if (!isActive()) { + return null; + } + if (currentFilter == null) { + currentFilter = new SemanticSimilarityFilter(similarItemIds); + } + return currentFilter.getQuery(); + } +} diff --git a/iped-app/src/main/java/iped/app/ui/TreeListener.java b/iped-app/src/main/java/iped/app/ui/TreeListener.java index 427f14d19c..aad61cedc0 100644 --- a/iped-app/src/main/java/iped/app/ui/TreeListener.java +++ b/iped-app/src/main/java/iped/app/ui/TreeListener.java @@ -81,7 +81,13 @@ public void valueChanged(TreeSelectionEvent evt) { rootSelected = false; for (TreePath path : selection) { - if (((Node) path.getLastPathComponent()).docId == -1) { + Object last = path.getLastPathComponent(); + if (last instanceof Node) { + if (((Node) last).docId == -1) { + rootSelected = true; + break; + } + } else { rootSelected = true; break; } @@ -96,7 +102,11 @@ public void valueChanged(TreeSelectionEvent evt) { BooleanQuery.Builder recursiveQueryBuilder = new BooleanQuery.Builder(); for (TreePath path : selection) { - Document doc = ((Node) path.getLastPathComponent()).getDoc(); + Object last = path.getLastPathComponent(); + if (!(last instanceof Node)) { + continue; + } + Document doc = ((Node) last).getDoc(); String parentId = doc.get(IndexItem.ID); diff --git a/iped-app/src/main/java/iped/app/ui/UICaseDataLoader.java b/iped-app/src/main/java/iped/app/ui/UICaseDataLoader.java index 428db9cd18..8720cfb4ce 100644 --- a/iped-app/src/main/java/iped/app/ui/UICaseDataLoader.java +++ b/iped-app/src/main/java/iped/app/ui/UICaseDataLoader.java @@ -33,6 +33,7 @@ import iped.app.ui.ai.AIFiltersLoader; import iped.app.ui.columns.ColumnsManagerUI; import iped.engine.config.ConfigurationManager; +import iped.engine.rag.RAGService; import iped.engine.core.EvidenceStatus; import iped.engine.core.Manager; import iped.engine.data.IPEDMultiSource; @@ -90,11 +91,36 @@ protected Void doInBackground() { singleCase = new IPEDSource(App.get().casesPathFile, manager.getIndexWriter()); App.get().appCase = new IPEDMultiSource(singleCase); - } else + } else { App.get().appCase = new IPEDMultiSource(App.get().casesPathFile); + } checkIfProcessingFinished(App.get().appCase); + // Initialize RAGService for App UI if RAGConfig is available + try { + iped.engine.config.RAGConfig ragConfig = ConfigurationManager.get().findObject(iped.engine.config.RAGConfig.class); + if (ragConfig != null) { + java.io.File caseDir = App.get().isMultiCase ? App.get().casesPathFile : App.get().appCase.getAtomicSources().get(0).getCaseDir(); + RAGService.initialize(caseDir, ragConfig); + + // Inject the Lucene IndexSearcher so RAGService can perform local KNN/lexical searches. + RAGService ragService = RAGService.getInstance(); + if (ragService != null && "lucene".equalsIgnoreCase(ragConfig.getVectorStoreMode()) && App.get().appCase != null) { + try { + org.apache.lucene.search.IndexSearcher luceneSearcher = App.get().appCase.getSearcher(); + ragService.setLuceneSearcher(luceneSearcher); + ragService.setLuceneAnalyzer(App.get().appCase.getAnalyzer()); + } catch (Exception e) { + LOGGER.warn("Could not inject Lucene IndexSearcher into RAGService in UICaseDataLoader.", e); + } + } + } + } catch (Throwable t) { + LOGGER.error("Error initializing RAGService in UICaseDataLoader", t); + } + + App.get().appCase.checkImagePaths(); App.get().appCase.getMultiBookmarks().addSelectionListener(App.get().getViewerController().getHtmlLinkViewer()); App.get().getViewerController().notifyAppLoaded(); @@ -135,8 +161,8 @@ protected Void doInBackground() { } treeModel = new TreeViewModel(); - } catch (Throwable e) { + LOGGER.error("Error loading case in doInBackground", e); e.printStackTrace(); showErrorDialog(e); } @@ -203,12 +229,25 @@ public void run() { @Override public void done() { + try { + get(); + } catch (Throwable t) { + System.err.println("DIAGNOSTIC: Exception thrown by doInBackground():"); + t.printStackTrace(); + } try { CategoryTreeModel.install(); AIFiltersLoader.load(); App.get().filterManager.loadFilters(); BookmarksController.get().updateUIandHistory(); + try { + App.get().initRAGFeatures(); + } catch (Throwable t) { + LOGGER.error("Error showing RAG UI features", t); + } + + App.get().tree.setModel(treeModel); App.get().tree.setLargeModel(true); App.get().tree.setCellRenderer(new TreeCellRenderer()); diff --git a/iped-app/src/main/java/iped/app/ui/rag/RAGAssistantPanel.java b/iped-app/src/main/java/iped/app/ui/rag/RAGAssistantPanel.java new file mode 100644 index 0000000000..88bd4a245d --- /dev/null +++ b/iped-app/src/main/java/iped/app/ui/rag/RAGAssistantPanel.java @@ -0,0 +1,1266 @@ +package iped.app.ui.rag; + +/** + * UI Panel for the IPED AI Assistant (RAG Chatbot), providing multi-turn chat, + * evidence source navigation, history management, and context window warnings. + * + * @author Rui Sant'Ana Junior + */ +import java.awt.BorderLayout; +import java.awt.Color; +import java.awt.Dimension; +import java.awt.FlowLayout; +import java.awt.event.MouseAdapter; +import java.awt.event.MouseEvent; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CancellationException; + +import javax.swing.BorderFactory; +import javax.swing.BoxLayout; +import javax.swing.DefaultListModel; +import javax.swing.JButton; +import javax.swing.JCheckBox; +import javax.swing.JEditorPane; +import javax.swing.JList; +import javax.swing.JOptionPane; +import javax.swing.JPanel; +import javax.swing.JScrollPane; +import javax.swing.JSplitPane; +import javax.swing.JTable; +import javax.swing.JTextArea; +import javax.swing.ListSelectionModel; +import javax.swing.SwingUtilities; +import javax.swing.SwingWorker; +import javax.swing.UIManager; +import javax.swing.table.AbstractTableModel; +import javax.swing.table.DefaultTableColumnModel; +import javax.swing.table.TableColumn; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import iped.app.ui.App; +import iped.app.ui.FileProcessor; +import iped.app.ui.Messages; +import iped.data.IIPEDSource; +import iped.data.IItemId; +import iped.engine.data.IPEDMultiSource; +import iped.engine.data.IPEDSource; +import iped.engine.data.ItemId; +import iped.engine.rag.RAGService; +import iped.engine.rag.RAGService.RAGSourceDoc; + +/** + * Swing docking panel for the RAG (Retrieval-Augmented Generation) AI + * Assistant. + * + * Layout (vertical split inside a horizontal split): + * + * 
+ * +------------------------------------+--------------------------------------------+
+ * |  [History list ??? left panel]       |  [Question text area + Ask / Cancel btns]  |
+ * |                                    |--------------------------------------------|
+ * |                                    |  [HTML response viewer ??? centre-right]      |
+ * |                                    |--------------------------------------------|
+ * |                                    |  [Sources table ??? bottom-right]             |
+ * +------------------------------------+--------------------------------------------+
+ *
+ * + * Double-clicking a row in the sources table calls + * {@code new FileProcessor(luceneId, false).execute()} to open the evidence + * item in the standard IPED viewer. + */ +public class RAGAssistantPanel extends JPanel { + + private static final long serialVersionUID = 1L; + private static final Logger LOGGER = LoggerFactory.getLogger(RAGAssistantPanel.class); + + // ------------------------------------------------------------------ widgets + private final DefaultListModel historyListModel = new DefaultListModel<>(); + private final JList historyList = new JList<>(historyListModel); + private final JTextArea questionArea; + private final JButton askButton; + private final JButton cancelButton; + private final JCheckBox chkScopeOnly; + private final JEditorPane responsePane; + private final SourcesTableModel sourcesModel = new SourcesTableModel(); + private final JTable sourcesTable = new JTable(sourcesModel); + + // ------------------------------------------------------------------ state + private SwingWorker currentWorker; + private final List sessions = new ArrayList<>(); + private ChatSession activeSession; + + // ---------------------------------------------------------------- constructor + + public RAGAssistantPanel() { + super(new BorderLayout()); + + // ---------- LEFT PANEL: chat history -------------------------------- + historyList.setSelectionMode(ListSelectionModel.SINGLE_SELECTION); + historyList.setFont(UIManager.getFont("Label.font")); + JScrollPane historyScroll = new JScrollPane(historyList); + historyScroll.setBorder(null); + + JButton newChatBtn = new JButton("+ " + Messages.getString("RAGAssistant.NewChat")); + newChatBtn.setFont(UIManager.getFont("Button.font")); + newChatBtn.addActionListener(e -> startNewSession()); + + JPanel leftPanel = new JPanel(new BorderLayout(5, 5)); + leftPanel.setBorder(BorderFactory.createTitledBorder( + Messages.getString("RAGAssistant.History"))); + leftPanel.add(newChatBtn, BorderLayout.NORTH); + leftPanel.add(historyScroll, BorderLayout.CENTER); + leftPanel.setPreferredSize(new Dimension(220, 0)); + + // ---------- TOP-RIGHT: question bar --------------------------------- + questionArea = new JTextArea(3, 40); + questionArea.setLineWrap(true); + questionArea.setWrapStyleWord(true); + questionArea.setFont(UIManager.getFont("TextField.font")); + String placeholder = Messages.getString("RAGAssistant.Question"); + questionArea.setText(placeholder); + questionArea.setForeground(Color.GRAY); + questionArea.addFocusListener(new java.awt.event.FocusAdapter() { + @Override + public void focusGained(java.awt.event.FocusEvent e) { + if (questionArea.getText().equals(placeholder)) { + questionArea.setText(""); + questionArea.setForeground(UIManager.getColor("TextField.foreground")); + } + } + + @Override + public void focusLost(java.awt.event.FocusEvent e) { + if (questionArea.getText().trim().isEmpty()) { + questionArea.setForeground(Color.GRAY); + questionArea.setText(placeholder); + } + } + }); + questionArea.addKeyListener(new java.awt.event.KeyAdapter() { + @Override + public void keyPressed(java.awt.event.KeyEvent e) { + if (e.getKeyCode() == java.awt.event.KeyEvent.VK_ENTER) { + if (e.isShiftDown()) { + questionArea.insert("\n", questionArea.getCaretPosition()); + } else { + e.consume(); + onAsk(); + } + } + } + }); + JScrollPane questionScroll = new JScrollPane(questionArea); + + askButton = new JButton(Messages.getString("RAGAssistant.Ask")); + cancelButton = new JButton(Messages.getString("RAGAssistant.Cancel")); + cancelButton.setEnabled(false); + + chkScopeOnly = new JCheckBox(getMessageSafely("RAGAssistant.ScopeSelected", "Restringir ?? sele????o")); + chkScopeOnly.setToolTipText(getMessageSafely("RAGAssistant.ScopeSelectedTooltip", + "Restringe a busca RAG apenas aos itens marcados, selecionados na tabela ou vis??veis nos filtros atuais.")); + chkScopeOnly.setFont(UIManager.getFont("Label.font")); + + askButton.addActionListener(e -> onAsk()); + cancelButton.addActionListener(e -> onCancel()); + + JPanel actionBtns = new JPanel(new FlowLayout(FlowLayout.RIGHT, 4, 0)); + actionBtns.add(askButton); + actionBtns.add(cancelButton); + + JPanel chkHolder = new JPanel(new FlowLayout(FlowLayout.RIGHT, 0, 4)); + chkHolder.add(chkScopeOnly); + + JPanel btnPanel = new JPanel(); + btnPanel.setLayout(new BoxLayout(btnPanel, BoxLayout.Y_AXIS)); + btnPanel.add(actionBtns); + btnPanel.add(chkHolder); + + JPanel topRight = new JPanel(new BorderLayout(4, 4)); + topRight.setBorder(BorderFactory.createEmptyBorder(4, 4, 4, 4)); + topRight.add(questionScroll, BorderLayout.CENTER); + topRight.add(btnPanel, BorderLayout.EAST); + topRight.setPreferredSize(new Dimension(0, 110)); + topRight.setMaximumSize(new Dimension(Integer.MAX_VALUE, 110)); + + // ---------- CENTRE-RIGHT: response area ----------------------------- + responsePane = new JEditorPane("text/html", ""); + responsePane.setEditable(false); + responsePane.setBackground(UIManager.getColor("TextArea.background")); + responsePane.addHyperlinkListener(e -> { + if (e.getEventType() == javax.swing.event.HyperlinkEvent.EventType.ACTIVATED) { + String description = e.getDescription(); + if (description != null && description.startsWith("sources:")) { + try { + int entryIndex = Integer.parseInt(description.substring(8)); + if (activeSession != null && entryIndex >= 0 && entryIndex < activeSession.entries.size()) { + HistoryEntry entry = activeSession.entries.get(entryIndex); + if (entry != null && entry.displaySources != null) { + sourcesModel.setRows(new ArrayList<>(entry.displaySources)); + } + } + } catch (NumberFormatException nfe) { + // ignore + } + } + } + }); + JScrollPane responseScroll = new JScrollPane(responsePane); + responseScroll.setBorder(BorderFactory.createEmptyBorder(0, 4, 0, 4)); + + // ---------- BOTTOM-RIGHT: sources table ----------------------------- + configureSourcesTable(); + JScrollPane sourcesScroll = new JScrollPane(sourcesTable); + sourcesScroll.setBorder(BorderFactory.createTitledBorder( + Messages.getString("RAGAssistant.Sources"))); + sourcesScroll.setPreferredSize(new Dimension(0, 140)); + + // ---------- RIGHT SPLIT: response + sources ------------------------- + JPanel rightContent = new JPanel(); + rightContent.setLayout(new BoxLayout(rightContent, BoxLayout.Y_AXIS)); + rightContent.add(topRight); + rightContent.add(responseScroll); + rightContent.add(sourcesScroll); + + JSplitPane rightSplit = new JSplitPane(JSplitPane.VERTICAL_SPLIT, responseScroll, sourcesScroll); + rightSplit.setResizeWeight(0.65); + rightSplit.setBorder(null); + + JPanel rightPanel = new JPanel(new BorderLayout()); + rightPanel.add(topRight, BorderLayout.NORTH); + rightPanel.add(rightSplit, BorderLayout.CENTER); + + // ---------- MAIN SPLIT: history + right ----------------------------- + JSplitPane mainSplit = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, leftPanel, rightPanel); + mainSplit.setResizeWeight(0.0); + mainSplit.setDividerLocation(220); + mainSplit.setBorder(null); + + add(mainSplit, BorderLayout.CENTER); + + // History selection listener ??? added here so that questionArea and responsePane + // (final fields) are guaranteed to be initialized before the lambda captures + // them. + historyList.addListSelectionListener(e -> { + if (!e.getValueIsAdjusting()) { + int idx = historyList.getSelectedIndex(); + if (idx >= 0 && idx < sessions.size()) { + activeSession = sessions.get(idx); + + // Render the conversation thread for this session + String conversationHtml = renderConversation(activeSession); + responsePane.setContentType("text/html"); + responsePane.setText(conversationHtml); + + // Scroll to the bottom of the conversation + SwingUtilities + .invokeLater(() -> responsePane.setCaretPosition(responsePane.getDocument().getLength())); + + // Restore sources of the last entry in the session (if any) + if (!activeSession.entries.isEmpty()) { + HistoryEntry lastEntry = activeSession.entries.get(activeSession.entries.size() - 1); + sourcesModel.setRows(lastEntry.displaySources); + } else { + sourcesModel.setRows(new ArrayList<>()); + } + + // Keep the question area clear for the next question (chat continuation) + String placeholderMsg = Messages.getString("RAGAssistant.Question"); + if (questionArea.hasFocus()) { + questionArea.setText(""); + questionArea.setForeground(UIManager.getColor("TextField.foreground")); + } else { + questionArea.setForeground(Color.GRAY); + questionArea.setText(placeholderMsg); + } + } + } + }); + + // Popup menu for history list operations (rename / delete) + javax.swing.JPopupMenu popupMenu = new javax.swing.JPopupMenu(); + javax.swing.JMenuItem renameItem = new javax.swing.JMenuItem(Messages.getString("RAGAssistant.RenameChat")); + renameItem.addActionListener(e -> renameSelectedHistory()); + popupMenu.add(renameItem); + javax.swing.JMenuItem deleteItem = new javax.swing.JMenuItem(Messages.getString("RAGAssistant.DeleteChat")); + deleteItem.addActionListener(e -> deleteSelectedHistory()); + popupMenu.add(deleteItem); + + historyList.addMouseListener(new MouseAdapter() { + @Override + public void mousePressed(MouseEvent e) { + showPopup(e); + } + + @Override + public void mouseReleased(MouseEvent e) { + showPopup(e); + } + + private void showPopup(MouseEvent e) { + if (e.isPopupTrigger()) { + int index = historyList.locationToIndex(e.getPoint()); + if (index >= 0) { + historyList.setSelectedIndex(index); + popupMenu.show(e.getComponent(), e.getX(), e.getY()); + } + } + } + }); + + // Key listener for Delete key on history list + historyList.addKeyListener(new java.awt.event.KeyAdapter() { + @Override + public void keyPressed(java.awt.event.KeyEvent e) { + if (e.getKeyCode() == java.awt.event.KeyEvent.VK_DELETE) { + deleteSelectedHistory(); + } + } + }); + + loadHistoryFromDatabase(); + } + + // ---------------------------------------------------------------- table setup + + private void configureSourcesTable() { + sourcesTable.setAutoResizeMode(JTable.AUTO_RESIZE_OFF); + sourcesTable.setSelectionMode(ListSelectionModel.SINGLE_SELECTION); + sourcesTable.setFillsViewportHeight(true); + sourcesTable.setFont(UIManager.getFont("Label.font")); + + DefaultTableColumnModel colModel = new DefaultTableColumnModel(); + String[] headers = { + Messages.getString("RAGAssistant.Header.Id"), + Messages.getString("RAGAssistant.Header.Name"), + Messages.getString("RAGAssistant.Header.Score"), + Messages.getString("RAGAssistant.Header.Path") + }; + int[] widths = { 50, 240, 70, 500 }; + for (int i = 0; i < headers.length; i++) { + TableColumn col = new TableColumn(i, widths[i]); + col.setHeaderValue(headers[i]); + col.setPreferredWidth(widths[i]); + colModel.addColumn(col); + } + sourcesTable.setColumnModel(colModel); + + // Double-click to open the item in IPED viewer + sourcesTable.addMouseListener(new MouseAdapter() { + @Override + public void mouseClicked(MouseEvent e) { + if (e.getClickCount() == 2) { + int row = sourcesTable.rowAtPoint(e.getPoint()); + if (row >= 0) { + openSourceItem(row); + } + } + } + }); + } + + // ---------------------------------------------------------------- actions + + private void onAsk() { + String question = questionArea.getText().trim(); + String placeholder = Messages.getString("RAGAssistant.Question"); + if (question.isEmpty() || question.equals(placeholder)) { + return; + } + + RAGService rag = RAGService.getInstance(); + if (rag == null || !rag.getConfig().isEnabled()) { + JOptionPane.showMessageDialog( + App.get(), + Messages.getString("RAGAssistant.NotEnabled"), + Messages.getString("RAGAssistant.ErrorTitle"), + JOptionPane.WARNING_MESSAGE); + return; + } + + if (activeSession == null) { + startNewSession(); + } + + final String q = question; + final String currentSessionId = activeSession.sessionId; + final boolean isFirstQuestion = activeSession.entries.isEmpty(); + + // Assemble chat history context (last 5 turns) from active session + final List historyTurns = new ArrayList<>(); + int entryCount = activeSession.entries.size(); + if (entryCount > 0) { + int startIdx = Math.max(0, entryCount - 5); + for (int i = startIdx; i < entryCount; i++) { + HistoryEntry he = activeSession.entries.get(i); + String cleanAns = he.rawAnswer != null ? he.rawAnswer : stripHtml(he.answerHtml); + historyTurns.add(new RAGService.HistoryTurn(he.question, cleanAns)); + } + } + + // Display current conversation thread + the new question + Thinking... + StringBuilder tempHtml = new StringBuilder(); + tempHtml.append(""); + if (entryCount > 0) { + tempHtml.append(extractBodyContent(renderConversation(activeSession))); + tempHtml.append("
"); + } + // User Question + tempHtml.append("
"); + tempHtml.append("").append(Messages.getString("RAGAssistant.YouLabel")).append(": ") + .append(markdownToHtml(q)); + tempHtml.append("
"); + // Thinking... + tempHtml.append("
"); + tempHtml.append("").append(Messages.getString("RAGAssistant.AssistantLabel")).append(": ") + .append(Messages.getString("RAGAssistant.Thinking")).append(""); + tempHtml.append("
"); + tempHtml.append(""); + + setUIBusy(true); + responsePane.setContentType("text/html"); + responsePane.setText(tempHtml.toString()); + SwingUtilities.invokeLater(() -> responsePane.setCaretPosition(responsePane.getDocument().getLength())); + sourcesModel.setRows(new ArrayList<>()); + + // Collect allowed item IDs here, on the EDT, before handing off to the worker. + // This guarantees that all Swing component reads happen on the Event Dispatch + // Thread. + final Set scopeAllowedIds = collectAllowedParentIds(); + + currentWorker = new SwingWorker() { + private List foundDocs = new ArrayList<>(); + private int omittedChunksCount = 0; + private boolean contextLimitReached = false; + private int contextWindowValue = 4096; + private int totalRetrievedChunksCount = 0; + + @Override + protected String doInBackground() throws Exception { + // 1. Prepare contextualized search query for multi-turn questions + String searchQuery = extractSearchQueryWithHistory(q, historyTurns); + + // 2. Embed the question + float[] qvec = rag.getEmbedding(searchQuery); + if (qvec == null) { + throw new IOException("Could not generate embedding for the question."); + } + + // 3. Hybrid search using contextualized search query and scope filter + List docs = rag.performHybridSearch(searchQuery, qvec, scopeAllowedIds); + foundDocs.addAll(docs); + + if (docs.isEmpty()) { + return Messages.getString("RAGAssistant.NoContextFound"); + } + + // 3. Build context + StringBuilder ctx = new StringBuilder(); + int docIndex = 1; + int omittedCount = 0; + boolean limitReached = false; + + int contextWindow = rag.getConfig().getLlmContextWindow(); + contextWindowValue = contextWindow; + int maxChars; + if ("local".equalsIgnoreCase(rag.getConfig().getLlmProvider())) { + int reservedTokens = Math.min(1024, contextWindow / 2); + // 1 token ??? 3.8 characters in Portuguese/English. Multiplier of 3.8 accurately + // reflects the character budget without premature truncation. + maxChars = (int) ((contextWindow - reservedTokens) * 3.8); + } else { + maxChars = Integer.MAX_VALUE; + } + + // Base length of prompt (system prompt + user prompt boilerplate + question) is + // roughly 2500 chars + int baseLength = 2500 + (q != null ? q.length() : 0); + totalRetrievedChunksCount = docs.size(); + + for (RAGSourceDoc doc : docs) { + if (doc.content != null && !doc.content.isEmpty()) { + String docName = getDocName(doc); + StringBuilder tempChunk = new StringBuilder(); + tempChunk.append("[").append(doc.itemId).append("]\n"); + if (docName != null) { + tempChunk.append("Nome do Arquivo: ").append(docName).append("\n"); + } + tempChunk.append("Score H??brido: ") + .append(String.format(java.util.Locale.US, "%.3f", doc.score)).append("\n"); + tempChunk.append("---\n"); + tempChunk.append(doc.content).append("\n\n========================================\n\n"); + + if (baseLength + ctx.length() + tempChunk.length() > maxChars) { + limitReached = true; + omittedCount++; + continue; + } + + ctx.append(tempChunk); + docIndex++; + } + } + + omittedChunksCount = omittedCount; + contextLimitReached = limitReached; + + // 4. Ask the LLM + return rag.getLlmProvider().generateAnswer(q, ctx.toString(), historyTurns); + } + + @Override + protected void done() { + setUIBusy(false); + try { + String answer = get(); + String htmlAnswer = markdownToHtml(answer); + if (contextLimitReached) { + String formattedWarning = String.format( + Messages.getString("RAGAssistant.ContextWindowWarning"), + contextWindowValue, omittedChunksCount, totalRetrievedChunksCount); + htmlAnswer = htmlAnswer + + "
" + + formattedWarning + "
"; + } + String html = "" + + htmlAnswer + ""; + + // Determine which sources to display (filtered) vs. full context (unchanged) + List displayDocs = filterTopSources(foundDocs, q, answer); + + // Save to database + long dbId = -1; + if (rag != null) { + dbId = rag.saveHistoryEntry(currentSessionId, q, html, displayDocs); + } + + // Add to active session + HistoryEntry entry = new HistoryEntry(dbId, q, html, answer, displayDocs); + activeSession.entries.add(entry); + + if (isFirstQuestion) { + if (activeSession.title == null + || activeSession.title.equals(Messages.getString("RAGAssistant.NewChat"))) { + String title = q.length() > 30 ? q.substring(0, 27) + "..." : q; + activeSession.title = title; + } + sessions.add(activeSession); + historyListModel.addElement(activeSession.title); + // Selecting triggers automatic list selection listener logic + int newIdx = sessions.size() - 1; + historyList.setSelectedIndex(newIdx); + } else { + // Just re-render conversation thread + responsePane.setContentType("text/html"); + responsePane.setText(renderConversation(activeSession)); + SwingUtilities.invokeLater( + () -> responsePane.setCaretPosition(responsePane.getDocument().getLength())); + sourcesModel.setRows(new ArrayList<>(displayDocs)); + } + + // Clear question area + String placeholderMsg = Messages.getString("RAGAssistant.Question"); + if (questionArea.hasFocus()) { + questionArea.setText(""); + questionArea.setForeground(UIManager.getColor("TextField.foreground")); + } else { + questionArea.setForeground(Color.GRAY); + questionArea.setText(placeholderMsg); + } + + } catch (CancellationException ce) { + // Restore conversation thread without thinking state + responsePane.setContentType("text/html"); + responsePane.setText(renderConversation(activeSession)); + } catch (Exception ex) { + LOGGER.error("RAG answer generation error", ex); + // Append error message to thread + StringBuilder errHtml = new StringBuilder(); + errHtml.append(""); + if (!activeSession.entries.isEmpty()) { + errHtml.append(extractBodyContent(renderConversation(activeSession))); + errHtml.append("
"); + } + errHtml.append( + "
"); + errHtml.append("").append(Messages.getString("RAGAssistant.YouLabel")).append(": ") + .append(markdownToHtml(q)); + errHtml.append("
"); + errHtml.append("
"); + errHtml.append("").append(Messages.getString("RAGAssistant.ErrorTitle")).append(":
") + .append(ex.getMessage()); + errHtml.append("
"); + errHtml.append(""); + responsePane.setContentType("text/html"); + responsePane.setText(errHtml.toString()); + } + } + }; + currentWorker.execute(); + } + + private void onCancel() { + if (currentWorker != null && !currentWorker.isDone()) { + currentWorker.cancel(true); + } + } + + private void setUIBusy(boolean busy) { + askButton.setEnabled(!busy); + cancelButton.setEnabled(busy); + questionArea.setEnabled(!busy); + } + + private void openSourceItem(int row) { + RAGSourceDoc doc = sourcesModel.getRow(row); + if (doc == null || doc.itemId == null) + return; + try { + int ipedId = Integer.parseInt(doc.itemId); + // Search for the lucene docId matching this IPED itemId + App app = App.get(); + if (app.appCase instanceof IPEDMultiSource) { + for (IPEDSource src : ((IPEDMultiSource) app.appCase).getAtomicSources()) { + if (ipedId <= src.getLastId()) { + IItemId itemId = new ItemId(src.getSourceId(), ipedId); + int luceneId = app.appCase.getLuceneId(itemId); + if (luceneId >= 0) { + new FileProcessor(luceneId, false).execute(); + return; + } + } + } + } + } catch (NumberFormatException nfe) { + // ignore + } + } + + private void loadHistoryFromDatabase() { + RAGService rag = RAGService.getInstance(); + if (rag != null && rag.getConfig().isEnabled()) { + Map customTitles = rag.loadSessionTitles(); + List records = rag.loadHistoryEntries(); + Map sessionMap = new java.util.LinkedHashMap<>(); + for (RAGService.HistoryEntryRecord rec : records) { + String sid = rec.sessionId; + if (sid == null || sid.trim().isEmpty()) { + sid = "legacy_" + rec.id; + } + + HistoryEntry entry = new HistoryEntry(rec.id, rec.question, rec.answerHtml, null, rec.displaySources); + + ChatSession session = sessionMap.get(sid); + if (session == null) { + String title = customTitles.get(sid); + if (title == null || title.trim().isEmpty()) { + title = rec.question.length() > 30 ? rec.question.substring(0, 27) + "..." : rec.question; + } + session = new ChatSession(sid, title); + sessionMap.put(sid, session); + sessions.add(session); + } + session.entries.add(entry); + } + + for (ChatSession session : sessions) { + historyListModel.addElement(session.title); + } + } + + if (sessions.isEmpty()) { + startNewSession(); + } else { + historyList.setSelectedIndex(sessions.size() - 1); + } + } + + private void startNewSession() { + String newSid = java.util.UUID.randomUUID().toString(); + activeSession = new ChatSession(newSid, Messages.getString("RAGAssistant.NewChat")); + + responsePane.setContentType("text/html"); + responsePane.setText(""); + sourcesModel.setRows(new ArrayList<>()); + + historyList.clearSelection(); + + String placeholderMsg = Messages.getString("RAGAssistant.Question"); + if (questionArea.hasFocus()) { + questionArea.setText(""); + questionArea.setForeground(UIManager.getColor("TextField.foreground")); + } else { + questionArea.setForeground(Color.GRAY); + questionArea.setText(placeholderMsg); + } + questionArea.requestFocusInWindow(); + } + + private void renameSelectedHistory() { + int idx = historyList.getSelectedIndex(); + if (idx >= 0 && idx < sessions.size()) { + ChatSession session = sessions.get(idx); + String newTitle = JOptionPane.showInputDialog( + this, + Messages.getString("RAGAssistant.RenameChatPrompt"), + session.title); + if (newTitle != null && !newTitle.trim().isEmpty() && !newTitle.equals(session.title)) { + newTitle = newTitle.trim(); + session.title = newTitle; + historyListModel.set(idx, newTitle); + + if (session.sessionId != null && !session.sessionId.startsWith("legacy_")) { + final String finalTitle = newTitle; + RAGService rag = RAGService.getInstance(); + if (rag != null) { + new SwingWorker() { + @Override + protected Void doInBackground() throws Exception { + rag.updateSessionTitle(session.sessionId, finalTitle); + return null; + } + }.execute(); + } + } + } + } + } + + private void deleteSelectedHistory() { + int idx = historyList.getSelectedIndex(); + if (idx >= 0 && idx < sessions.size()) { + int option = JOptionPane.showConfirmDialog( + this, + Messages.getString("RAGAssistant.ConfirmDeleteChat"), + Messages.getString("RAGAssistant.ConfirmDeleteTitle"), + JOptionPane.YES_NO_OPTION); + if (option == JOptionPane.YES_OPTION) { + ChatSession session = sessions.remove(idx); + historyListModel.remove(idx); + + if (sessions.isEmpty()) { + startNewSession(); + } else { + int newSel = Math.max(0, idx - 1); + historyList.setSelectedIndex(newSel); + } + + if (session.sessionId != null && !session.sessionId.startsWith("legacy_")) { + RAGService rag = RAGService.getInstance(); + if (rag != null) { + new SwingWorker() { + @Override + protected Void doInBackground() throws Exception { + rag.deleteHistorySession(session.sessionId); + return null; + } + }.execute(); + } + } + } + } + } + + private String renderConversation(ChatSession session) { + if (session == null || session.entries.isEmpty()) { + return ""; + } + StringBuilder html = new StringBuilder(); + html.append(""); + + for (int i = 0; i < session.entries.size(); i++) { + HistoryEntry entry = session.entries.get(i); + + // User Question + html.append("
"); + html.append("").append(Messages.getString("RAGAssistant.YouLabel")).append(": ") + .append(markdownToHtml(entry.question)); + html.append("
"); + + // Assistant Answer + html.append("
"); + html.append("").append(Messages.getString("RAGAssistant.AssistantLabel")).append(": ").append( + entry.answerHtml != null ? extractBodyContent(entry.answerHtml) : markdownToHtml(entry.rawAnswer)); + + if (entry.displaySources != null && !entry.displaySources.isEmpty()) { + int count = entry.displaySources.size(); + boolean isPt = "pt".equalsIgnoreCase(iped.localization.LocaleResolver.getLocale().getLanguage()); + String label = isPt + ? (count == 1 ? "1 fonte consultada" : count + " fontes consultadas") + : (count == 1 ? "1 source consulted" : count + " sources consulted"); + String hint = isPt ? "(clique para carregar no grid)" : "(click to load in grid)"; + html.append("
"); + html.append("📁 ") + .append(label).append(" ").append(hint).append(""); + html.append("
"); + } + html.append("
"); + + if (i < session.entries.size() - 1) { + html.append("
"); + } + } + + html.append(""); + return html.toString(); + } + + private String extractBodyContent(String html) { + if (html == null) + return ""; + int bodyStart = html.indexOf("", bodyStart); + if (bodyClose == -1) + return html; + int bodyEnd = html.indexOf(""); + if (bodyEnd == -1) { + return html.substring(bodyClose + 1); + } + return html.substring(bodyClose + 1, bodyEnd); + } + + // ---------------------------------------------------------------- helpers + + private List deduplicateByItemId(List list) { + if (list == null || list.isEmpty()) + return list; + java.util.Map uniqueDocs = new java.util.LinkedHashMap<>(); + for (RAGSourceDoc doc : list) { + if (doc.itemId != null) { + if (!uniqueDocs.containsKey(doc.itemId)) { + uniqueDocs.put(doc.itemId, doc); + } else { + if (doc.score > uniqueDocs.get(doc.itemId).score) { + uniqueDocs.put(doc.itemId, doc); + } + } + } + } + return new ArrayList<>(uniqueDocs.values()); + } + + /** + * Detects whether the LLM answer is a negative response indicating that no + * relevant information was found in the retrieved context. + * + * Uses a broad regex pattern instead of a fixed list of literal strings so + * that variations in phrasing across different LLM models and both PT-BR/EN + * are handled uniformly. + */ + private List filterTopSources(List docs, String query, String answer) { + if (docs == null || docs.isEmpty() || answer == null) { + return new ArrayList<>(); + } + + Set citedIds = new java.util.HashSet<>(); + + // 1. Match bracket citations like [5], [5, 6] + java.util.regex.Pattern pBrackets = java.util.regex.Pattern.compile("\\[([\\d,\\s]+)\\]"); + java.util.regex.Matcher m = pBrackets.matcher(answer); + while (m.find()) { + for (String part : m.group(1).split(",")) { + String id = part.trim(); + if (!id.isEmpty()) + citedIds.add(id); + } + } + + List citedList = new ArrayList<>(); + Set added = new java.util.HashSet<>(); + + // Match docs cited by bracket ID [5] + for (RAGSourceDoc doc : docs) { + if (doc.itemId != null && citedIds.contains(doc.itemId)) { + if (added.add(doc.itemId)) { + citedList.add(doc); + } + } + } + + // 2. Also match docs cited by filename/title in answer text (e.g. + // "PGR-00149505/2023" or "Relotacao.pdf") + String lowerAnswer = answer.toLowerCase(java.util.Locale.ROOT); + for (RAGSourceDoc doc : docs) { + if (doc.itemId != null && !added.contains(doc.itemId)) { + String name = getDocName(doc); + if (name != null && !name.trim().isEmpty()) { + String baseName = name; + int dotIdx = name.lastIndexOf('.'); + if (dotIdx > 0) { + baseName = name.substring(0, dotIdx); + } + if (baseName.length() >= 3 && lowerAnswer.contains(baseName.toLowerCase(java.util.Locale.ROOT))) { + added.add(doc.itemId); + citedList.add(doc); + } + } + } + } + + if (citedList.isEmpty()) { + return deduplicateByItemId(docs); + } + return citedList; + } + + private String extractSearchQueryWithHistory(String question, List historyTurns) { + if (historyTurns == null || historyTurns.isEmpty() || question == null) { + return question; + } + StringBuilder searchStr = new StringBuilder(question); + RAGService.HistoryTurn lastTurn = historyTurns.get(historyTurns.size() - 1); + if (lastTurn != null && lastTurn.question != null) { + String[] tokens = lastTurn.question.split("[^a-zA-Z0-9????????????????????????????????????????????????????]"); + for (String token : tokens) { + String trimmed = token.trim(); + if (trimmed.length() >= 3 && !isStopWord(trimmed.toLowerCase())) { + if (!question.toLowerCase().contains(trimmed.toLowerCase())) { + searchStr.append(" ").append(trimmed); + } + } + } + } + return searchStr.toString(); + } + + private Set collectAllowedParentIds() { + if (chkScopeOnly == null || !chkScopeOnly.isSelected()) { + return null; + } + Set allowedItemIds = new HashSet<>(); + try { + App app = App.get(); + if (app != null && app.appCase != null) { + // 1. Check for checked items in bookmarks + if (app.appCase instanceof IPEDMultiSource) { + for (IPEDSource src : ((IPEDMultiSource) app.appCase).getAtomicSources()) { + int lastId = src.getLastId(); + for (int id = 1; id <= lastId; id++) { + if (src.getBookmarks() != null && src.getBookmarks().isChecked(id)) { + allowedItemIds.add(String.valueOf(id)); + } + } + } + } else if (app.appCase.getBookmarks() != null) { + int lastId = app.appCase.getLastId(); + for (int id = 1; id <= lastId; id++) { + if (app.appCase.getBookmarks().isChecked(id)) { + allowedItemIds.add(String.valueOf(id)); + } + } + } + + // 2. If no items are checked, check for selected/highlighted rows in + // resultsTable + javax.swing.JTable tbl = app.getResultsTable(); + iped.search.IMultiSearchResult res = app.getResults(); + if (allowedItemIds.isEmpty() && tbl != null && res != null) { + int[] selectedRows = tbl.getSelectedRows(); + if (selectedRows != null) { + for (int row : selectedRows) { + int rowModel = tbl.convertRowIndexToModel(row); + IItemId iid = res.getItem(rowModel); + if (iid != null) { + allowedItemIds.add(String.valueOf(iid.getId())); + } + } + } + } + + // 3. If still empty, use all currently filtered items visible in ipedResult + if (allowedItemIds.isEmpty() && res != null) { + int count = res.getLength(); + LOGGER.warn( + "RAG scope filter: no items checked or selected; falling back to all {} currently visible items.", + count); + for (int r = 0; r < count; r++) { + IItemId iid = res.getItem(r); + if (iid != null) { + allowedItemIds.add(String.valueOf(iid.getId())); + } + } + } + } + } catch (Exception e) { + LOGGER.error("Error collecting allowed parent IDs for RAG scope filter", e); + } + return allowedItemIds.isEmpty() ? null : allowedItemIds; + } + + private String getMessageSafely(String key, String defaultValue) { + try { + return Messages.getString(key); + } catch (Exception e) { + return defaultValue; + } + } + + private String getDocName(RAGSourceDoc doc) { + if (doc == null || doc.itemId == null) + return null; + try { + int id = Integer.parseInt(doc.itemId); + App app = App.get(); + if (app != null && app.appCase instanceof IPEDMultiSource) { + for (IPEDSource src : ((IPEDMultiSource) app.appCase).getAtomicSources()) { + if (id <= src.getLastId()) { + IItemId iid = new ItemId(src.getSourceId(), id); + iped.data.IItem item = app.appCase.getItemByItemId(iid); + if (item != null) + return item.getName(); + } + } + } + } catch (Exception ignore) { + } + return null; + } + + private static Set extractNormalizedNumbers(String text) { + Set numbers = new HashSet<>(); + if (text == null) + return numbers; + java.util.regex.Matcher matcher = java.util.regex.Pattern.compile("\\b\\d+[.,\\d]*\\b").matcher(text); + while (matcher.find()) { + String numStr = matcher.group(); + String cleanNum = numStr.replaceAll("[.,]", ""); + if (cleanNum.length() >= 3) { + numbers.add(cleanNum); + } + } + return numbers; + } + + private static boolean isYear(String numStr) { + if (numStr.length() == 4) { + try { + int val = Integer.parseInt(numStr); + return val >= 1700 && val <= 2100; + } catch (NumberFormatException e) { + return false; + } + } + return false; + } + + private static Set tokenizeAndClean(String text) { + Set words = new HashSet<>(); + if (text == null) + return words; + String cleanText = text.replaceAll("<[^>]*>", " "); + // Normalize accents/diacritics + cleanText = java.text.Normalizer.normalize(cleanText, java.text.Normalizer.Form.NFD); + cleanText = cleanText.replaceAll("[\\p{InCombiningDiacriticalMarks}]", ""); + + String[] tokens = cleanText.toLowerCase().split("[^a-zA-Z0-9]"); + for (String t : tokens) { + String trimmed = t.trim(); + if (trimmed.length() >= 3 && !isStopWord(trimmed)) { + words.add(trimmed); + } + } + return words; + } + + /** + * Stop-words used to filter out noise tokens when building the lexical query + * context. + */ + private static final Set STOP_WORDS = new HashSet<>(java.util.Arrays.asList( + // Portuguese function words + "como", "para", "com", "uma", "mais", "este", "esta", "estes", "estas", + "esse", "essa", "esses", "essas", "aquele", "aquela", + "pelo", "pela", "pelos", "pelas", "seus", "suas", + "onde", "quando", "quem", "qual", "quais", + "todo", "toda", "todos", "todas", "tudo", + "sobre", "entre", "sempre", "nunca", "talvez", + "ainda", "muito", "pouco", "tanto", "cada", + "outro", "outra", "outros", "outras", + "mesmo", "mesma", "pode", "podem", + "base", "documento", "documentos", "fornecido", "fornecidos", + "arquivo", "arquivos", "conteudo", "resumir", "listar", "descritos", + "tamanho", "tipo", "existem", "dois", "identificados", "nomes", + "explicitos", "vez", + "nos", "nas", "num", "numa", "mas", "nao", + "sim", "aos", "sua", "seu", "ele", "ela", "eles", "elas", + "isso", "isto", "aquilo", + "ser", "sao", "era", "eram", "seja", "sejam", "foi", "foram", + "ter", "tem", "tinha", "tinham", + "poderia", "poderiam", "deve", "devem", + "primeiro", "segundo", "terceiro", + "bloco", "blocos", "linha", "linhas", + "item", "itens", "id", "ids", + "foto", "fotos", "imagem", "imagens", + "contexto", "contextos", "metadado", "metadados", + // English function words + "out", "this", "that", "with", "from", "have", "been", + "type", "regular", "folder", "name")); + + private static boolean isStopWord(String word) { + return STOP_WORDS.contains(word); + } + + /** + * Minimal markdown-to-HTML conversion for LLM responses: + * bold, italic, code blocks, inline code and newlines. + */ + private static String markdownToHtml(String md) { + if (md == null) + return ""; + String html = md + .replace("&", "&") + .replace("<", "<") + .replace(">", ">"); + // code blocks + html = html.replaceAll("(?s)```[a-z]*\n?(.*?)```", + "
$1
"); + // inline code + html = html.replaceAll("`([^`]+)`", + "$1"); + // bold + html = html.replaceAll("\\*\\*([^*]+)\\*\\*", "$1"); + html = html.replaceAll("__([^_]+)__", "$1"); + // italic + html = html.replaceAll("\\*([^*]+)\\*", "$1"); + html = html.replaceAll("_([^_]+)_", "$1"); + // newlines + html = html.replace("\n", "
"); + return html; + } + + private static String stripHtml(String html) { + if (html == null) + return ""; + String text = html.replaceAll("(?i)", "\n"); + text = text.replaceAll("<[^>]*>", ""); + text = text.replace("<", "<") + .replace(">", ">") + .replace("&", "&"); + return text.trim(); + } + + // ================================================================ inner + // classes + + private static class ChatSession { + final String sessionId; + String title; + final List entries = new ArrayList<>(); + + ChatSession(String sessionId, String title) { + this.sessionId = sessionId; + this.title = title; + } + } + + /** + * Stores a single history entry: the original question text, the rendered + * HTML response and the filtered list of source documents to display. + */ + private static class HistoryEntry { + final Long dbId; + final String question; + final String answerHtml; + final String rawAnswer; + final List displaySources; + + HistoryEntry(Long dbId, String question, String answerHtml, String rawAnswer, + List displaySources) { + this.dbId = dbId; + this.question = question; + this.answerHtml = answerHtml; + this.rawAnswer = rawAnswer; + this.displaySources = displaySources; + } + } + + private static class SourcesTableModel extends AbstractTableModel { + + private static final long serialVersionUID = 1L; + private List rows = new ArrayList<>(); + + void setRows(List rows) { + this.rows = rows == null ? new ArrayList<>() : rows; + fireTableDataChanged(); + } + + RAGSourceDoc getRow(int rowIndex) { + if (rowIndex < 0 || rowIndex >= rows.size()) + return null; + return rows.get(rowIndex); + } + + @Override + public int getRowCount() { + return rows.size(); + } + + @Override + public int getColumnCount() { + return 4; + } + + @Override + public Object getValueAt(int rowIndex, int columnIndex) { + RAGSourceDoc doc = rows.get(rowIndex); + switch (columnIndex) { + case 0: + return doc.itemId; + case 1: { + // Try to resolve the item name from the appCase + try { + int id = Integer.parseInt(doc.itemId); + App app = App.get(); + if (app != null && app.appCase instanceof IPEDMultiSource) { + for (IPEDSource src : ((IPEDMultiSource) app.appCase).getAtomicSources()) { + if (id <= src.getLastId()) { + IItemId iid = new ItemId(src.getSourceId(), id); + iped.data.IItem item = app.appCase.getItemByItemId(iid); + if (item != null) + return item.getName(); + } + } + } + } catch (Exception ignore) { + } + return doc.itemId; + } + case 2: + return String.format("%.3f", doc.score); + case 3: { + try { + int id = Integer.parseInt(doc.itemId); + App app = App.get(); + if (app != null && app.appCase instanceof IPEDMultiSource) { + for (IPEDSource src : ((IPEDMultiSource) app.appCase).getAtomicSources()) { + if (id <= src.getLastId()) { + IItemId iid = new ItemId(src.getSourceId(), id); + iped.data.IItem item = app.appCase.getItemByItemId(iid); + if (item != null) + return item.getPath(); + } + } + } + } catch (Exception ignore) { + } + return ""; + } + default: + return ""; + } + } + + @Override + public Class getColumnClass(int columnIndex) { + return String.class; + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/config/AbstractTaskConfig.java b/iped-engine/src/main/java/iped/engine/config/AbstractTaskConfig.java index 79fe1f348a..05a4d13e7a 100644 --- a/iped-engine/src/main/java/iped/engine/config/AbstractTaskConfig.java +++ b/iped-engine/src/main/java/iped/engine/config/AbstractTaskConfig.java @@ -45,11 +45,14 @@ public void processConfig(Path resource) throws IOException { @Override public boolean isEnabled() { - return enabledProp.isEnabled(); + return enabledProp != null ? enabledProp.isEnabled() : true; } @Override public void setEnabled(boolean enabled) { + if (enabledProp == null) { + enabledProp = new EnableTaskProperty(getTaskEnableProperty()); + } enabledProp.setEnabled(enabled); } diff --git a/iped-engine/src/main/java/iped/engine/config/Configuration.java b/iped-engine/src/main/java/iped/engine/config/Configuration.java index e217b76d37..cc0dbc5b08 100644 --- a/iped-engine/src/main/java/iped/engine/config/Configuration.java +++ b/iped-engine/src/main/java/iped/engine/config/Configuration.java @@ -206,6 +206,7 @@ public void loadConfigurables(String configPathStr, boolean loadAll) throws IOEx configManager.addObject(new AnalysisConfig()); configManager.addObject(new AIFiltersConfig()); configManager.addObject(new ProcessingPriorityConfig()); + configManager.addObject(new RAGConfig()); configManager.addObject(new EnableTaskProperty(FaceRecognitionConfig.enableParam)); configManager.addObject(new EnableTaskProperty(AgeEstimationConfig.enableParam)); diff --git a/iped-engine/src/main/java/iped/engine/config/ElasticSearchTaskConfig.java b/iped-engine/src/main/java/iped/engine/config/ElasticSearchTaskConfig.java index 4dc52d9fc2..9c9b6792aa 100644 --- a/iped-engine/src/main/java/iped/engine/config/ElasticSearchTaskConfig.java +++ b/iped-engine/src/main/java/iped/engine/config/ElasticSearchTaskConfig.java @@ -28,6 +28,8 @@ public class ElasticSearchTaskConfig extends AbstractTaskPropertiesConfig { private static final String VALIDATE_SSL = "validateSSL"; private static final String TERM_VECTOR = "useTermVector"; private static final String RETRIES = "retries"; + private static final String USERNAME_KEY = "username"; + private static final String PASSWORD_KEY = "password"; private static final int DEFAULT_RETRIES = 1; @@ -47,6 +49,8 @@ public class ElasticSearchTaskConfig extends AbstractTaskPropertiesConfig { private boolean validateSSL = true; private boolean termVector = false; private int retries; + private String username; + private String password; public String getHost() { return host; @@ -134,7 +138,10 @@ public void processProperties(UTF8Properties props) { } retries = Integer.parseInt(properties.getProperty(RETRIES, Integer.toString(DEFAULT_RETRIES)).trim()); - + username = props.getProperty(USERNAME_KEY); + username = username == null ? null : username.trim(); + password = props.getProperty(PASSWORD_KEY); + password = password == null ? null : password.trim(); } public boolean getValidateSSL() { @@ -149,4 +156,12 @@ public int getRetries() { return retries; } + public String getUsername() { + return username; + } + + public String getPassword() { + return password; + } + } diff --git a/iped-engine/src/main/java/iped/engine/config/IndexTaskConfig.java b/iped-engine/src/main/java/iped/engine/config/IndexTaskConfig.java index e2da1e750f..c848256b6c 100644 --- a/iped-engine/src/main/java/iped/engine/config/IndexTaskConfig.java +++ b/iped-engine/src/main/java/iped/engine/config/IndexTaskConfig.java @@ -95,6 +95,11 @@ public void processProperties(UTF8Properties properties) { commitIntervalSeconds = Integer.parseInt(value.trim()); } + value = properties.getProperty("textOverlapSize"); //$NON-NLS-1$ + if (value != null && !value.trim().isEmpty()) { + textOverlapSize = Integer.valueOf(value.trim()); + } + } private int[] convertExtraCharsToIndex(String chars) { diff --git a/iped-engine/src/main/java/iped/engine/config/RAGConfig.java b/iped-engine/src/main/java/iped/engine/config/RAGConfig.java new file mode 100644 index 0000000000..b75874c376 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/config/RAGConfig.java @@ -0,0 +1,384 @@ +package iped.engine.config; + +import iped.utils.UTF8Properties; + +/** + * Configuration class for the RAG (Retrieval-Augmented Generation) engine. + * Reads and exposes all parameters defined in RAGConfig.txt. + * + * @author Rui Sant'Ana Junior + */ +public class RAGConfig extends AbstractTaskPropertiesConfig { + + private static final long serialVersionUID = 1L; + + private static final String CONF_FILE_NAME = "RAGConfig.txt"; + private static final String ENABLED_KEY = "enableRAG"; + + private String embeddingProvider = "local"; + private String embeddingModel = "bge-m3"; + private String embeddingEndpoint = "http://localhost:11434/api/embeddings"; + private int embeddingDimensions = 1024; + private int chunkSize = 3500; + private int chunkOverlap = 700; + private boolean enableLLMQueries = true; + private float chunkSimilarityThreshold = 0.62f; + private int maxRetrievedChunks = 5; + private float vectorSearchWeight = 0.4f; + private float lexicalSearchWeight = 0.6f; + private String llmProvider = "local"; + private String llmModel = "qwen3:4b"; + private String llmEndpoint = "http://localhost:11434/v1"; + private String llmApiKey = ""; + private int llmContextWindow = 9216; + private String vectorStoreMode = "lucene"; // "lucene" | "opensearch" + private String embeddingFilterMode = "blacklist"; + private String embeddingCategoryBlacklist = "^(Programs and Libraries|Compressed Archives|Image Disks|Unallocated|File Slacks|Main Registry Files|Other Registry Files|Fonts|Geographic Data|Databases|XML Files|Other)$"; + private String embeddingMimeTypeBlacklist = "^(application/octet-stream" + + "|application/x-dosexec|application/x-msdownload|application/x-ms-installer|application/x-msi" + + "|application/x-sharedlib" + // .so, .dex, .odex (bin??rios Android/Linux) + "|application/x-dex" + // Dalvik Executable (.dex) + "|application/vnd\\.android\\.dex" + // variante de MIME para .dex + "|application/vnd\\.android\\.package-archive" + // pacotes APK + "|application/java-archive|application/x-java-archive" + // arquivos JAR/WAR/EAR + "|application/x-object" + // arquivos objeto compilados (.o) + "|application/x-archive" + // bibliotecas est??ticas (.a) + "|application/x-executable" + // execut??veis ELF gen??ricos + "|application/x-windows-registry-main|application/x-windows-registry|application/x-sqlite3|application/x-edb|application/x-msaccess" + + "|text/xml|application/xml" + // Dumps XML / minireport.xml + "|font/.*" + + "|application/zip|application/x-7z-compressed|application/x-rar-compressed|application/x-tar|application/x-gzip|application/x-bzip2" + + "|application/x-ms-pdb|application/x-win-lnk" + + ")$"; + private String embeddingCategoryWhitelist = "^(Documents|Emails and Mailboxes|Chats|SMS Messages|Instant Messages Artifacts|Notes|Plain Texts|HTML and Web pages|PDF Documents|Office Documents|Audio)$"; + private String embeddingMimeTypeWhitelist = "^(text/plain|text/html|application/xhtml\\+xml|text/csv|text/tab-separated-values|message/rfc822|application/pdf|application/msword|application/vnd\\.openxmlformats-officedocument\\..*|application/vnd\\.ms-excel|application/vnd\\.ms-powerpoint|application/vnd\\.oasis\\.opendocument\\..*|application/rtf|application/json|audio/.*)$"; + private transient java.util.regex.Pattern embeddingCategoryBlacklistPattern; + private transient java.util.regex.Pattern embeddingMimeTypeBlacklistPattern; + private transient java.util.regex.Pattern embeddingCategoryWhitelistPattern; + private transient java.util.regex.Pattern embeddingMimeTypeWhitelistPattern; + + public String getEmbeddingProvider() { + return embeddingProvider; + } + + public String getEmbeddingModel() { + return embeddingModel; + } + + public String getEmbeddingEndpoint() { + return embeddingEndpoint; + } + + public int getEmbeddingDimensions() { + return embeddingDimensions; + } + + public int getChunkSize() { + return chunkSize; + } + + public int getChunkOverlap() { + return chunkOverlap; + } + + public boolean isEnableLLMQueries() { + return enableLLMQueries; + } + + public float getChunkSimilarityThreshold() { + return chunkSimilarityThreshold; + } + + public int getMaxRetrievedChunks() { + return maxRetrievedChunks; + } + + /** + * Weight applied to the vector (semantic) side of the RRF fusion score. + * Must be in [0.0, 1.0]. The lexical weight is automatically 1 - vectorWeight. + * Default: 0.4 (40/60 split between vector and lexical). + * For forensic use with many entity queries (names, CPFs, codes), consider + * shifting towards lexical, e.g. vectorSearchWeight = 0.4. + */ + public float getVectorSearchWeight() { + return vectorSearchWeight; + } + + /** + * Weight applied to the lexical (BM25) side of the RRF fusion score. + * Automatically derived as 1.0 - vectorSearchWeight. + */ + public float getLexicalSearchWeight() { + return lexicalSearchWeight; + } + + public String getLlmProvider() { + return llmProvider; + } + + public String getLlmModel() { + return llmModel; + } + + public String getLlmEndpoint() { + return llmEndpoint; + } + + public String getLlmApiKey() { + return llmApiKey; + } + + public int getLlmContextWindow() { + return llmContextWindow; + } + + /** + * Returns the vector store backend mode. + * "lucene" : vectors are stored in the case's own Lucene index (portable, default). + * "opensearch" : vectors are stored in an external OpenSearch cluster (corporate/legacy). + */ + public String getVectorStoreMode() { + return vectorStoreMode; + } + + public void setVectorStoreMode(String vectorStoreMode) { + this.vectorStoreMode = vectorStoreMode; + } + + public String getEmbeddingFilterMode() { + return embeddingFilterMode; + } + + public String getEmbeddingCategoryBlacklist() { + return embeddingCategoryBlacklist; + } + + public String getEmbeddingMimeTypeBlacklist() { + return embeddingMimeTypeBlacklist; + } + + public String getEmbeddingCategoryWhitelist() { + return embeddingCategoryWhitelist; + } + + public String getEmbeddingMimeTypeWhitelist() { + return embeddingMimeTypeWhitelist; + } + + public java.util.regex.Pattern getEmbeddingCategoryBlacklistPattern() { + if (embeddingCategoryBlacklistPattern == null) { + synchronized (this) { + if (embeddingCategoryBlacklistPattern == null) { + String patternStr = embeddingCategoryBlacklist != null ? embeddingCategoryBlacklist.trim() : ""; + if (patternStr.isEmpty()) { + patternStr = "$^"; + } + embeddingCategoryBlacklistPattern = java.util.regex.Pattern.compile(patternStr, java.util.regex.Pattern.CASE_INSENSITIVE); + } + } + } + return embeddingCategoryBlacklistPattern; + } + + public java.util.regex.Pattern getEmbeddingMimeTypeBlacklistPattern() { + if (embeddingMimeTypeBlacklistPattern == null) { + synchronized (this) { + if (embeddingMimeTypeBlacklistPattern == null) { + String patternStr = embeddingMimeTypeBlacklist != null ? embeddingMimeTypeBlacklist.trim() : ""; + if (patternStr.isEmpty()) { + patternStr = "$^"; + } + embeddingMimeTypeBlacklistPattern = java.util.regex.Pattern.compile(patternStr, java.util.regex.Pattern.CASE_INSENSITIVE); + } + } + } + return embeddingMimeTypeBlacklistPattern; + } + + public java.util.regex.Pattern getEmbeddingCategoryWhitelistPattern() { + if (embeddingCategoryWhitelistPattern == null) { + synchronized (this) { + if (embeddingCategoryWhitelistPattern == null) { + String patternStr = embeddingCategoryWhitelist != null ? embeddingCategoryWhitelist.trim() : ""; + if (patternStr.isEmpty()) { + patternStr = "$^"; + } + embeddingCategoryWhitelistPattern = java.util.regex.Pattern.compile(patternStr, java.util.regex.Pattern.CASE_INSENSITIVE); + } + } + } + return embeddingCategoryWhitelistPattern; + } + + public java.util.regex.Pattern getEmbeddingMimeTypeWhitelistPattern() { + if (embeddingMimeTypeWhitelistPattern == null) { + synchronized (this) { + if (embeddingMimeTypeWhitelistPattern == null) { + String patternStr = embeddingMimeTypeWhitelist != null ? embeddingMimeTypeWhitelist.trim() : ""; + if (patternStr.isEmpty()) { + patternStr = "$^"; + } + embeddingMimeTypeWhitelistPattern = java.util.regex.Pattern.compile(patternStr, java.util.regex.Pattern.CASE_INSENSITIVE); + } + } + } + return embeddingMimeTypeWhitelistPattern; + } + + /** + * Decides whether an item with the given categories and MIME type should have an + * embedding generated. Returns {@code true} if ANY category/MIME combination + * passes the configured blacklist/whitelist filter. + */ + public boolean shouldEmbed(java.util.Collection categories, String mimeType) { + // In whitelist mode, at least one category OR the MIME type must match. + // In blacklist mode, no category AND no MIME type must be blocked. + if (categories == null || categories.isEmpty()) { + // No categories available ??? decide purely by MIME type using null category + return shouldEmbed((String) null, mimeType); + } + if ("whitelist".equalsIgnoreCase(embeddingFilterMode)) { + for (String category : categories) { + if (shouldEmbed(category, mimeType)) { + return true; + } + } + return false; + } else { + for (String category : categories) { + if (!shouldEmbed(category, mimeType)) { + return false; + } + } + return true; + } + } + + /** + * Decides whether an item with a single category and MIME type should have an + * embedding generated. This is the canonical implementation used by both overloads. + */ + public boolean shouldEmbed(String category, String mimeType) { + if ("whitelist".equalsIgnoreCase(embeddingFilterMode)) { + boolean categoryMatched = embeddingCategoryWhitelist != null + && !embeddingCategoryWhitelist.trim().isEmpty() + && category != null + && getEmbeddingCategoryWhitelistPattern().matcher(category).matches(); + boolean mimeTypeMatched = embeddingMimeTypeWhitelist != null + && !embeddingMimeTypeWhitelist.trim().isEmpty() + && mimeType != null + && getEmbeddingMimeTypeWhitelistPattern().matcher(mimeType).matches(); + return categoryMatched || mimeTypeMatched; + } else { + if (embeddingCategoryBlacklist != null && !embeddingCategoryBlacklist.trim().isEmpty() + && category != null + && getEmbeddingCategoryBlacklistPattern().matcher(category).matches()) { + return false; + } + if (embeddingMimeTypeBlacklist != null && !embeddingMimeTypeBlacklist.trim().isEmpty() + && mimeType != null + && getEmbeddingMimeTypeBlacklistPattern().matcher(mimeType).matches()) { + return false; + } + return true; + } + } + + + + @Override + public String getTaskEnableProperty() { + return ENABLED_KEY; + } + + @Override + public String getTaskConfigFileName() { + return CONF_FILE_NAME; + } + + @Override + public void processProperties(UTF8Properties props) { + String val; + + val = props.getProperty("embeddingProvider"); + if (val != null) embeddingProvider = val.trim(); + + val = props.getProperty("embeddingModel"); + if (val != null) embeddingModel = val.trim(); + + val = props.getProperty("embeddingEndpoint"); + if (val != null) embeddingEndpoint = val.trim(); + + val = props.getProperty("embeddingDimensions"); + if (val != null) embeddingDimensions = Integer.parseInt(val.trim()); + + val = props.getProperty("chunkSize"); + if (val != null) chunkSize = Integer.parseInt(val.trim()); + + val = props.getProperty("chunkOverlap"); + if (val != null) chunkOverlap = Integer.parseInt(val.trim()); + + val = props.getProperty("enableLLMQueries"); + if (val != null) enableLLMQueries = Boolean.parseBoolean(val.trim()); + + val = props.getProperty("chunkSimilarityThreshold"); + if (val != null) chunkSimilarityThreshold = Float.parseFloat(val.trim()); + + val = props.getProperty("maxRetrievedChunks"); + if (val != null) maxRetrievedChunks = Integer.parseInt(val.trim()); + + val = props.getProperty("vectorSearchWeight"); + if (val != null) { + float v = Float.parseFloat(val.trim()); + if (v < 0.0f || v > 1.0f) + throw new IllegalArgumentException("vectorSearchWeight must be between 0.0 and 1.0, got: " + v); + vectorSearchWeight = v; + lexicalSearchWeight = 1.0f - v; + } + + val = props.getProperty("llmProvider"); + if (val != null) llmProvider = val.trim(); + + val = props.getProperty("llmModel"); + if (val != null) llmModel = val.trim(); + + val = props.getProperty("llmEndpoint"); + if (val != null) llmEndpoint = val.trim(); + + val = props.getProperty("llmApiKey"); + if (val != null) llmApiKey = val.trim(); + + val = props.getProperty("llmContextWindow"); + if (val != null) llmContextWindow = Integer.parseInt(val.trim()); + + val = props.getProperty("vectorStoreMode"); + if (val != null) vectorStoreMode = val.trim().toLowerCase(); + + val = props.getProperty("embeddingFilterMode"); + if (val != null) embeddingFilterMode = val.trim(); + + val = props.getProperty("embeddingCategoryBlacklist"); + if (val != null) { + embeddingCategoryBlacklist = val.trim(); + embeddingCategoryBlacklistPattern = null; + } + + val = props.getProperty("embeddingMimeTypeBlacklist"); + if (val != null) { + embeddingMimeTypeBlacklist = val.trim(); + embeddingMimeTypeBlacklistPattern = null; + } + + val = props.getProperty("embeddingCategoryWhitelist"); + if (val != null) { + embeddingCategoryWhitelist = val.trim(); + embeddingCategoryWhitelistPattern = null; + } + + val = props.getProperty("embeddingMimeTypeWhitelist"); + if (val != null) { + embeddingMimeTypeWhitelist = val.trim(); + embeddingMimeTypeWhitelistPattern = null; + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/core/Statistics.java b/iped-engine/src/main/java/iped/engine/core/Statistics.java index 33cea61f65..02aaed0173 100644 --- a/iped-engine/src/main/java/iped/engine/core/Statistics.java +++ b/iped-engine/src/main/java/iped/engine/core/Statistics.java @@ -321,6 +321,14 @@ public void logStatistics(Manager manager) throws Exception { int indexed = (numDocs - getSplits() - previousIndexedFiles) / 2; LOGGER.info("Total Indexed: {}", indexed); //$NON-NLS-1$ + if (iped.engine.task.index.IndexTask.totalEmbeddedDocs.get() > 0) { + iped.engine.rag.RAGService ragService = iped.engine.rag.RAGService.getInstance(); + String ragModel = (ragService != null && ragService.getConfig() != null) ? ragService.getConfig().getEmbeddingModel() : "unknown"; + LOGGER.info("RAG Embeddings: {} vectors generated in {} documents (Model: {}).", + iped.engine.task.index.IndexTask.totalEmbeddedVectors.get(), + iped.engine.task.index.IndexTask.totalEmbeddedDocs.get(), + ragModel); + } LOGGER.info("Discovered volume: {} bytes", caseData.getDiscoveredVolume()); LOGGER.info("Processed volume: {} bytes", getVolume()); diff --git a/iped-engine/src/main/java/iped/engine/rag/ClaudeLLMProvider.java b/iped-engine/src/main/java/iped/engine/rag/ClaudeLLMProvider.java new file mode 100644 index 0000000000..030aae2126 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/ClaudeLLMProvider.java @@ -0,0 +1,169 @@ +package iped.engine.rag; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +/** + * LLM provider implementation that interacts with Anthropic's Claude API + * (e.g., claude-3-5-sonnet-20240620, etc.). + * + * @author Rui Sant'Ana Junior + */ +public class ClaudeLLMProvider implements LLMProvider { + + private final String model; + private final String apiKey; + private final String customEndpoint; + private final HttpClient client; + + public ClaudeLLMProvider(String customEndpoint, String model, String apiKey) { + this.customEndpoint = customEndpoint; + this.model = (model == null || model.isEmpty()) ? "claude-sonnet-4-6" : model; + this.apiKey = apiKey; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + @Override + public String generateAnswer(String question, String context) throws IOException, InterruptedException { + return generateAnswer(question, context, java.util.Collections.emptyList()); + } + + @Override + public String generateAnswer(String question, String context, java.util.List history) throws IOException, InterruptedException { + String systemPrompt = getSystemPrompt(question); + String userPrompt = "Context:\n" + context + "\n\nQuestion: " + question + "\n\nAnswer:"; + + StringBuilder messagesJson = new StringBuilder(); + messagesJson.append("["); + + boolean first = true; + if (history != null) { + for (RAGService.HistoryTurn turn : history) { + if (!first) messagesJson.append(","); + messagesJson.append("{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString("Question: " + turn.question)).append("}"); + messagesJson.append(",{\"role\":\"assistant\",\"content\":").append(LocalEmbeddingProvider.toJsonString(turn.answer)).append("}"); + first = false; + } + } + + if (!first) messagesJson.append(","); + messagesJson.append("{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString(userPrompt)).append("}"); + messagesJson.append("]"); + + String requestBody = "{" + + "\"model\":\"" + model + "\"," + + "\"max_tokens\":1024," + + "\"system\":" + LocalEmbeddingProvider.toJsonString(systemPrompt) + "," + + "\"messages\":" + messagesJson.toString() + "," + + "\"temperature\":0.0" + + "}"; + + + String targetUrl; + if (customEndpoint != null && !customEndpoint.trim().isEmpty()) { + targetUrl = customEndpoint; + } else { + targetUrl = "https://api.anthropic.com/v1/messages"; + } + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(targetUrl)) + .header("Content-Type", "application/json") + .header("x-api-key", apiKey) + .header("anthropic-version", "2023-06-01") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(60)) + .build(); + + int maxRetries = 3; + long backoffMs = 2000; + int lastStatusCode = -1; + String lastResponseBody = ""; + + for (int attempt = 0; attempt < maxRetries; attempt++) { + try { + HttpResponse response = sendPrivileged(client, request); + lastStatusCode = response.statusCode(); + lastResponseBody = response.body(); + + if (response.statusCode() == 200) { + return parseClaudeResponse(response.body()); + } else if (response.statusCode() == 429) { + Thread.sleep(backoffMs); + } else if (response.statusCode() >= 500) { + Thread.sleep(backoffMs); + } else { + throw new IOException("Claude LLM service returned status code " + response.statusCode() + ": " + response.body()); + } + } catch (IOException e) { + if (attempt == maxRetries - 1) { + throw e; + } + Thread.sleep(backoffMs); + } + backoffMs *= 2; + } + throw new IOException("Failed to generate answer from Claude LLM after " + maxRetries + " attempts. Last status: " + lastStatusCode + ", Response: " + lastResponseBody); + } + + public static String parseClaudeResponse(String json) throws IOException { + int index = json.indexOf("\"text\""); + if (index == -1) { + throw new IOException("Failed to parse Claude response. Could not find 'text' key in response JSON: " + json); + } + index = json.indexOf(":", index); + if (index == -1) { + throw new IOException("Failed to parse Claude response. Invalid JSON format near 'text' in: " + json); + } + index = json.indexOf("\"", index); + if (index == -1) { + throw new IOException("Failed to parse Claude response. Could not find start of text string in: " + json); + } + + StringBuilder sb = new StringBuilder(); + boolean escaped = false; + for (int i = index + 1; i < json.length(); i++) { + char c = json.charAt(i); + if (escaped) { + switch (c) { + case 'n': sb.append('\n'); break; + case 'r': sb.append('\r'); break; + case 't': sb.append('\t'); break; + case 'b': sb.append('\b'); break; + case 'f': sb.append('\f'); break; + case '"': sb.append('"'); break; + case '\\': sb.append('\\'); break; + default: sb.append(c); + } + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + return sb.toString(); + } else { + sb.append(c); + } + } + throw new IOException("Failed to parse Claude response. Unclosed text string in response JSON: " + json); + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/EmbeddingProvider.java b/iped-engine/src/main/java/iped/engine/rag/EmbeddingProvider.java new file mode 100644 index 0000000000..00bbe8c604 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/EmbeddingProvider.java @@ -0,0 +1,38 @@ +package iped.engine.rag; + +import java.io.IOException; + +/** + * Interface for vector embedding generation providers. + * + * @author Rui Sant'Ana Junior + */ +public interface EmbeddingProvider { + /** + * Generates a numeric vector representation (embedding) for the given text. + * + * @param text input text + * @return float array representation of the embedding + * @throws IOException if network or IO errors occur + * @throws InterruptedException if thread is interrupted during call + */ + float[] generateEmbedding(String text) throws IOException, InterruptedException; + + /** + * Generates numeric vector representations (embeddings) for a list of texts in batch. + * + * @param texts list of input texts + * @return list of float array representations matching input texts order + * @throws IOException if network or IO errors occur + * @throws InterruptedException if thread is interrupted during call + */ + default java.util.List generateEmbeddings(java.util.List texts) throws IOException, InterruptedException { + java.util.List list = new java.util.ArrayList<>(); + if (texts != null) { + for (String t : texts) { + list.add(generateEmbedding(t)); + } + } + return list; + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/GeminiEmbeddingProvider.java b/iped-engine/src/main/java/iped/engine/rag/GeminiEmbeddingProvider.java new file mode 100644 index 0000000000..25c166f62b --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/GeminiEmbeddingProvider.java @@ -0,0 +1,174 @@ +package iped.engine.rag; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; + +/** + * Embedding provider that calls the Google Generative Language API + * (Gemini Embeddings). + * + * @author Rui Sant'Ana Junior + * + *

Endpoint used: + * {@code POST https://generativelanguage.googleapis.com/v1beta/models/{model}:embedContent?key={apiKey}} + * + *

Request body: + * 

+ * {
+ *   "model": "models/text-embedding-004",
+ *   "content": { "parts": [{ "text": "..." }] }
+ * }
+ *
+ * + *

Response: + * 

+ * {
+ *   "embedding": { "values": [0.123, ...] }
+ * }
+ *
+ * + * Recommended model: {@code text-embedding-004} (768 dimensions). + * Set {@code embeddingDimensions = 768} in RAGConfig.txt when using this model. + */ +public class GeminiEmbeddingProvider implements EmbeddingProvider { + + private static final String BASE_URL = + "https://generativelanguage.googleapis.com/v1beta/models/"; + + private final String model; + private final String apiKey; + private final String customEndpoint; + private final int outputDimensionality; + private final HttpClient client; + + public GeminiEmbeddingProvider(String customEndpoint, String model, String apiKey, int outputDimensionality) { + this.customEndpoint = customEndpoint; + this.model = (model == null || model.isBlank()) ? "gemini-embedding-001" : model; + this.apiKey = apiKey; + this.outputDimensionality = outputDimensionality; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + public GeminiEmbeddingProvider(String customEndpoint, String model, String apiKey) { + this(customEndpoint, model, apiKey, 768); + } + + @Override + public float[] generateEmbedding(String text) throws IOException, InterruptedException { + String targetUrl; + if (customEndpoint != null && !customEndpoint.isBlank()) { + targetUrl = customEndpoint; + } else { + // model may be given as "text-embedding-004" or "models/text-embedding-004" + String modelPath = model.startsWith("models/") ? model : "models/" + model; + targetUrl = BASE_URL + modelPath.replaceFirst("models/", "") + + ":embedContent?key=" + apiKey; + } + + String modelName = model.startsWith("models/") ? model.substring(7) : model; + String requestBody = "{" + + "\"model\":\"models/" + modelName + "\"," + + "\"content\":{\"parts\":[{\"text\":" + LocalEmbeddingProvider.toJsonString(text) + "}]}"; + if (outputDimensionality > 0 && !"embedding-001".equalsIgnoreCase(modelName)) { + requestBody += ",\"outputDimensionality\":" + outputDimensionality; + } + requestBody += "}"; + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(targetUrl)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(30)) + .build(); + + int maxRetries = 5; + long backoffMs = 2000; + + for (int attempt = 0; attempt < maxRetries; attempt++) { + try { + HttpResponse response = sendPrivileged(client, request); + + if (response.statusCode() == 200) { + return parseGeminiEmbeddingResponse(response.body()); + } else if (response.statusCode() == 429) { + String retryAfter = response.headers().firstValue("Retry-After").orElse(""); + long sleepMs = backoffMs; + if (!retryAfter.isBlank()) { + try { sleepMs = Long.parseLong(retryAfter) * 1000; } + catch (NumberFormatException ignore) {} + } + Thread.sleep(sleepMs); + } else if (response.statusCode() >= 500) { + Thread.sleep(backoffMs); + } else { + throw new IOException("Gemini Embedding API returned status " + + response.statusCode() + ": " + response.body()); + } + } catch (IOException e) { + if (attempt == maxRetries - 1) throw e; + Thread.sleep(backoffMs); + } + backoffMs *= 2; + } + throw new IOException("Failed to generate Gemini embedding after " + maxRetries + " attempts."); + } + + /** + * Parses the {@code embedding.values} array from the Gemini API JSON response. + * + * Expected structure: + * 
{"embedding":{"values":[0.123,-0.456,...]}}
+ */ + static float[] parseGeminiEmbeddingResponse(String json) throws IOException { + // Find "values":[ + int valuesIdx = json.indexOf("\"values\""); + if (valuesIdx == -1) { + throw new IOException("Gemini embedding response missing 'values' field: " + json); + } + int arrStart = json.indexOf('[', valuesIdx); + int arrEnd = json.indexOf(']', arrStart); + if (arrStart == -1 || arrEnd == -1) { + throw new IOException("Gemini embedding response has malformed 'values' array: " + json); + } + + String arrContent = json.substring(arrStart + 1, arrEnd).trim(); + if (arrContent.isEmpty()) { + throw new IOException("Gemini embedding response returned empty values array."); + } + + String[] parts = arrContent.split(","); + List floats = new ArrayList<>(parts.length); + for (String part : parts) { + String t = part.trim(); + if (!t.isEmpty()) { + floats.add(Float.parseFloat(t)); + } + } + + float[] result = new float[floats.size()]; + for (int i = 0; i < floats.size(); i++) result[i] = floats.get(i); + return result; + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/GeminiLLMProvider.java b/iped-engine/src/main/java/iped/engine/rag/GeminiLLMProvider.java new file mode 100644 index 0000000000..49a85cd5a3 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/GeminiLLMProvider.java @@ -0,0 +1,163 @@ +package iped.engine.rag; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +public class GeminiLLMProvider implements LLMProvider { + + private final String model; + private final String apiKey; + private final String customEndpoint; + private final HttpClient client; + + public GeminiLLMProvider(String customEndpoint, String model, String apiKey) { + this.customEndpoint = customEndpoint; + this.model = (model == null || model.isEmpty()) ? "gemini-1.5-flash" : model; + this.apiKey = apiKey; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + @Override + public String generateAnswer(String question, String context) throws IOException, InterruptedException { + return generateAnswer(question, context, java.util.Collections.emptyList()); + } + + @Override + public String generateAnswer(String question, String context, java.util.List history) throws IOException, InterruptedException { + String systemPrompt = getSystemPrompt(question); + String userPrompt = "Context:\n" + context + "\n\nQuestion: " + question + "\n\nAnswer:"; + + StringBuilder contentsJson = new StringBuilder(); + contentsJson.append("["); + + boolean first = true; + if (history != null) { + for (RAGService.HistoryTurn turn : history) { + if (!first) contentsJson.append(","); + contentsJson.append("{\"role\":\"user\",\"parts\":[{\"text\":").append(LocalEmbeddingProvider.toJsonString("Question: " + turn.question)).append("}]}"); + contentsJson.append(",{\"role\":\"model\",\"parts\":[{\"text\":").append(LocalEmbeddingProvider.toJsonString(turn.answer)).append("}]}"); + first = false; + } + } + + if (!first) contentsJson.append(","); + contentsJson.append("{\"role\":\"user\",\"parts\":[{\"text\":").append(LocalEmbeddingProvider.toJsonString(userPrompt)).append("}]}"); + contentsJson.append("]"); + + String requestBody = "{" + + "\"systemInstruction\":{" + + " \"parts\":[{\"text\":" + LocalEmbeddingProvider.toJsonString(systemPrompt) + "}]" + + "}," + + "\"contents\":" + contentsJson.toString() + "," + + "\"generationConfig\":{" + + " \"temperature\":0.0" + + "}" + + "}"; + + + String targetUrl; + if (customEndpoint != null && !customEndpoint.trim().isEmpty()) { + targetUrl = customEndpoint; + } else { + targetUrl = "https://generativelanguage.googleapis.com/v1beta/models/" + model + ":generateContent?key=" + apiKey; + } + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(targetUrl)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(60)) + .build(); + + int maxRetries = 3; + long backoffMs = 2000; + int lastStatusCode = -1; + String lastResponseBody = ""; + + for (int attempt = 0; attempt < maxRetries; attempt++) { + try { + HttpResponse response = sendPrivileged(client, request); + lastStatusCode = response.statusCode(); + lastResponseBody = response.body(); + + if (response.statusCode() == 200) { + return parseGeminiResponse(response.body()); + } else if (response.statusCode() == 429) { + Thread.sleep(backoffMs); + } else if (response.statusCode() >= 500) { + Thread.sleep(backoffMs); + } else { + throw new IOException("Gemini LLM service returned status code " + response.statusCode() + ": " + response.body()); + } + } catch (IOException e) { + if (attempt == maxRetries - 1) { + throw e; + } + Thread.sleep(backoffMs); + } + backoffMs *= 2; + } + throw new IOException("Failed to generate answer from Gemini LLM after " + maxRetries + " attempts. Last status: " + lastStatusCode + ", Response: " + lastResponseBody); + } + + public static String parseGeminiResponse(String json) throws IOException { + int index = json.indexOf("\"text\""); + if (index == -1) { + throw new IOException("Failed to parse Gemini response. Could not find 'text' key in response JSON: " + json); + } + index = json.indexOf(":", index); + if (index == -1) { + throw new IOException("Failed to parse Gemini response. Invalid JSON format near 'text' in: " + json); + } + index = json.indexOf("\"", index); + if (index == -1) { + throw new IOException("Failed to parse Gemini response. Could not find start of text string in: " + json); + } + + StringBuilder sb = new StringBuilder(); + boolean escaped = false; + for (int i = index + 1; i < json.length(); i++) { + char c = json.charAt(i); + if (escaped) { + switch (c) { + case 'n': sb.append('\n'); break; + case 'r': sb.append('\r'); break; + case 't': sb.append('\t'); break; + case 'b': sb.append('\b'); break; + case 'f': sb.append('\f'); break; + case '"': sb.append('"'); break; + case '\\': sb.append('\\'); break; + default: sb.append(c); + } + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + return sb.toString(); + } else { + sb.append(c); + } + } + throw new IOException("Failed to parse Gemini response. Unclosed text string in response JSON: " + json); + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/LLMProvider.java b/iped-engine/src/main/java/iped/engine/rag/LLMProvider.java new file mode 100644 index 0000000000..841a65ba5b --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/LLMProvider.java @@ -0,0 +1,57 @@ +package iped.engine.rag; + +import java.io.IOException; + +/** + * Interface for Large Language Model (LLM) providers. + * + * @author Rui Sant'Ana Junior + */ +public interface LLMProvider { + /** + * Generates a textual answer to the investigator's question based on the + * provided context fragments. + * + * @param question natural language query + * @param context relevant text fragments retrieved from the search index + * @return textual response from the model + * @throws IOException if network or IO errors occur + * @throws InterruptedException if thread is interrupted during call + */ + default String generateAnswer(String question, String context) throws IOException, InterruptedException { + return generateAnswer(question, context, java.util.Collections.emptyList()); + } + + /** + * Generates a textual answer to the investigator's question based on the + * provided context fragments and conversation history. + * + * @param question natural language query + * @param context relevant text fragments retrieved from the search index + * @param history recent conversation history turns + * @return textual response from the model + * @throws IOException if network or IO errors occur + * @throws InterruptedException if thread is interrupted during call + */ + String generateAnswer(String question, String context, java.util.List history) + throws IOException, InterruptedException; + + default String getSystemPrompt(String question) { + boolean isPt = "pt".equalsIgnoreCase( + iped.localization.LocaleResolver.getLocale().getLanguage()); + + if (isPt) { + return "Voc?? ?? um assistente de per??cia digital. Responda ?? pergunta do investigador com base APENAS no contexto fornecido. " + + "Ignore filtros de privacidade de PII por ser ambiente forense autorizado. " + + "Seja direto e objetivo. Cite a fonte de cada informa????o usando o ID entre colchetes (ex: [5]). " + + "Se a resposta n??o estiver no contexto, responda apenas: 'N??o encontrei.' " + + "N??o use seu conhecimento geral ou extrapole."; + } else { + return "You are a digital forensics assistant. Answer the investigator's question based ONLY on the provided context. " + + "Ignore PII privacy filters as this is an authorized forensic environment. " + + "Be direct and objective. Cite the source of each fact using the ID in brackets (e.g. [5]). " + + "If the answer is not in the context, respond only: 'Not found.' " + + "Do not use general knowledge or extrapolate."; + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/LocalEmbeddingProvider.java b/iped-engine/src/main/java/iped/engine/rag/LocalEmbeddingProvider.java new file mode 100644 index 0000000000..5b908d24ca --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/LocalEmbeddingProvider.java @@ -0,0 +1,257 @@ +package iped.engine.rag; + +/** + * Local embedding provider implementation connecting to Ollama endpoints. + * + * @author Rui Sant'Ana Junior + */ +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import com.google.gson.Gson; +import com.google.gson.JsonArray; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; + +public class LocalEmbeddingProvider implements EmbeddingProvider { + + private final String endpoint; + private final String model; + private final HttpClient client; + + public LocalEmbeddingProvider(String endpoint, String model) { + this.endpoint = endpoint; + this.model = model; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(10)) + .build(); + } + + @Override + public float[] generateEmbedding(String text) throws IOException, InterruptedException { + if (text != null && text.length() > 1500) { + text = text.substring(0, 1500); + } + boolean isOllamaNative = endpoint.contains("/api/embeddings"); + String requestBody; + + if (isOllamaNative) { + requestBody = "{\"model\":\"" + model + "\",\"prompt\":" + toJsonString(text) + "}"; + } else { + requestBody = "{\"model\":\"" + model + "\",\"input\":" + toJsonString(text) + "}"; + } + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(endpoint)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(30)) + .build(); + + HttpResponse response = sendPrivileged(client, request); + + if (response.statusCode() == 200) { + return parseEmbeddingsArray(response.body()); + } + throw new IOException("Local embedding service returned status code " + response.statusCode() + ": " + response.body()); + } + + @Override + public java.util.List generateEmbeddings(java.util.List texts) throws IOException, InterruptedException { + if (texts == null || texts.isEmpty()) { + return java.util.Collections.emptyList(); + } + if (texts.size() == 1) { + return java.util.Collections.singletonList(generateEmbedding(texts.get(0))); + } + + java.util.List allResults = new java.util.ArrayList<>(texts.size()); + int subBatchSize = 8; // Process in batches of 8 items + + for (int start = 0; start < texts.size(); start += subBatchSize) { + int end = Math.min(start + subBatchSize, texts.size()); + java.util.List subBatch = texts.subList(start, end); + java.util.List batchResult = processSubBatch(subBatch, 30); + + if (batchResult != null && batchResult.size() == subBatch.size()) { + allResults.addAll(batchResult); + } else { + // Fallback for items individually without retries or sleeps + for (String txt : subBatch) { + allResults.add(generateEmbedding(txt)); + } + } + } + + return allResults; + } + + private java.util.List processSubBatch(java.util.List subBatch, int timeoutSeconds) throws IOException, InterruptedException { + String batchUrl = endpoint; + if (endpoint.contains("/api/embeddings") || endpoint.contains("/api/embed")) { + int idx = endpoint.indexOf("/api/"); + if (idx != -1) { + batchUrl = endpoint.substring(0, idx) + "/v1/embeddings"; + } + } + + StringBuilder inputsJson = new StringBuilder("["); + for (int i = 0; i < subBatch.size(); i++) { + String t = subBatch.get(i); + if (t != null && t.length() > 1500) { + t = t.substring(0, 1500); + } + inputsJson.append(toJsonString(t)); + if (i < subBatch.size() - 1) { + inputsJson.append(","); + } + } + inputsJson.append("]"); + + String requestBody = "{\"model\":\"" + model + "\",\"input\":" + inputsJson.toString() + "}"; + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(batchUrl)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(timeoutSeconds)) + .build(); + + HttpResponse response = sendPrivileged(client, request); + + if (response.statusCode() == 200) { + java.util.List res = parseBatchEmbeddings(response.body()); + if (res != null && res.size() == subBatch.size()) { + return res; + } + } + throw new IOException("Local embedding batch returned status code " + response.statusCode() + ": " + response.body()); + } + + private static final Gson GSON = new Gson(); + + public static java.util.List parseBatchEmbeddings(String json) throws IOException { + if (json == null) return null; + java.util.List list = new java.util.ArrayList<>(); + try { + JsonObject obj = GSON.fromJson(json, JsonObject.class); + + // OpenAI format: "data": [ {"embedding": [...], "index": 0}, ... ] + if (obj.has("data")) { + JsonArray data = obj.getAsJsonArray("data"); + java.util.Map indexMap = new java.util.TreeMap<>(); + for (int i = 0; i < data.size(); i++) { + JsonObject item = data.get(i).getAsJsonObject(); + int index = item.has("index") ? item.get("index").getAsInt() : i; + if (item.has("embedding")) { + JsonArray emb = item.getAsJsonArray("embedding"); + float[] vec = new float[emb.size()]; + for (int j = 0; j < emb.size(); j++) { + vec[j] = emb.get(j).getAsFloat(); + } + indexMap.put(index, vec); + } + } + for (float[] vec : indexMap.values()) { + list.add(vec); + } + return list; + } + + // Ollama native format: "embeddings": [[...], [...]] + if (obj.has("embeddings")) { + JsonArray embs = obj.getAsJsonArray("embeddings"); + for (int i = 0; i < embs.size(); i++) { + JsonArray emb = embs.get(i).getAsJsonArray(); + float[] vec = new float[emb.size()]; + for (int j = 0; j < emb.size(); j++) { + vec[j] = emb.get(j).getAsFloat(); + } + list.add(vec); + } + return list; + } + + return null; + } catch (Exception e) { + throw new IOException("Failed to parse batch embeddings from JSON: " + json, e); + } + } + + public static String toJsonString(String text) { + if (text == null) return "null"; + StringBuilder sb = new StringBuilder(); + sb.append("\""); + for (int i = 0; i < text.length(); i++) { + char c = text.charAt(i); + switch (c) { + case '"': sb.append("\\\""); break; + case '\\': sb.append("\\\\"); break; + case '\b': sb.append("\\b"); break; + case '\f': sb.append("\\f"); break; + case '\n': sb.append("\\n"); break; + case '\r': sb.append("\\r"); break; + case '\t': sb.append("\\t"); break; + default: + if (c < ' ') { + String t = "000" + Integer.toHexString(c); + sb.append("\\u" + t.substring(t.length() - 4)); + } else { + sb.append(c); + } + } + } + sb.append("\""); + return sb.toString(); + } + + public static float[] parseEmbeddingsArray(String json) throws IOException { + try { + JsonObject obj = GSON.fromJson(json, JsonObject.class); + JsonArray arr = null; + if (obj.has("embedding")) { + arr = obj.getAsJsonArray("embedding"); + } else if (obj.has("embeddings")) { + JsonArray outer = obj.getAsJsonArray("embeddings"); + if (outer.size() > 0) { + arr = outer.get(0).getAsJsonArray(); + } + } else { + for (java.util.Map.Entry entry : obj.entrySet()) { + if (entry.getValue().isJsonArray()) { + arr = entry.getValue().getAsJsonArray(); + break; + } + } + } + + if (arr != null) { + float[] result = new float[arr.size()]; + for (int i = 0; i < arr.size(); i++) { + result[i] = arr.get(i).getAsFloat(); + } + return result; + } + throw new IOException("No embedding array found in JSON: " + json); + } catch (Exception e) { + throw new IOException("Failed to parse embedding vector from JSON: " + json, e); + } + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/LocalLLMProvider.java b/iped-engine/src/main/java/iped/engine/rag/LocalLLMProvider.java new file mode 100644 index 0000000000..860901502b --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/LocalLLMProvider.java @@ -0,0 +1,149 @@ +package iped.engine.rag; + +/** + * Local LLM provider implementation connecting to Ollama chat endpoints. + * + * @author Rui Sant'Ana Junior + */ +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +public class LocalLLMProvider implements LLMProvider { + + private final String endpoint; + private final String model; + private final HttpClient client; + + public LocalLLMProvider(String endpoint, String model) { + this.endpoint = endpoint; + this.model = model; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + @Override + public String generateAnswer(String question, String context) throws IOException, InterruptedException { + return generateAnswer(question, context, java.util.Collections.emptyList()); + } + + @Override + public String generateAnswer(String question, String context, java.util.List history) throws IOException, InterruptedException { + String systemPrompt = getSystemPrompt(question); + + int contextWindow = 4096; + try { + RAGService rag = RAGService.getInstance(); + if (rag != null && rag.getConfig() != null) { + contextWindow = rag.getConfig().getLlmContextWindow(); + } + } catch (Exception ignore) {} + + StringBuilder messagesJson = new StringBuilder(); + messagesJson.append("["); + + // 1. Add system prompt as a system message + messagesJson.append("{\"role\":\"system\",\"content\":").append(LocalEmbeddingProvider.toJsonString(systemPrompt)).append("}"); + + // 2. Add history turns + if (history != null) { + for (RAGService.HistoryTurn turn : history) { + messagesJson.append(",{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString("Question: " + turn.question)).append("}"); + messagesJson.append(",{\"role\":\"assistant\",\"content\":").append(LocalEmbeddingProvider.toJsonString(turn.answer)).append("}"); + } + } + + // 3. Add clean user prompt + String userPrompt = "Context:\n" + context + "\n\nQuestion: " + question + "\n\nAnswer:"; + messagesJson.append(",{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString(userPrompt)).append("}"); + messagesJson.append("]"); + + String requestBody = "{" + + "\"model\":\"" + model + "\"," + + "\"messages\":" + messagesJson.toString() + "," + + "\"options\":{\"num_ctx\":" + contextWindow + ",\"temperature\":0.0}," + + "\"stream\":false" + + "}"; + + String resolvedUri = endpoint; + if (resolvedUri.contains("/v1")) { + resolvedUri = resolvedUri.replace("/v1", "/api/chat"); + } else if (!resolvedUri.endsWith("/api/chat")) { + resolvedUri = resolvedUri.endsWith("/") ? resolvedUri + "api/chat" : resolvedUri + "/api/chat"; + } + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(resolvedUri)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(300)) + .build(); + + HttpResponse response = sendPrivileged(client, request); + + if (response.statusCode() != 200) { + throw new IOException("Local LLM service returned status code " + response.statusCode() + ": " + response.body()); + } + + return parseChatCompletionResponse(response.body()); + } + + public static String parseChatCompletionResponse(String json) throws IOException { + int index = json.indexOf("\"content\""); + if (index == -1) { + throw new IOException("Failed to parse LLM response. Could not find 'content' key in response JSON: " + json); + } + index = json.indexOf(":", index); + if (index == -1) { + throw new IOException("Failed to parse LLM response. Invalid JSON format near 'content' in: " + json); + } + index = json.indexOf("\"", index); + if (index == -1) { + throw new IOException("Failed to parse LLM response. Could not find start of content string in: " + json); + } + + StringBuilder sb = new StringBuilder(); + boolean escaped = false; + for (int i = index + 1; i < json.length(); i++) { + char c = json.charAt(i); + if (escaped) { + switch (c) { + case 'n': sb.append('\n'); break; + case 'r': sb.append('\r'); break; + case 't': sb.append('\t'); break; + case 'b': sb.append('\b'); break; + case 'f': sb.append('\f'); break; + case '"': sb.append('"'); break; + case '\\': sb.append('\\'); break; + default: sb.append(c); + } + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + return sb.toString(); + } else { + sb.append(c); + } + } + throw new IOException("Failed to parse LLM response. Unclosed content string in response JSON: " + json); + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/RAGService.java b/iped-engine/src/main/java/iped/engine/rag/RAGService.java new file mode 100644 index 0000000000..c238ec5949 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/RAGService.java @@ -0,0 +1,1058 @@ +package iped.engine.rag; + +/** + * Core engine service for native IPED RAG (Retrieval-Augmented Generation), + * hybrid search (BM25 + KNN RRF), embedding caching, and multi-turn chat management. + * + * @author Rui Sant'Ana Junior + */ +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.security.MessageDigest; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.lang.reflect.Type; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import com.google.gson.Gson; +import com.google.gson.reflect.TypeToken; + +import org.apache.http.HttpHost; +import org.apache.http.auth.AuthScope; +import org.apache.http.auth.UsernamePasswordCredentials; +import org.apache.http.client.CredentialsProvider; +import org.apache.http.conn.ssl.NoopHostnameVerifier; +import org.apache.http.impl.client.BasicCredentialsProvider; +import iped.engine.util.SSLFix; +import org.apache.lucene.document.Document; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.client.RequestOptions; +import org.opensearch.client.RestClient; +import org.opensearch.client.RestClientBuilder; + +import org.opensearch.client.RestHighLevelClient; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.builder.SearchSourceBuilder; + +import iped.engine.config.ConfigurationManager; +import iped.engine.config.ElasticSearchTaskConfig; +import iped.engine.config.RAGConfig; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +public class RAGService { + + private static final Logger LOGGER = LoggerFactory.getLogger(RAGService.class); + private static final org.apache.logging.log4j.Level MSG = org.apache.logging.log4j.Level.getLevel("MSG"); + private static final org.apache.logging.log4j.Logger LOG4J_LOGGER = org.apache.logging.log4j.LogManager.getLogger(RAGService.class); + private static RAGService instance; + + private final File caseDir; + private final RAGConfig config; + private final EmbeddingProvider embeddingProvider; + private final LLMProvider llmProvider; + private final boolean available; + private Connection dbConnection; + private RestHighLevelClient openSearchClient; + private String indexName; + + /** Lucene IndexSearcher injected by the SearchApp when a case is opened (lucene mode only). */ + private volatile IndexSearcher luceneSearcher; + + /** Lucene Analyzer injected by the SearchApp when a case is opened. */ + private volatile org.apache.lucene.analysis.Analyzer luceneAnalyzer; + + public static synchronized void initialize(File caseDir, RAGConfig config) { + if (instance == null || (caseDir != null && instance.caseDir != null && !instance.caseDir.equals(caseDir))) { + instance = new RAGService(caseDir, config); + } + } + + public static synchronized RAGService getInstance() { + return instance; + } + + protected RAGService(File caseDir, RAGConfig config) { + this.caseDir = caseDir; + this.config = config; + + // Initialize Embedding Provider + String embProv = config.getEmbeddingProvider(); + if ("local".equalsIgnoreCase(embProv)) { + this.embeddingProvider = new LocalEmbeddingProvider(config.getEmbeddingEndpoint(), config.getEmbeddingModel()); + } else if ("gemini".equalsIgnoreCase(embProv)) { + this.embeddingProvider = new GeminiEmbeddingProvider(config.getEmbeddingEndpoint(), config.getEmbeddingModel(), config.getLlmApiKey(), config.getEmbeddingDimensions()); + } else { + this.embeddingProvider = new RemoteEmbeddingProvider(config.getEmbeddingEndpoint(), config.getEmbeddingModel(), config.getLlmApiKey()); + } + + // Initialize LLM Provider + String provider = config.getLlmProvider(); + if ("local".equalsIgnoreCase(provider)) { + this.llmProvider = new LocalLLMProvider(config.getLlmEndpoint(), config.getLlmModel()); + } else if ("gemini".equalsIgnoreCase(provider)) { + this.llmProvider = new GeminiLLMProvider(config.getLlmEndpoint(), config.getLlmModel(), config.getLlmApiKey()); + } else if ("claude".equalsIgnoreCase(provider)) { + this.llmProvider = new ClaudeLLMProvider(config.getLlmEndpoint(), config.getLlmModel(), config.getLlmApiKey()); + } else { + this.llmProvider = new RemoteLLMProvider(config.getLlmEndpoint(), config.getLlmModel(), config.getLlmApiKey()); + } + + // Resolve Index Name + if (caseDir != null) { + this.indexName = caseDir.getName(); + } + + // Perform health check on embedding provider and vector store at startup + boolean embOk = false; + if (config != null && config.isEnabled()) { + try { + embOk = checkEmbeddingHealth(config); + } catch (Exception e) { + embOk = false; + } + } + + if (embOk) { + if ("opensearch".equalsIgnoreCase(config.getVectorStoreMode())) { + boolean osOk = checkOpenSearchHealth(config); + if (!osOk) { + String msg = "RAG: OpenSearch cluster unreachable. Falling back to local Lucene vector store."; + if (LOG4J_LOGGER != null && MSG != null) { + LOG4J_LOGGER.log(MSG, msg); + } else { + LOGGER.warn(msg); + } + config.setVectorStoreMode("lucene"); + } + } + this.available = true; + LOGGER.info("RAG embedding generation enabled."); + } else { + this.available = false; + if (config != null && config.isEnabled()) { + String msg = "RAG: Embedding service unreachable. Skipping embeddings."; + if (LOG4J_LOGGER != null && MSG != null) { + LOG4J_LOGGER.log(MSG, msg); + } else { + LOGGER.warn(msg); + } + } + } + } + + private boolean checkOpenSearchHealth(RAGConfig config) { + if (!"opensearch".equalsIgnoreCase(config.getVectorStoreMode())) { + return true; + } + try { + RestHighLevelClient client = getOpenSearchClient(); + if (client != null) { + return client.ping(RequestOptions.DEFAULT); + } + } catch (Exception e) { + return false; + } + return false; + } + + public boolean isAvailable() { + return available; + } + + private static boolean checkEmbeddingHealth(RAGConfig config) { + String endpoint = config.getEmbeddingEndpoint(); + if (endpoint == null || endpoint.trim().isEmpty()) { + return false; + } + try { + java.net.http.HttpClient client = java.net.http.HttpClient.newBuilder() + .connectTimeout(java.time.Duration.ofSeconds(2)) + .build(); + + String pingUrl = endpoint; + if (endpoint.contains("/api/embeddings")) { + pingUrl = endpoint.substring(0, endpoint.indexOf("/api/embeddings")); + } else if (endpoint.contains("/v1/embeddings")) { + pingUrl = endpoint.substring(0, endpoint.indexOf("/v1/embeddings")); + } + if (pingUrl.trim().isEmpty()) { + pingUrl = endpoint; + } + + java.net.http.HttpRequest request = java.net.http.HttpRequest.newBuilder() + .uri(java.net.URI.create(pingUrl)) + .GET() + .timeout(java.time.Duration.ofSeconds(2)) + .build(); + + java.net.http.HttpResponse response = sendPrivilegedPing(client, request); + return response.statusCode() < 500; + } catch (Exception e) { + return false; + } + } + + private static java.net.http.HttpResponse sendPrivilegedPing(java.net.http.HttpClient client, java.net.http.HttpRequest request) throws Exception { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, java.net.http.HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + throw e.getException(); + } + } + + public RAGConfig getConfig() { + return config; + } + + public EmbeddingProvider getEmbeddingProvider() { + return embeddingProvider; + } + + public LLMProvider getLlmProvider() { + return llmProvider; + } + + /** + * Injects the Lucene IndexSearcher from the opened case. Must be called by + * the SearchApp after the case index is opened, when vectorStoreMode=lucene. + * + * @param searcher the IndexSearcher from IPEDSource.getSearcher() + */ + public void setLuceneSearcher(IndexSearcher searcher) { + if (searcher != null) { + searcher.setSimilarity(new org.apache.lucene.search.similarities.BM25Similarity()); + } + this.luceneSearcher = searcher; + } + + public void setLuceneAnalyzer(org.apache.lucene.analysis.Analyzer analyzer) { + this.luceneAnalyzer = analyzer; + } + + private File resolveDbFile(File baseDir) { + if (baseDir == null) return new File("iped/data/embedding_cache.db"); + + // Preserve legacy DB location if it already exists in an old case + File oldDb = new File(baseDir, "storage/embedding_cache.db"); + if (oldDb.exists()) return oldDb; + File oldIpedDb = new File(baseDir, "iped/storage/embedding_cache.db"); + if (oldIpedDb.exists()) return oldIpedDb; + File oldIpedDataDb = new File(baseDir, "iped/data/embedding_cache.db"); + if (oldIpedDataDb.exists()) return oldIpedDataDb; + + // Always resolve into the internal iped/ folder + File ipedDir = baseDir; + if (!baseDir.getName().equalsIgnoreCase("iped") && new File(baseDir, "iped").isDirectory()) { + ipedDir = new File(baseDir, "iped"); + } + + File dataDir = new File(ipedDir, "data"); + return new File(dataDir, "embedding_cache.db"); + } + + private synchronized Connection getConnection() throws SQLException { + if (dbConnection == null || dbConnection.isClosed()) { + File dbFile = resolveDbFile(caseDir); + dbFile.getParentFile().mkdirs(); + org.sqlite.SQLiteConfig sqliteConfig = new org.sqlite.SQLiteConfig(); + sqliteConfig.setJournalMode(org.sqlite.SQLiteConfig.JournalMode.WAL); + dbConnection = sqliteConfig.createConnection("jdbc:sqlite:" + dbFile.getAbsolutePath()); + try (Statement stmt = dbConnection.createStatement()) { + stmt.executeUpdate("CREATE TABLE IF NOT EXISTS embedding_cache (" + + "text_hash TEXT, " + + "model_name TEXT, " + + "embedding BLOB, " + + "PRIMARY KEY(text_hash, model_name))"); + stmt.executeUpdate("CREATE TABLE IF NOT EXISTS chat_history (" + + "id INTEGER PRIMARY KEY AUTOINCREMENT, " + + "timestamp INTEGER, " + + "session_id TEXT, " + + "question TEXT, " + + "answer_html TEXT, " + + "sources_json TEXT)"); + stmt.executeUpdate("CREATE TABLE IF NOT EXISTS chat_session_titles (" + + "session_id TEXT PRIMARY KEY, " + + "title TEXT)"); + } + } + return dbConnection; + } + + public synchronized float[] getCachedEmbedding(String text, String modelName) { + if (text == null || text.isBlank()) { + return null; + } + String hash = computeSHA256(text); + String sql = "SELECT embedding FROM embedding_cache WHERE text_hash = ? AND model_name = ?"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, hash); + pstmt.setString(2, modelName); + try (ResultSet rs = pstmt.executeQuery()) { + if (rs.next()) { + byte[] bytes = rs.getBytes(1); + if (bytes != null) { + return byteArrayToFloatArray(bytes); + } + } + } + } catch (SQLException e) { + // Log or ignore cache read errors + } + return null; + } + + public synchronized void cacheEmbedding(String text, String modelName, float[] embedding) { + String hash = computeSHA256(text); + String sql = "INSERT OR REPLACE INTO embedding_cache (text_hash, model_name, embedding) VALUES (?, ?, ?)"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, hash); + pstmt.setString(2, modelName); + pstmt.setBytes(3, floatArrayToByteArray(embedding)); + pstmt.executeUpdate(); + } catch (SQLException e) { + // Log or ignore cache write errors + } + } + + public synchronized long saveHistoryEntry(String sessionId, String question, String answerHtml, List sources) { + String sql = "INSERT INTO chat_history (timestamp, session_id, question, answer_html, sources_json) VALUES (?, ?, ?, ?, ?)"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS)) { + pstmt.setLong(1, System.currentTimeMillis()); + pstmt.setString(2, sessionId); + pstmt.setString(3, question); + pstmt.setString(4, answerHtml); + String json = new Gson().toJson(sources); + pstmt.setString(5, json); + pstmt.executeUpdate(); + try (ResultSet rs = pstmt.getGeneratedKeys()) { + if (rs.next()) { + return rs.getLong(1); + } + } + } catch (SQLException e) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error saving chat history entry", e); + } + return -1; + } + + public synchronized List loadHistoryEntries() { + List list = new ArrayList<>(); + String sql = "SELECT id, session_id, question, answer_html, sources_json FROM chat_history ORDER BY id ASC"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql); + ResultSet rs = pstmt.executeQuery()) { + Gson gson = new Gson(); + Type type = new TypeToken>(){}.getType(); + while (rs.next()) { + long id = rs.getLong(1); + String sessionId = rs.getString(2); + String question = rs.getString(3); + String answerHtml = rs.getString(4); + String sourcesJson = rs.getString(5); + List sources = gson.fromJson(sourcesJson, type); + list.add(new HistoryEntryRecord(id, sessionId, question, answerHtml, sources)); + } + } catch (SQLException e) { + // If table layout doesn't match (missing session_id on legacy tables), recreate it + try (Connection conn = getConnection(); + Statement stmt = conn.createStatement()) { + stmt.executeUpdate("DROP TABLE IF EXISTS chat_history"); + stmt.executeUpdate("CREATE TABLE IF NOT EXISTS chat_history (" + + "id INTEGER PRIMARY KEY AUTOINCREMENT, " + + "timestamp INTEGER, " + + "session_id TEXT, " + + "question TEXT, " + + "answer_html TEXT, " + + "sources_json TEXT)"); + } catch (SQLException ex) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error recreating chat_history table", ex); + } + } + return list; + } + + public synchronized void deleteHistorySession(String sessionId) { + String sql = "DELETE FROM chat_history WHERE session_id = ?"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, sessionId); + pstmt.executeUpdate(); + } catch (SQLException e) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error deleting chat history session", e); + } + String sqlTitle = "DELETE FROM chat_session_titles WHERE session_id = ?"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sqlTitle)) { + pstmt.setString(1, sessionId); + pstmt.executeUpdate(); + } catch (SQLException e) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error deleting chat session title", e); + } + } + + public synchronized void updateSessionTitle(String sessionId, String title) { + String sql = "INSERT OR REPLACE INTO chat_session_titles (session_id, title) VALUES (?, ?)"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, sessionId); + pstmt.setString(2, title); + pstmt.executeUpdate(); + } catch (SQLException e) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error updating chat session title", e); + } + } + + public synchronized Map loadSessionTitles() { + Map titles = new HashMap<>(); + String sql = "SELECT session_id, title FROM chat_session_titles"; + try (Connection conn = getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql); + ResultSet rs = pstmt.executeQuery()) { + while (rs.next()) { + titles.put(rs.getString(1), rs.getString(2)); + } + } catch (SQLException e) { + org.slf4j.LoggerFactory.getLogger(RAGService.class).error("Error loading chat session titles", e); + } + return titles; + } + + public static class HistoryEntryRecord { + public final long id; + public final String sessionId; + public final String question; + public final String answerHtml; + public final List displaySources; + + public HistoryEntryRecord(long id, String sessionId, String question, String answerHtml, List displaySources) { + this.id = id; + this.sessionId = sessionId; + this.question = question; + this.answerHtml = answerHtml; + this.displaySources = displaySources; + } + } + + public static class HistoryTurn { + public final String question; + public final String answer; + + public HistoryTurn(String question, String answer) { + this.question = question; + this.answer = answer; + } + } + + + public float[] getEmbedding(String text) throws IOException, InterruptedException { + String modelName = config.getEmbeddingModel(); + float[] vector = getCachedEmbedding(text, modelName); + if (vector == null && isAvailable()) { + vector = embeddingProvider.generateEmbedding(text); + if (vector != null) { + cacheEmbedding(text, modelName, vector); + } + } + return vector; + } + + public List getEmbeddings(List texts) throws IOException, InterruptedException { + if (texts == null || texts.isEmpty()) { + return java.util.Collections.emptyList(); + } + String modelName = config.getEmbeddingModel(); + List result = new ArrayList<>(texts.size()); + List missingIndices = new ArrayList<>(); + List missingTexts = new ArrayList<>(); + + for (int i = 0; i < texts.size(); i++) { + String txt = texts.get(i); + float[] vec = getCachedEmbedding(txt, modelName); + result.add(vec); + if (vec == null) { + missingIndices.add(i); + missingTexts.add(txt); + } + } + + if (!missingTexts.isEmpty() && isAvailable()) { + List newVecs = embeddingProvider.generateEmbeddings(missingTexts); + if (newVecs != null && newVecs.size() == missingTexts.size()) { + for (int k = 0; k < missingTexts.size(); k++) { + int idx = missingIndices.get(k); + float[] vec = newVecs.get(k); + result.set(idx, vec); + if (vec != null) { + cacheEmbedding(missingTexts.get(k), modelName, vec); + } + } + } else { + for (int k = 0; k < missingTexts.size(); k++) { + int idx = missingIndices.get(k); + float[] vec = getEmbedding(missingTexts.get(k)); + result.set(idx, vec); + } + } + } + return result; + } + + private synchronized RestHighLevelClient getOpenSearchClient() throws IOException { + if (openSearchClient == null) { + ElasticSearchTaskConfig esConfig = ConfigurationManager.get().findObject(ElasticSearchTaskConfig.class); + if (esConfig == null) { + throw new IOException("ElasticSearchConfig.txt is not loaded."); + } + + String protocol = esConfig.getProtocol(); + String hostProp = esConfig.getHost(); + int port = esConfig.getPort(); + + String[] hosts = hostProp.split("[,;\\s]+"); + HttpHost[] httpHosts = new HttpHost[hosts.length]; + for (int i = 0; i < hosts.length; i++) { + httpHosts[i] = new HttpHost(hosts[i].trim(), port, protocol); + } + + RestClientBuilder clientBuilder = RestClient.builder(httpHosts) + .setRequestConfigCallback(requestConfigBuilder -> requestConfigBuilder + .setConnectTimeout(esConfig.getConnect_timeout()) + .setSocketTimeout(esConfig.getTimeout_millis())); + + String user = System.getProperty("elastic.user"); + String password = System.getProperty("elastic.password"); + + if (user == null || password == null) { + String cmdFields = System.getenv("elastic"); + if (cmdFields == null) { + cmdFields = System.getProperty("elastic"); + } + if (cmdFields != null) { + String[] entries = cmdFields.split(";"); + for (String entry : entries) { + String[] pair = entry.split(":", 2); + if (pair.length == 2) { + if ("user".equals(pair[0])) user = pair[1]; + else if ("password".equals(pair[0])) password = pair[1]; + else if ("indexName".equals(pair[0])) this.indexName = pair[1]; + } + } + } + } + + if (user == null) { + user = esConfig.getUsername(); + } + if (password == null) { + password = esConfig.getPassword(); + } + + final String finalUser = user; + final String finalPassword = password; + clientBuilder.setHttpClientConfigCallback(httpClientBuilder -> { + if (finalUser != null && finalPassword != null) { + CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); + credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(finalUser, finalPassword)); + httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider); + } + if (!esConfig.getValidateSSL()) { + httpClientBuilder.setSSLContext(SSLFix.getUnsecureSSLContext()).setSSLHostnameVerifier(NoopHostnameVerifier.INSTANCE); + } + return httpClientBuilder; + }); + + openSearchClient = new RestHighLevelClient(clientBuilder); + } + return openSearchClient; + } + + public List performVectorSearch(float[] queryVector, int k) throws IOException { + return performVectorSearch(queryVector, k, null); + } + + public List performVectorSearch(float[] queryVector, int k, java.util.Collection allowedItemIds) throws IOException { + if ("lucene".equalsIgnoreCase(config.getVectorStoreMode())) { + return performVectorSearchLucene(queryVector, k, allowedItemIds); + } + return performVectorSearchOpenSearch(queryVector, k, allowedItemIds); + } + + // --- Lucene KNN backend --- + + private List performVectorSearchLucene(float[] queryVector, int k, + java.util.Collection allowedItemIds) throws IOException { + IndexSearcher searcher = luceneSearcher; + if (searcher == null) { + throw new IOException("Lucene IndexSearcher not initialised. " + + "Call RAGService.setLuceneSearcher() after opening the case."); + } + + // Retrieve k*5 candidates then re-rank (consistent with OpenSearch behaviour) + int candidates = k * 5; + Query knnQuery = new KnnVectorQuery( + iped.engine.task.index.IndexTask.CONTENT_EMBEDDING, queryVector, candidates); + + Query finalQuery = knnQuery; + if (allowedItemIds != null && !allowedItemIds.isEmpty()) { + BooleanQuery.Builder b = new BooleanQuery.Builder(); + b.add(knnQuery, BooleanClause.Occur.MUST); + // Filter: fragment's parent ID must match one of the allowed item IDs. + BooleanQuery.Builder idFilter = new BooleanQuery.Builder(); + for (String id : allowedItemIds) { + idFilter.add(new TermQuery( + new Term("fragParentNumericId", id)), + BooleanClause.Occur.SHOULD); + } + b.add(idFilter.build(), BooleanClause.Occur.MUST); + finalQuery = b.build(); + } + + TopDocs topDocs = searcher.search(finalQuery, Math.min(k, 1000)); + + List hits = new ArrayList<>(); + for (ScoreDoc scoreDoc : topDocs.scoreDocs) { + Document doc = searcher.doc(scoreDoc.doc); + String fragTrackId = doc.get(iped.engine.task.index.IndexTask.FRAG_TRACK_ID); + String content = doc.get(iped.engine.task.index.IndexTask.CONTENT_STORED); + if (fragTrackId == null || content == null) { + continue; + } + // parentGlobalId is everything before the last '#' + int sep = fragTrackId.lastIndexOf('#'); + String parentId = sep > 0 ? fragTrackId.substring(0, sep) : fragTrackId; + + String parentNumericId = doc.get("fragParentNumericId"); + + RAGSourceDoc ragDoc = new RAGSourceDoc(); + ragDoc.itemId = parentNumericId != null ? parentNumericId : parentId; + ragDoc.content = content; + ragDoc.contenttrackID = fragTrackId; + ragDoc.score = scoreDoc.score; + hits.add(ragDoc); + } + return hits; + } + + // --- OpenSearch backend (legacy) --- + + private List performVectorSearchOpenSearch(float[] queryVector, int k, + java.util.Collection allowedItemIds) throws IOException { + RestHighLevelClient client = getOpenSearchClient(); + SearchRequest searchRequest = new SearchRequest(indexName); + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + + StringBuilder vecJson = new StringBuilder("["); + for (int i = 0; i < queryVector.length; i++) { + vecJson.append(queryVector[i]); + if (i < queryVector.length - 1) vecJson.append(","); + } + vecJson.append("]"); + + String knnJson = "{\"knn\": {\"content_embedding\": {\"vector\": " + vecJson.toString() + ", \"k\": " + k + "}}}"; + QueryBuilder knnQuery = QueryBuilders.wrapperQuery(knnJson); + QueryBuilder queryBuilder = knnQuery; + + if (allowedItemIds != null && !allowedItemIds.isEmpty()) { + queryBuilder = QueryBuilders.boolQuery() + .must(knnQuery) + .filter(QueryBuilders.termsQuery("id", allowedItemIds)); + } + + sourceBuilder.query(queryBuilder); + sourceBuilder.size(k); + sourceBuilder.fetchSource(new String[]{"id", "content", "contenttrackID"}, null); + searchRequest.source(sourceBuilder); + + SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT); + List hits = new ArrayList<>(); + for (SearchHit hit : response.getHits().getHits()) { + Map source = hit.getSourceAsMap(); + if (source != null) { + RAGSourceDoc doc = new RAGSourceDoc(); + doc.itemId = source.get("id") != null ? String.valueOf(source.get("id")) : null; + doc.content = source.get("content") != null ? String.valueOf(source.get("content")) : null; + doc.contenttrackID = source.get("contenttrackID") != null ? String.valueOf(source.get("contenttrackID")) : null; + doc.score = hit.getScore(); + hits.add(doc); + } + } + return hits; + } + + public List performLexicalSearch(String question, int k) throws IOException { + return performLexicalSearch(question, k, null); + } + + public List performLexicalSearch(String question, int k, java.util.Collection allowedItemIds) throws IOException { + if ("lucene".equalsIgnoreCase(config.getVectorStoreMode())) { + return performLexicalSearchLucene(question, k, allowedItemIds); + } + return performLexicalSearchOpenSearch(question, k, allowedItemIds); + } + + // --- Lucene BM25 backend --- + + private List performLexicalSearchLucene(String question, int k, + java.util.Collection allowedItemIds) throws IOException { + IndexSearcher searcher = luceneSearcher; + if (searcher == null) { + throw new IOException("Lucene IndexSearcher not initialised. " + + "Call RAGService.setLuceneSearcher() after opening the case."); + } + + // Only search over fragment documents that have been embedded + Query hasEmbeddingFilter = new TermQuery( + new Term(iped.engine.task.index.IndexTask.HAS_EMBEDDING_FRAG, "1")); + + String cleanedQuestion = question; + if (cleanedQuestion != null) { + cleanedQuestion = cleanedQuestion.trim(); + while (cleanedQuestion.endsWith("?") || cleanedQuestion.endsWith(".") || cleanedQuestion.endsWith("!")) { + cleanedQuestion = cleanedQuestion.substring(0, cleanedQuestion.length() - 1).trim(); + } + cleanedQuestion = cleanedQuestion.replaceAll("[\\+\\-\\&\\|\\!\\(\\)\\{\\}\\[\\]\\^\\\"\\~\\*\\?\\:\\\\\\/]", " "); + cleanedQuestion = cleanedQuestion.trim().replaceAll("\\s+", " "); + } + + Query wordsQuery; + try { + org.apache.lucene.analysis.Analyzer analyzer = luceneAnalyzer; + if (analyzer == null) { + analyzer = new org.apache.lucene.analysis.standard.StandardAnalyzer(); + } + org.apache.lucene.queryparser.flexible.standard.StandardQueryParser parser = + new org.apache.lucene.queryparser.flexible.standard.StandardQueryParser(analyzer); + parser.setMultiFields(new String[]{ iped.properties.BasicProps.CONTENT }); + wordsQuery = parser.parse(cleanedQuestion, iped.properties.BasicProps.CONTENT); + + if (wordsQuery instanceof BooleanQuery) { + BooleanQuery bq = (BooleanQuery) wordsQuery; + int clauseCount = bq.clauses().size(); + int minMatch = Math.max(1, (int) (clauseCount * 0.4)); + + BooleanQuery.Builder bqBuilder = new BooleanQuery.Builder(); + for (BooleanClause clause : bq.clauses()) { + bqBuilder.add(clause); + } + bqBuilder.setMinimumNumberShouldMatch(minMatch); + wordsQuery = bqBuilder.build(); + } + } catch (org.apache.lucene.queryparser.flexible.core.QueryNodeException e) { + throw new IOException("Failed to parse lexical query: " + question, e); + } + + BooleanQuery.Builder boolBuilder = new BooleanQuery.Builder(); + boolBuilder.add(hasEmbeddingFilter, BooleanClause.Occur.MUST); + boolBuilder.add(wordsQuery, BooleanClause.Occur.MUST); + + if (allowedItemIds != null && !allowedItemIds.isEmpty()) { + BooleanQuery.Builder idFilter = new BooleanQuery.Builder(); + for (String id : allowedItemIds) { + idFilter.add(new TermQuery( + new Term("fragParentNumericId", id)), + BooleanClause.Occur.SHOULD); + } + boolBuilder.add(idFilter.build(), BooleanClause.Occur.MUST); + } + + TopDocs topDocs = searcher.search(boolBuilder.build(), Math.min(k, 1000)); + + List hits = new ArrayList<>(); + for (ScoreDoc scoreDoc : topDocs.scoreDocs) { + Document doc = searcher.doc(scoreDoc.doc); + String fragTrackId = doc.get(iped.engine.task.index.IndexTask.FRAG_TRACK_ID); + String content = doc.get(iped.engine.task.index.IndexTask.CONTENT_STORED); + if (fragTrackId == null || content == null) { + continue; + } + int sep = fragTrackId.lastIndexOf('#'); + String parentId = sep > 0 ? fragTrackId.substring(0, sep) : fragTrackId; + + String parentNumericId = doc.get("fragParentNumericId"); + + RAGSourceDoc ragDoc = new RAGSourceDoc(); + ragDoc.itemId = parentNumericId != null ? parentNumericId : parentId; + ragDoc.content = content; + ragDoc.contenttrackID = fragTrackId; + ragDoc.score = scoreDoc.score; + hits.add(ragDoc); + } + return hits; + } + + // --- OpenSearch BM25 backend (legacy) --- + + private List performLexicalSearchOpenSearch(String question, int k, + java.util.Collection allowedItemIds) throws IOException { + RestHighLevelClient client = getOpenSearchClient(); + SearchRequest searchRequest = new SearchRequest(indexName); + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + + QueryBuilder matchQuery = QueryBuilders.matchQuery("content", question) + .minimumShouldMatch("40%"); + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery() + .must(matchQuery) + .filter(QueryBuilders.termQuery(iped.engine.task.index.IndexTask.HAS_EMBEDDING_FRAG, "1")); + + if (allowedItemIds != null && !allowedItemIds.isEmpty()) { + boolQuery.filter(QueryBuilders.termsQuery("id", allowedItemIds)); + } + + sourceBuilder.query(boolQuery); + sourceBuilder.size(k); + sourceBuilder.fetchSource(new String[]{"id", "content", "contenttrackID"}, null); + searchRequest.source(sourceBuilder); + + SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT); + List hits = new ArrayList<>(); + for (SearchHit hit : response.getHits().getHits()) { + Map source = hit.getSourceAsMap(); + if (source != null) { + RAGSourceDoc doc = new RAGSourceDoc(); + doc.itemId = source.get("id") != null ? String.valueOf(source.get("id")) : null; + doc.content = source.get("content") != null ? String.valueOf(source.get("content")) : null; + doc.contenttrackID = source.get("contenttrackID") != null ? String.valueOf(source.get("contenttrackID")) : null; + doc.score = hit.getScore(); + hits.add(doc); + } + } + return hits; + } + + public List performHybridSearch(String question, float[] queryVector) throws IOException { + return performHybridSearch(question, queryVector, null); + } + + public List performHybridSearch(String question, float[] queryVector, java.util.Collection allowedItemIds) throws IOException { + int maxChunks = config.getMaxRetrievedChunks(); + List vectorHits = performVectorSearch(queryVector, maxChunks * 5, allowedItemIds); + List lexicalHits = performLexicalSearch(question, maxChunks * 5, allowedItemIds); + + // Raw similarity fallback: if there are no lexical matches and the best vector + // match has a raw similarity score below the threshold, treat the whole query as noise. + // This prevents RRF normalization from boosting weak/spurious vector hits to 1.0. + float rawThreshold = config.getChunkSimilarityThreshold(); + if (lexicalHits.isEmpty() && !vectorHits.isEmpty() && rawThreshold > 0.0f) { + double topRawScore = vectorHits.get(0).score; + if (topRawScore < rawThreshold) { + return new ArrayList<>(); + } + } + + Map rrfScores = new HashMap<>(); + Map docMap = new HashMap<>(); + + // Rank vector hits - Weighted by vectorSearchWeight. + // Apply raw threshold before RRF: only vector hits >= threshold participate. + double vectorWeight = config.getVectorSearchWeight(); + double lexicalWeight = config.getLexicalSearchWeight(); + float threshold = config.getChunkSimilarityThreshold(); + int rank = 1; + for (RAGSourceDoc doc : vectorHits) { + if (threshold > 0.0f && doc.score < threshold) { + continue; // Apply threshold before RRF + } + rrfScores.put(doc.contenttrackID, rrfScores.getOrDefault(doc.contenttrackID, 0.0) + vectorWeight * (1.0 / (60.0 + rank))); + docMap.put(doc.contenttrackID, doc); + rank++; + } + + // Rank lexical hits - Weighted by lexicalSearchWeight + rank = 1; + for (RAGSourceDoc doc : lexicalHits) { + rrfScores.put(doc.contenttrackID, rrfScores.getOrDefault(doc.contenttrackID, 0.0) + lexicalWeight * (1.0 / (60.0 + rank))); + if (!docMap.containsKey(doc.contenttrackID)) { + docMap.put(doc.contenttrackID, doc); + } + rank++; + } + + // Sort by RRF score descending + List sortedIds = new ArrayList<>(rrfScores.keySet()); + sortedIds.sort((id1, id2) -> Double.compare(rrfScores.get(id2), rrfScores.get(id1))); + + // Collect top-N results. + // Scores are normalized so the best document = 1.0 and others are proportional + // for displaying relative scores in the UI. + List result = new ArrayList<>(); + int limit = Math.min(sortedIds.size(), maxChunks); + double topRrfScore = limit > 0 ? rrfScores.get(sortedIds.get(0)) : 1.0; + for (int i = 0; i < limit; i++) { + String id = sortedIds.get(i); + float normalizedScore = (float) (rrfScores.get(id) / topRrfScore); // [0.0, 1.0] + RAGSourceDoc doc = docMap.get(id); + doc.score = normalizedScore; + result.add(doc); + } + return result; + } + + + public float[] getAverageEmbeddingForItem(String parentNumericId) throws IOException { + if ("lucene".equalsIgnoreCase(config.getVectorStoreMode())) { + return getAverageEmbeddingForItemLucene(parentNumericId); + } + return getAverageEmbeddingForItemOpenSearch(parentNumericId); + } + + private float[] getAverageEmbeddingForItemLucene(String parentNumericId) throws IOException { + IndexSearcher searcher = luceneSearcher; + if (searcher == null) { + throw new IOException("Lucene IndexSearcher not initialised."); + } + Query filter = new TermQuery(new Term("fragParentNumericId", parentNumericId)); + TopDocs topDocs = searcher.search(filter, 200); + + for (ScoreDoc sd : topDocs.scoreDocs) { + Document doc = searcher.doc(sd.doc); + // KnnVectorField values are not stored as retrievable bytes in Lucene 9; + // we can only check presence via HAS_EMBEDDING_FRAG marker. + // Return null to signal that the caller must regenerate the embedding. + String hasEmb = doc.get(iped.engine.task.index.IndexTask.HAS_EMBEDDING_FRAG); + if ("1".equals(hasEmb)) { + // Embedding is present but cannot be retrieved from Lucene stored fields. + // Callers should use getEmbedding(text) on the stored content instead. + return null; + } + } + return null; + } + + private float[] getAverageEmbeddingForItemOpenSearch(String parentId) throws IOException { + RestHighLevelClient client = getOpenSearchClient(); + SearchRequest searchRequest = new SearchRequest(indexName); + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); + sourceBuilder.query(QueryBuilders.termQuery("document_content.parent", parentId)); + sourceBuilder.size(100); + sourceBuilder.fetchSource(new String[]{"content_embedding"}, null); + searchRequest.source(sourceBuilder); + + SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT); + List embeddings = new java.util.ArrayList<>(); + for (SearchHit hit : response.getHits().getHits()) { + Map source = hit.getSourceAsMap(); + if (source != null && source.containsKey("content_embedding")) { + Object embObj = source.get("content_embedding"); + if (embObj instanceof List) { + List list = (List) embObj; + float[] emb = new float[list.size()]; + for (int i = 0; i < list.size(); i++) { + emb[i] = ((Number) list.get(i)).floatValue(); + } + embeddings.add(emb); + } + } + } + + if (embeddings.isEmpty()) { + return null; + } + + int dim = embeddings.get(0).length; + float[] avg = new float[dim]; + for (float[] emb : embeddings) { + for (int i = 0; i < dim; i++) { + avg[i] += emb[i]; + } + } + for (int i = 0; i < dim; i++) { + avg[i] /= embeddings.size(); + } + return avg; + } + + public List findSimilarDocumentIds(float[] queryVector, int maxResults) throws IOException { + List fragmentHits = performVectorSearch(queryVector, maxResults * 3); + Map parentScores = new LinkedHashMap<>(); + for (RAGSourceDoc hit : fragmentHits) { + if (hit.itemId != null && !parentScores.containsKey(hit.itemId)) { + parentScores.put(hit.itemId, hit.score); + } + } + List sortedParentIds = new ArrayList<>(parentScores.keySet()); + if (sortedParentIds.size() > maxResults) { + return sortedParentIds.subList(0, maxResults); + } + return sortedParentIds; + } + + public void close() { + try { + if (dbConnection != null && !dbConnection.isClosed()) { + dbConnection.close(); + } + } catch (SQLException ignored) {} + try { + if (openSearchClient != null) { + openSearchClient.close(); + } + } catch (IOException ignored) {} + } + + public static class RAGSourceDoc { + public String itemId; + public String content; + public String contenttrackID; + public float score; + } + + private static String computeSHA256(String text) { + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + byte[] hash = digest.digest(text.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + StringBuilder hexString = new StringBuilder(); + for (byte b : hash) { + String hex = Integer.toHexString(0xff & b); + if (hex.length() == 1) hexString.append('0'); + hexString.append(hex); + } + return hexString.toString(); + } catch (Exception ex) { + throw new RuntimeException(ex); + } + } + + private static byte[] floatArrayToByteArray(float[] floats) { + ByteBuffer buffer = ByteBuffer.allocate(floats.length * 4); + for (float f : floats) { + buffer.putFloat(f); + } + return buffer.array(); + } + + private static float[] byteArrayToFloatArray(byte[] bytes) { + ByteBuffer buffer = ByteBuffer.wrap(bytes); + float[] floats = new float[bytes.length / 4]; + for (int i = 0; i < floats.length; i++) { + floats[i] = buffer.getFloat(); + } + return floats; + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/RemoteEmbeddingProvider.java b/iped-engine/src/main/java/iped/engine/rag/RemoteEmbeddingProvider.java new file mode 100644 index 0000000000..2ce7aad05b --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/RemoteEmbeddingProvider.java @@ -0,0 +1,95 @@ +package iped.engine.rag; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +/** + * Remote embedding provider implementation. + * + * @author Rui Sant'Ana Junior + */ +public class RemoteEmbeddingProvider implements EmbeddingProvider { + + private final String endpoint; + private final String model; + private final String apiKey; + private final HttpClient client; + + public RemoteEmbeddingProvider(String endpoint, String model, String apiKey) { + this.endpoint = endpoint; + this.model = model; + this.apiKey = apiKey; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + @Override + public float[] generateEmbedding(String text) throws IOException, InterruptedException { + String requestBody = "{\"model\":\"" + model + "\",\"input\":" + LocalEmbeddingProvider.toJsonString(text) + "}"; + + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder() + .uri(URI.create(endpoint)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(30)); + + if (endpoint.contains("azure.com")) { + requestBuilder.header("api-key", apiKey); + } else if (apiKey != null && !apiKey.trim().isEmpty()) { + requestBuilder.header("Authorization", "Bearer " + apiKey.trim()); + } + + HttpRequest request = requestBuilder.build(); + int maxRetries = 5; + long backoffMs = 2000; + + for (int attempt = 0; attempt < maxRetries; attempt++) { + try { + HttpResponse response = sendPrivileged(client, request); + + if (response.statusCode() == 200) { + return LocalEmbeddingProvider.parseEmbeddingsArray(response.body()); + } else if (response.statusCode() == 429) { + String retryAfter = response.headers().firstValue("Retry-After").orElse(""); + long sleepMs = backoffMs; + if (!retryAfter.isEmpty()) { + try { + sleepMs = Long.parseLong(retryAfter) * 1000; + } catch (NumberFormatException ignored) {} + } + Thread.sleep(sleepMs); + } else if (response.statusCode() >= 500) { + Thread.sleep(backoffMs); + } else { + throw new IOException("Remote embedding service returned status code " + response.statusCode() + ": " + response.body()); + } + } catch (IOException e) { + if (attempt == maxRetries - 1) { + throw e; + } + Thread.sleep(backoffMs); + } + backoffMs *= 2; + } + throw new IOException("Failed to generate embedding after " + maxRetries + " attempts due to rate limit or connection errors."); + } + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/rag/RemoteLLMProvider.java b/iped-engine/src/main/java/iped/engine/rag/RemoteLLMProvider.java new file mode 100644 index 0000000000..f67a5b2b32 --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/rag/RemoteLLMProvider.java @@ -0,0 +1,131 @@ +package iped.engine.rag; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +/** + * Remote LLM provider implementation. + * + * @author Rui Sant'Ana Junior + */ +public class RemoteLLMProvider implements LLMProvider { + + private final String endpoint; + private final String model; + private final String apiKey; + private final HttpClient client; + + public RemoteLLMProvider(String endpoint, String model, String apiKey) { + this.endpoint = endpoint; + this.model = model; + this.apiKey = apiKey; + this.client = HttpClient.newBuilder() + .connectTimeout(Duration.ofSeconds(15)) + .build(); + } + + @Override + public String generateAnswer(String question, String context) throws IOException, InterruptedException { + return generateAnswer(question, context, java.util.Collections.emptyList()); + } + + @Override + public String generateAnswer(String question, String context, java.util.List history) throws IOException, InterruptedException { + String systemPrompt = getSystemPrompt(question); + String userPrompt = "Context:\n" + context + "\n\nQuestion: " + question + "\n\nAnswer:"; + + StringBuilder messagesJson = new StringBuilder(); + messagesJson.append("["); + messagesJson.append("{\"role\":\"system\",\"content\":").append(LocalEmbeddingProvider.toJsonString(systemPrompt)).append("}"); + + if (history != null) { + for (RAGService.HistoryTurn turn : history) { + messagesJson.append(",{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString("Question: " + turn.question)).append("}"); + messagesJson.append(",{\"role\":\"assistant\",\"content\":").append(LocalEmbeddingProvider.toJsonString(turn.answer)).append("}"); + } + } + + messagesJson.append(",{\"role\":\"user\",\"content\":").append(LocalEmbeddingProvider.toJsonString(userPrompt)).append("}"); + messagesJson.append("]"); + + StringBuilder bodyJson = new StringBuilder(); + bodyJson.append("{\"model\":\"").append(model).append("\",\"messages\":").append(messagesJson.toString()); + if (!model.startsWith("o1") && !model.startsWith("o3")) { + bodyJson.append(",\"temperature\":0.0"); + } + bodyJson.append("}"); + String requestBody = bodyJson.toString(); + + String targetUrl = endpoint.endsWith("/chat/completions") ? endpoint : endpoint + "/chat/completions"; + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder() + .uri(URI.create(targetUrl)) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofSeconds(60)); + + if (endpoint.contains("azure.com")) { + requestBuilder.header("api-key", apiKey); + } else if (apiKey != null && !apiKey.trim().isEmpty()) { + requestBuilder.header("Authorization", "Bearer " + apiKey.trim()); + } + + HttpRequest request = requestBuilder.build(); + int maxRetries = 4; + long backoffMs = 2000; + int lastStatusCode = -1; + String lastResponseBody = ""; + + for (attempt = 0; attempt < maxRetries; attempt++) { + try { + HttpResponse response = sendPrivileged(client, request); + lastStatusCode = response.statusCode(); + lastResponseBody = response.body(); + + if (response.statusCode() == 200) { + return LocalLLMProvider.parseChatCompletionResponse(response.body()); + } else if (response.statusCode() == 429) { + String retryAfter = response.headers().firstValue("Retry-After").orElse(""); + long sleepMs = backoffMs; + if (!retryAfter.isEmpty()) { + try { + sleepMs = Long.parseLong(retryAfter) * 1000; + } catch (NumberFormatException ignored) {} + } + Thread.sleep(sleepMs); + } else if (response.statusCode() >= 500) { + Thread.sleep(backoffMs); + } else { + throw new IOException("Remote LLM service returned status code " + response.statusCode() + ": " + response.body()); + } + } catch (IOException e) { + if (attempt == maxRetries - 1) { + throw e; + } + Thread.sleep(backoffMs); + } + backoffMs *= 2; + } + throw new IOException("Failed to generate answer from LLM after " + maxRetries + " attempts. Last status: " + lastStatusCode + ", Response: " + lastResponseBody); + } + + // Declared attempt for loop scope clarity in case compilations check this + private int attempt; + + private static HttpResponse sendPrivileged(HttpClient client, HttpRequest request) throws IOException, InterruptedException { + try { + return java.security.AccessController.doPrivileged( + (java.security.PrivilegedExceptionAction>) () -> client.send(request, HttpResponse.BodyHandlers.ofString()) + ); + } catch (java.security.PrivilegedActionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) throw (IOException) cause; + if (cause instanceof InterruptedException) throw (InterruptedException) cause; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + throw new IOException(cause); + } + } +} diff --git a/iped-engine/src/main/java/iped/engine/task/index/ElasticSearchIndexTask.java b/iped-engine/src/main/java/iped/engine/task/index/ElasticSearchIndexTask.java index 68ef7cd97e..d76cd37975 100644 --- a/iped-engine/src/main/java/iped/engine/task/index/ElasticSearchIndexTask.java +++ b/iped-engine/src/main/java/iped/engine/task/index/ElasticSearchIndexTask.java @@ -54,6 +54,8 @@ import iped.engine.config.ConfigurationManager; import iped.engine.config.ElasticSearchTaskConfig; import iped.engine.config.IndexTaskConfig; +import iped.engine.config.RAGConfig; +import iped.engine.rag.RAGService; import iped.engine.io.FragmentingReader; import iped.engine.task.AbstractTask; import iped.engine.task.MinIOTask.MinIODataRef; @@ -121,13 +123,12 @@ public boolean isEnabled() { } public List> getConfigurables() { - return Arrays.asList(new ElasticSearchTaskConfig()); + return Arrays.asList(new ElasticSearchTaskConfig(), new RAGConfig()); } @Override public void init(ConfigurationManager configurationManager) throws Exception { - taskInstances.add(this); elasticConfig = configurationManager.findObject(ElasticSearchTaskConfig.class); retries = elasticConfig.getRetries(); @@ -136,6 +137,21 @@ public void init(ConfigurationManager configurationManager) throws Exception { return; } + RAGConfig ragConfig = configurationManager.findObject(RAGConfig.class); + if (ragConfig != null && ragConfig.isEnabled()) { + RAGService.initialize(output, ragConfig); + } + + // In lucene vector mode the OpenSearch client is not needed for indexing. + // Skip the connection entirely to avoid a ConnectException at startup. + if (ragConfig != null && ragConfig.isEnabled() + && "lucene".equalsIgnoreCase(ragConfig.getVectorStoreMode())) { + isEnabled = false; + LOGGER.info("ElasticSearchIndexTask disabled: vectorStoreMode=lucene, " + + "OpenSearch client not required."); + return; + } + if (client != null) { return; } @@ -146,12 +162,29 @@ public void init(ConfigurationManager configurationManager) throws Exception { parseCmdLineFields(System.getProperty(CMD_FIELDS_KEY)); parseCmdLineFields(args.getExtraParams().get(CMD_FIELDS_KEY)); + if (user == null) { + user = elasticConfig.getUsername(); + } + if (password == null) { + password = elasticConfig.getPassword(); + } + if (indexName == null) { indexName = output.getParentFile().getName(); } + String protocol = elasticConfig.getProtocol(); + String hostProp = elasticConfig.getHost(); + int port = elasticConfig.getPort(); + + String[] hosts = hostProp.split("[,;\\s]+"); + HttpHost[] httpHosts = new HttpHost[hosts.length]; + for (int i = 0; i < hosts.length; i++) { + httpHosts[i] = new HttpHost(hosts[i].trim(), port, protocol); + } + RestClientBuilder clientBuilder = RestClient - .builder(new HttpHost(elasticConfig.getHost(), elasticConfig.getPort(), elasticConfig.getProtocol())) + .builder(httpHosts) .setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() { @Override public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) { @@ -160,28 +193,36 @@ public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder reques } }); - if (user != null && password != null) { - CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); - credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(user, password)); - clientBuilder.setHttpClientConfigCallback(new HttpClientConfigCallback() { - @Override - public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) { - var custombuilder=httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider); - // disable validation if configured - if (!elasticConfig.getValidateSSL()) { - custombuilder.setSSLContext(SSLFix.getUnsecureSSLContext()).setSSLHostnameVerifier(NoopHostnameVerifier.INSTANCE); - } - return custombuilder; + final String finalUser = user; + final String finalPassword = password; + clientBuilder.setHttpClientConfigCallback(new HttpClientConfigCallback() { + @Override + public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) { + if (finalUser != null && finalPassword != null) { + CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); + credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(finalUser, finalPassword)); + httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider); } - }); - } + if (!elasticConfig.getValidateSSL()) { + httpClientBuilder.setSSLContext(SSLFix.getUnsecureSSLContext()).setSSLHostnameVerifier(NoopHostnameVerifier.INSTANCE); + } + return httpClientBuilder; + } + }); client = new RestHighLevelClient(clientBuilder); - boolean ping = client.ping(RequestOptions.DEFAULT); + boolean ping = false; + try { + ping = client.ping(RequestOptions.DEFAULT); + } catch (Exception e) { + ping = false; + } if (!ping) { - throw new IOException("ElasticSearch cluster at " + elasticConfig.getHost() + ":" + elasticConfig.getPort() - + " not responding to ping."); + LOGGER.warn("ElasticSearch/OpenSearch cluster at " + elasticConfig.getHost() + ":" + elasticConfig.getPort() + + " unreachable. Disabling cluster indexing task."); + isEnabled = false; + return; } if (args.isRestart()) { @@ -198,6 +239,7 @@ public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpCli throw new IPEDException("ElasticSearch index does not exist: " + indexName); } + taskInstances.add(this); } private void parseCmdLineFields(String cmdFields) { @@ -282,12 +324,12 @@ private void createFieldMappings(String indexName) throws IOException { properties.put(ExtraProperties.LOCATIONS, Collections.singletonMap("type", "geo_point")); properties.put("extraAttributes." + ImageSimilarityTask.IMAGE_FEATURES, Map.of("type", "knn_vector", "dimension", ImageSimilarity.numFeatures, "method", - Map.of("name", "hnsw", "space_type", "l2", "engine", "nmslib"))); + Map.of("name", "hnsw", "space_type", "l2", "engine", "lucene"))); // mapping faces as nested field var faces_mapping = new HashMap(); faces_mapping.put("face_encoding", Map.of("type", "knn_vector", "dimension", FACE_SIZE, "method", - Map.of("name", "hnsw", "space_type", "l2", "engine", "nmslib"))); + Map.of("name", "hnsw", "space_type", "l2", "engine", "lucene"))); faces_mapping.put("face_location", Collections.singletonMap("type", "short")); properties.put("faces", Map.of("type", "nested", "properties", faces_mapping)); @@ -299,6 +341,15 @@ private void createFieldMappings(String indexName) throws IOException { properties.put(BasicProps.CONTENT, contentMapping); + RAGConfig ragConfig = ConfigurationManager.get().findObject(RAGConfig.class); + if (ragConfig != null && ragConfig.isEnabled()) { + properties.put("content_embedding", Map.of("type", "knn_vector", "dimension", ragConfig.getEmbeddingDimensions(), "method", + Map.of("name", "hnsw", "space_type", "cosinesimil", "engine", "lucene"))); + properties.put("embedding_model", Collections.singletonMap("type", "keyword")); + properties.put("embedding_timestamp", Collections.singletonMap("type", "date")); + properties.put(IndexTask.HAS_EMBEDDING_FRAG, Collections.singletonMap("type", "keyword")); + } + // mapping the parent-child relation /* * "document_content": { "type": "join", "relations": { "document": "content" } @@ -403,52 +454,104 @@ protected synchronized void process(IItem item) throws Exception { } IndexTaskConfig indexConfig = ConfigurationManager.get().findObject(IndexTaskConfig.class); - FragmentingReader fragReader = new FragmentingReader(textReader, indexConfig.getTextSplitSize(), - indexConfig.getTextOverlapSize()); + int splitSize = indexConfig.getTextSplitSize(); + int overlapSize = indexConfig.getTextOverlapSize(); + RAGConfig ragConfig = ConfigurationManager.get().findObject(RAGConfig.class); + if (ragConfig != null && ragConfig.isEnabled()) { + String mimeType = item.getMediaType() != null ? item.getMediaType().toString() : ""; + if (ragConfig.shouldEmbed(item.getCategorySet(), mimeType)) { + splitSize = 4000; + overlapSize = 1000; + } + } + FragmentingReader fragReader = new FragmentingReader(textReader, splitSize, overlapSize); int fragNum = fragReader.estimateNumberOfFrags(); if (fragNum == -1) { fragNum = 1; } // used for parent items in elastic to store just metadata info - // globalID works like an 'UUID' and should be unique across cases String parentId = (String) item.getExtraAttribute(ExtraProperties.GLOBAL_ID); try { + RAGService ragService = RAGService.getInstance(); + boolean doRAG = ragService != null + && ragService.getConfig().isEnabled() + && ragService.isAvailable() + && "opensearch".equalsIgnoreCase(ragService.getConfig().getVectorStoreMode()); + boolean shouldEmbedItem = false; + if (doRAG) { + String mimeType = item.getMediaType() != null ? item.getMediaType().toString() : ""; + shouldEmbedItem = ragService.getConfig().shouldEmbed(item.getCategorySet(), mimeType); + if (shouldEmbedItem) { + item.setExtraAttribute(ExtraProperties.HAS_EMBEDDING, Boolean.TRUE.toString()); + } + } + // creates the father; XContentBuilder jsonMetadata = getJsonMetadataBuilder(item); IndexRequest parentIndexRequest = createIndexRequest(parentId, parentId, jsonMetadata); bulkRequest.add(parentIndexRequest); idToPath.put(parentId, item.getPath()); + List frags = new ArrayList<>(); + do { - // used for children items in elastic to store text content String contenttrackID = Util.generatetrackIDForTextFrag(parentId, fragNum); + String content = getStringFromReader(fragReader); + FragHolder fh = new FragHolder(); + fh.contenttrackID = contenttrackID; + fh.fragNum = fragNum--; + fh.content = content; + frags.add(fh); + } while (!Thread.currentThread().isInterrupted() && fragReader.nextFragment()); + + List embeddings = null; + if (doRAG && shouldEmbedItem && !frags.isEmpty()) { + List textsToEmbed = new ArrayList<>(); + for (FragHolder fh : frags) { + if (fh.content != null && !fh.content.trim().isEmpty()) { + textsToEmbed.add(fh.content); + } + } + if (!textsToEmbed.isEmpty()) { + try { + embeddings = ragService.getEmbeddings(textsToEmbed); + if (embeddings != null && !embeddings.isEmpty()) { + IndexTask.totalEmbeddedDocs.incrementAndGet(); + IndexTask.totalEmbeddedVectors.addAndGet(embeddings.size()); + } + } catch (Exception e) { + LOGGER.error("Failed to generate batch embeddings for item " + item.getId(), e); + } + } + } - // creates the json _source of the fragment - XContentBuilder jsonContent = getJsonFragmentBuilder(item, fragReader, parentId, contenttrackID, - fragNum--); + int embedIdx = 0; + for (FragHolder fh : frags) { + float[] emb = null; + if (embeddings != null && fh.content != null && !fh.content.trim().isEmpty()) { + if (embedIdx < embeddings.size()) { + emb = embeddings.get(embedIdx++); + } + } - // creates the request - IndexRequest contentRequest = createIndexRequest(contenttrackID, parentId, jsonContent); + XContentBuilder jsonContent = getJsonFragmentBuilder(item, fh.content, parentId, fh.contenttrackID, fh.fragNum, emb); + IndexRequest contentRequest = createIndexRequest(fh.contenttrackID, parentId, jsonContent); bulkRequest.add(contentRequest); - - idToPath.put(contenttrackID, item.getPath()); + idToPath.put(fh.contenttrackID, item.getPath()); LOGGER.debug("Added to bulk request {}", item.getPath()); if (bulkRequest.estimatedSizeInBytes() >= elasticConfig.getMin_bulk_size() || bulkRequest.numberOfActions() >= elasticConfig.getMin_bulk_items()) { - - // do not send more requests while commit is going on while (this.onCommit.get()) { this.wait(); } sendBulkRequest(); } - - } while (!Thread.currentThread().isInterrupted() && fragReader.nextFragment()); + } } finally { fragReader.close(); @@ -671,14 +774,69 @@ private XContentBuilder getJsonFragmentBuilder(IItem item, Reader textReader, St document_content.put("name", "content"); document_content.put("parent", parentID); + String content = getStringFromReader(textReader); + builder.startObject().field(BasicProps.EVIDENCE_UUID, item.getDataSource().getUUID()) .field(BasicProps.ID, item.getId()).field("document_content", document_content) .field("contenttrackID", contenttrackID).field(IndexTask.FRAG_NUM, fragNum) - .field(BasicProps.CONTENT, getStringFromReader(textReader)); + .field(BasicProps.CONTENT, content); + + RAGService ragService = RAGService.getInstance(); + if (ragService != null && ragService.getConfig().isEnabled()) { + try { + if (content != null && !content.trim().isEmpty()) { + java.util.HashSet categories = item.getCategorySet(); + String mimeType = item.getMediaType().toString(); + if (ragService.getConfig().shouldEmbed(categories, mimeType)) { + float[] embedding = ragService.getEmbedding(content); + if (embedding != null) { + builder.array("content_embedding", embedding); + builder.field("embedding_model", ragService.getConfig().getEmbeddingModel()); + builder.field("embedding_timestamp", System.currentTimeMillis()); + } + } + } + } catch (Exception e) { + LOGGER.error("Failed to generate embedding for fragment " + contenttrackID + " of item " + item.getId(), e); + } + } return builder.endObject(); } + private XContentBuilder getJsonFragmentBuilder(IItem item, String content, String parentID, + String contenttrackID, int fragNum, float[] embedding) throws IOException { + + XContentBuilder builder = XContentFactory.jsonBuilder(); + + HashMap document_content = new HashMap<>(); + document_content.put("name", "content"); + document_content.put("parent", parentID); + + builder.startObject().field(BasicProps.EVIDENCE_UUID, item.getDataSource().getUUID()) + .field(BasicProps.ID, item.getId()).field("document_content", document_content) + .field("contenttrackID", contenttrackID).field(IndexTask.FRAG_NUM, fragNum) + .field(BasicProps.CONTENT, content); + + if (embedding != null) { + RAGService ragService = RAGService.getInstance(); + builder.array("content_embedding", embedding); + builder.field(IndexTask.HAS_EMBEDDING_FRAG, "1"); + if (ragService != null && ragService.getConfig() != null) { + builder.field("embedding_model", ragService.getConfig().getEmbeddingModel()); + } + builder.field("embedding_timestamp", System.currentTimeMillis()); + } + + return builder.endObject(); + } + + private static class FragHolder { + String contenttrackID; + int fragNum; + String content; + } + private String getStringFromReader(Reader reader) throws IOException { StringBuilder sb = new StringBuilder(); int i = 0; diff --git a/iped-engine/src/main/java/iped/engine/task/index/IndexTask.java b/iped-engine/src/main/java/iped/engine/task/index/IndexTask.java index 8c3792bb76..f9524156e4 100644 --- a/iped-engine/src/main/java/iped/engine/task/index/IndexTask.java +++ b/iped-engine/src/main/java/iped/engine/task/index/IndexTask.java @@ -9,11 +9,14 @@ import java.util.List; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.FieldType; import org.apache.lucene.document.IntPoint; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexOptions; import org.apache.tika.io.TikaInputStream; import org.apache.tika.metadata.Metadata; @@ -26,19 +29,23 @@ import iped.engine.CmdLineArgs; import iped.engine.config.ConfigurationManager; import iped.engine.config.IndexTaskConfig; +import iped.engine.config.RAGConfig; import iped.engine.core.Worker.STATE; import iped.engine.data.IPEDSource; import iped.engine.data.Item; import iped.engine.io.CloseFilterReader; import iped.engine.io.FragmentingReader; import iped.engine.io.ParsingReader; +import iped.engine.rag.RAGService; import iped.engine.task.AbstractTask; import iped.engine.task.ParsingTask; import iped.engine.task.SkipCommitedTask; import iped.engine.task.carver.BaseCarveTask; import iped.engine.util.Util; +import iped.engine.util.UIPropertyListenerProvider; import iped.exception.IPEDException; import iped.parsers.standard.StandardParser; +import iped.properties.ExtraProperties; import iped.utils.IOUtil; /** @@ -60,12 +67,30 @@ public class IndexTask extends AbstractTask { public static final String FRAG_PARENT_ID = "fragParentId"; public static final String extraAttrFilename = "extraAttributes.dat"; //$NON-NLS-1$ + /** Stored field: unique track-ID of a fragment document (used to locate it). */ + public static final String FRAG_TRACK_ID = "fragTrackId"; + + /** KNN vector field: stores the text embedding of a fragment. */ + public static final String CONTENT_EMBEDDING = "content_embedding"; + + /** Stored field: stores the original fragment text so RAGService can retrieve it during search. */ + public static final String CONTENT_STORED = "content_stored"; + + /** Marker field: set to "1" on fragment child documents that have an embedding vector. */ + public static final String HAS_EMBEDDING_FRAG = "hasEmbeddingFrag"; + private static final AtomicBoolean finished = new AtomicBoolean(); private static final AtomicBoolean lastIDLoaded = new AtomicBoolean(); + /** Global RAG statistics counters for final Statistics summary report. */ + public static final java.util.concurrent.atomic.AtomicInteger totalEmbeddedDocs = + new java.util.concurrent.atomic.AtomicInteger(0); + public static final java.util.concurrent.atomic.AtomicInteger totalEmbeddedVectors = + new java.util.concurrent.atomic.AtomicInteger(0); + private static FieldType contentField; - private static final FieldType getContentFieldType() { + public static final FieldType getContentFieldType() { if (contentField == null) { FieldType field = new FieldType(); field.setIndexOptions(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS); @@ -119,6 +144,11 @@ public void process(IItem evidence) throws IOException { stats.updateLastId(evidence.getId()); + RAGConfig ragConfig = ConfigurationManager.get().findObject(RAGConfig.class); + boolean luceneVectorMode = ragConfig != null + && ragConfig.isEnabled() + && "lucene".equalsIgnoreCase(ragConfig.getVectorStoreMode()); + if (textReader == null) { if (indexConfig.isIndexFileContents() && (indexConfig.isIndexUnallocated() || !BaseCarveTask.UNALLOCATED_MIMETYPE.equals(evidence.getMediaType()))) { @@ -143,10 +173,18 @@ public void process(IItem evidence) throws IOException { if (textReader == null) textReader = new StringReader(""); //$NON-NLS-1$ - FragmentingReader fragReader = new FragmentingReader(textReader, indexConfig.getTextSplitSize(), - indexConfig.getTextOverlapSize()); + int splitSize = indexConfig.getTextSplitSize(); + int overlapSize = indexConfig.getTextOverlapSize(); + if (ragConfig != null && ragConfig.isEnabled()) { + String mimeType = evidence.getMediaType() != null ? evidence.getMediaType().toString() : ""; + if (ragConfig.shouldEmbed(evidence.getCategorySet(), mimeType)) { + splitSize = ragConfig.getChunkSize(); + overlapSize = ragConfig.getChunkOverlap(); + } + } + FragmentingReader fragReader = new FragmentingReader(textReader, splitSize, overlapSize); try { - worker.writer.addDocuments(new DocumentsIterable(evidence, fragReader)); + worker.writer.addDocuments(new DocumentsIterable(evidence, fragReader, luceneVectorMode, ragConfig)); } catch (IOException e) { if (IOUtil.isDiskFull(e)) @@ -163,61 +201,215 @@ public void process(IItem evidence) throws IOException { private class DocumentsIterable implements Iterable { private IItem item; - private FragmentingReader fragReader; - private boolean hasMoreContentFrags, parentIndexed = false; + private java.util.List fragmentTexts; + private java.util.List vectors; + private boolean parentIndexed = false; private int numFrags = 0; - - private DocumentsIterable(IItem item, FragmentingReader fragReader) { + private final boolean luceneVectorMode; + private final boolean shouldEmbed; + private final int embDimension; + private final String globalId; + private FragmentingReader lazyReader; // for non-vector mode + + private DocumentsIterable(IItem item, FragmentingReader fragReader, + boolean luceneVectorMode, RAGConfig ragConfig) { this.item = item; - this.fragReader = fragReader; + this.luceneVectorMode = luceneVectorMode; + + // Pre-compute embedding eligibility for this item + boolean embed = false; + int dim = 0; + String gid = null; + if (luceneVectorMode && ragConfig != null) { + String mimeType = item.getMediaType() != null ? item.getMediaType().toString() : ""; + embed = ragConfig.shouldEmbed(item.getCategorySet(), mimeType); + dim = ragConfig.getEmbeddingDimensions(); + gid = (String) item.getExtraAttribute(ExtraProperties.GLOBAL_ID); + } + this.shouldEmbed = embed; + this.embDimension = dim; + this.globalId = gid; + + if (luceneVectorMode) { + // Read all fragments eagerly + java.util.List frags = new java.util.ArrayList<>(); + try { + frags.add(captureText(fragReader)); + while (fragReader.nextFragment()) { + frags.add(captureText(fragReader)); + } + } catch (IOException e) { + throw new RuntimeException(e); + } + this.fragmentTexts = frags; + + // Batch generate embeddings synchronously + RAGService ragService = RAGService.getInstance(); + if (shouldEmbed && !frags.isEmpty() && ragService != null && ragService.isAvailable()) { + try { + this.vectors = ragService.getEmbeddings(frags); + if (this.vectors != null && !this.vectors.isEmpty()) { + totalEmbeddedDocs.incrementAndGet(); + totalEmbeddedVectors.addAndGet(this.vectors.size()); + } + LOGGER.debug("RAG: generated {} embeddings for {}", this.vectors != null ? this.vectors.size() : 0, item.getPath()); + } catch (Exception e) { + LOGGER.error("RAG: Error generating embeddings for {}", item.getPath(), e); + } + } + } else { + this.lazyReader = fragReader; + } } - public Iterator iterator() { - return new Iterator() { + public java.util.Iterator iterator() { + return new java.util.Iterator() { + private int index = 0; + private boolean hasMoreContentFrags = false; public boolean hasNext() { try { while (worker.state != STATE.RUNNING) { + synchronized (worker) { + if (worker.state == STATE.PAUSING) { + worker.state = STATE.PAUSED; + } + } Thread.sleep(1000); } if (Thread.interrupted()) { throw new InterruptedException(); } - hasMoreContentFrags = (numFrags == 0 || fragReader.nextFragment()); - return hasMoreContentFrags || !parentIndexed; - + if (luceneVectorMode) { + return index < fragmentTexts.size() || !parentIndexed; + } else { + hasMoreContentFrags = (numFrags == 0 || lazyReader.nextFragment()); + return hasMoreContentFrags || !parentIndexed; + } } catch (InterruptedException | IOException e) { throw new RuntimeException(e); } } public Document next() { - if (hasMoreContentFrags) { - if (++numFrags > 1) { - stats.incSplits(); - LOGGER.info("{} Splitting text of {}", Thread.currentThread().getName(), item.getPath()); //$NON-NLS-1$ + if (luceneVectorMode) { + if (index < fragmentTexts.size()) { + String fragText = fragmentTexts.get(index); + numFrags = index + 1; + if (numFrags > 1) { + stats.incSplits(); + if (shouldEmbed) { + LOGGER.debug("{} Splitting RAG text of {}", Thread.currentThread().getName(), item.getPath()); //$NON-NLS-1$ + } else { + LOGGER.info("{} Splitting text of {}", Thread.currentThread().getName(), item.getPath()); //$NON-NLS-1$ + } + } + + // child (content) document + Document doc = new Document(); + doc.add(new IntPoint(FRAG_NUM, numFrags)); + doc.add(new IntPoint(FRAG_PARENT_ID, item.getId())); + + if (shouldEmbed && globalId != null) { + String fragTrackId = globalId + "#" + numFrags; + + // StoredField so the text can be retrieved at search time + doc.add(new StoredField(CONTENT_STORED, fragText)); + + // StringField (stored) so it can be uniquely identified + doc.add(new StringField(FRAG_TRACK_ID, fragTrackId, Field.Store.YES)); + + // StringField (stored) containing the numeric parent ID + doc.add(new StringField("fragParentNumericId", String.valueOf(item.getId()), Field.Store.YES)); + + // Add pre-generated vector if available + if (vectors != null && index < vectors.size()) { + float[] vector = vectors.get(index); + if (vector != null && vector.length > 0) { + doc.add(new org.apache.lucene.document.KnnVectorField( + CONTENT_EMBEDDING, vector, + org.apache.lucene.index.VectorSimilarityFunction.COSINE)); + doc.add(new StringField(HAS_EMBEDDING_FRAG, "1", Field.Store.YES)); + } + } + } + + // Index the text for BM25 lexical search + doc.add(new Field(IndexItem.CONTENT, new StringReader(fragText), + getContentFieldType())); + + index++; + return doc; + } else { + if (numFrags > 1) { + item.setExtraAttribute(TEXT_SPLITTED, Boolean.TRUE.toString()); + } + // Mark the parent item as having embeddings so it can be + // filtered in the IPED UI (follows the textSplitted pattern). + // Only set if at least one non-null vector was actually generated + // (mirrors the same guard used when adding KnnVectorField to children). + if (shouldEmbed && vectors != null + && vectors.stream().anyMatch(v -> v != null && v.length > 0)) { + item.setExtraAttribute(ExtraProperties.HAS_EMBEDDING, Boolean.TRUE.toString()); + } + long totalSize = 0; + for (String s : fragmentTexts) { + totalSize += s.length(); + } + item.setExtraAttribute(TEXT_SIZE, totalSize); + // parent (metadata) document + Document doc = IndexItem.Document(item, output); + parentIndexed = true; + return doc; } - // child (content) document - Document doc = new Document(); - doc.add(new IntPoint(FRAG_NUM, numFrags)); - doc.add(new IntPoint(FRAG_PARENT_ID, item.getId())); - doc.add(new Field(IndexItem.CONTENT, new CloseFilterReader(fragReader), getContentFieldType())); - return doc; } else { - if (numFrags > 1) { - item.setExtraAttribute(TEXT_SPLITTED, Boolean.TRUE.toString()); + // Standard lazy mode + if (hasMoreContentFrags) { + if (++numFrags > 1) { + stats.incSplits(); + LOGGER.info("{} Splitting text of {}", Thread.currentThread().getName(), item.getPath()); //$NON-NLS-1$ + } + Document doc = new Document(); + doc.add(new IntPoint(FRAG_NUM, numFrags)); + doc.add(new IntPoint(FRAG_PARENT_ID, item.getId())); + doc.add(new Field(IndexItem.CONTENT, new CloseFilterReader(lazyReader), + getContentFieldType())); + return doc; + } else { + if (numFrags > 1) { + item.setExtraAttribute(TEXT_SPLITTED, Boolean.TRUE.toString()); + } + item.setExtraAttribute(TEXT_SIZE, lazyReader.getTotalTextSize()); + // parent (metadata) document + Document doc = IndexItem.Document(item, output); + parentIndexed = true; + return doc; } - item.setExtraAttribute(TEXT_SIZE, fragReader.getTotalTextSize()); - // parent (metadata) document - Document doc = IndexItem.Document(item, output); - parentIndexed = true; - return doc; } } - + }; } + /** + * Reads the current fragment from the FragmentingReader into a String. + * The String is then used both as a StoredField and as the text for the + * BM25 Field (via a StringReader), avoiding any double-read of the stream. + */ + private String captureText(FragmentingReader reader) { + StringBuilder sb = new StringBuilder(4096); + char[] buf = new char[4096]; + int n; + try { + while ((n = reader.read(buf, 0, buf.length)) != -1) { + sb.append(buf, 0, n); + } + } catch (IOException e) { + LOGGER.warn("Error reading fragment text for embedding.", e); + } + return sb.toString(); + } + } private Metadata getMetadata(IItem evidence) { @@ -260,6 +452,15 @@ public void init(ConfigurationManager configurationManager) throws Exception { this.autoParser = new StandardParser(); + RAGConfig ragConfig = configurationManager.findObject(RAGConfig.class); + if (ragConfig != null && ragConfig.isEnabled()) { + try { + RAGService.initialize(output, ragConfig); + } catch (Exception e) { + LOGGER.error("RAG: Failed to initialise RAGService", e); + } + } + } @Override diff --git a/iped-engine/src/test/java/iped/engine/rag/RAGServiceTest.java b/iped-engine/src/test/java/iped/engine/rag/RAGServiceTest.java new file mode 100644 index 0000000000..f837641421 --- /dev/null +++ b/iped-engine/src/test/java/iped/engine/rag/RAGServiceTest.java @@ -0,0 +1,401 @@ +package iped.engine.rag; + +import static org.junit.Assert.*; + +import java.io.File; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; + +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import iped.engine.config.RAGConfig; +import iped.engine.rag.RAGService.RAGSourceDoc; + +/** + * Unit tests for {@link RAGService}. + * + * Network-dependent methods (OpenSearch calls) are not exercised here ??? they + * require an integration-test environment. Tests focus on: + *
    + *
  • SQLite embedding cache (write ??? read round-trip).
    • + *
  • Float ??? byte array conversion helpers (via cache round-trip).
    • + *
  • Reciprocal Rank Fusion (RRF) ranking logic in + * {@link RAGService#performHybridSearch(String, float[])} via a + * subclass that stubs the OpenSearch calls.
    • + *
  • Parent-ID deduplication in {@link RAGService#findSimilarDocumentIds}.
    • + *
  • Provider selection: correct provider type created per config.
    • + *
+ */ +public class RAGServiceTest { + + @Rule + public TemporaryFolder tempFolder = new TemporaryFolder(); + + private RAGConfig config; + + @Before + public void setUp() { + config = new RAGConfig(); + setField(config, "embeddingProvider", "local"); + setField(config, "embeddingModel", "bge-m3"); + setField(config, "embeddingEndpoint", "http://localhost:11434/api/embeddings"); + setField(config, "embeddingDimensions", 4); + setField(config, "llmProvider", "local"); + setField(config, "llmModel", "qwen3"); + setField(config, "llmEndpoint", "http://localhost:11434/v1"); + setField(config, "llmApiKey", ""); + setField(config, "chunkSimilarityThreshold", 0.0f); + setField(config, "maxRetrievedChunks", 10); + } + + // ---------------------------------------------------------------- cache tests + + @Test + public void testEmbeddingCacheRoundTrip() throws Exception { + File caseDir = tempFolder.newFolder("case1"); + RAGService.initialize(caseDir, config); + RAGService service = RAGService.getInstance(); + + float[] original = {0.1f, 0.2f, 0.3f, 0.4f}; + service.cacheEmbedding("hello world", "bge-m3", original); + + float[] retrieved = service.getCachedEmbedding("hello world", "bge-m3"); + assertNotNull("Cached embedding must not be null", retrieved); + assertEquals("Embedding length mismatch", original.length, retrieved.length); + for (int i = 0; i < original.length; i++) { + assertEquals("Value mismatch at " + i, original[i], retrieved[i], 1e-6f); + } + } + + @Test + public void testCacheMissReturnsNull() throws Exception { + File caseDir = tempFolder.newFolder("case2"); + RAGService.initialize(caseDir, config); + RAGService service = RAGService.getInstance(); + + assertNull("Cache miss should return null", + service.getCachedEmbedding("not-in-cache", "bge-m3")); + } + + @Test + public void testCacheDifferentModelsAreIsolated() throws Exception { + File caseDir = tempFolder.newFolder("case3"); + RAGService.initialize(caseDir, config); + RAGService service = RAGService.getInstance(); + + float[] emb1 = {1.0f, 0.0f, 0.0f, 0.0f}; + float[] emb2 = {0.0f, 1.0f, 0.0f, 0.0f}; + service.cacheEmbedding("text", "model-A", emb1); + service.cacheEmbedding("text", "model-B", emb2); + + float[] r1 = service.getCachedEmbedding("text", "model-A"); + float[] r2 = service.getCachedEmbedding("text", "model-B"); + assertNotNull(r1); + assertNotNull(r2); + assertEquals(1.0f, r1[0], 1e-6f); + assertEquals(1.0f, r2[1], 1e-6f); + } + + @Test + public void testCacheOverwriteSameKey() throws Exception { + File caseDir = tempFolder.newFolder("case4"); + RAGService.initialize(caseDir, config); + RAGService service = RAGService.getInstance(); + + float[] first = {1.0f, 0.0f, 0.0f, 0.0f}; + float[] second = {0.0f, 0.0f, 0.0f, 1.0f}; + service.cacheEmbedding("same-text", "model-X", first); + service.cacheEmbedding("same-text", "model-X", second); + + float[] result = service.getCachedEmbedding("same-text", "model-X"); + assertNotNull(result); + assertEquals("Latest value should win", 1.0f, result[3], 1e-6f); + } + + // ---------------------------------------------------------------- RRF tests + // + // We use a local subclass that stubs performVectorSearch / performLexicalSearch + // so we can exercise the RRF fusion logic without an OpenSearch cluster. + + @Test + public void testHybridSearchRRFCombinesRanks() throws Exception { + File caseDir = tempFolder.newFolder("case-rrf"); + RAGService.initialize(caseDir, config); + + RAGSourceDoc a = doc("1", "content-a", "a-001", 0.9f); + RAGSourceDoc b = doc("2", "content-b", "b-001", 0.8f); + RAGSourceDoc c = doc("3", "content-c", "c-001", 0.7f); + + // doc-A appears in both vector and lexical lists ??? higher RRF score + RAGService service = new RAGService(caseDir, config) { + @Override + public List performVectorSearch(float[] v, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(a, b); + } + @Override + public List performLexicalSearch(String q, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(a, c); + } + }; + + List result = service.performHybridSearch("question", new float[4]); + assertFalse("Result must not be empty", result.isEmpty()); + assertEquals("Doc-A (in both lists) should be ranked first", "a-001", + result.get(0).contenttrackID); + } + + @Test + public void testHybridSearchRespectMaxChunks() throws Exception { + File caseDir = tempFolder.newFolder("case-max"); + RAGConfig cfg = cloneConfig(config); + setField(cfg, "maxRetrievedChunks", 2); + RAGService.initialize(caseDir, cfg); + + RAGSourceDoc a = doc("1", "ca", "t1", 0.9f); + RAGSourceDoc b = doc("2", "cb", "t2", 0.8f); + RAGSourceDoc c = doc("3", "cc", "t3", 0.7f); + + RAGService service = new RAGService(caseDir, cfg) { + @Override + public List performVectorSearch(float[] v, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(a, b, c); + } + @Override + public List performLexicalSearch(String q, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(a, b, c); + } + }; + + List result = service.performHybridSearch("q", new float[4]); + assertTrue("Result should not exceed maxRetrievedChunks=2", result.size() <= 2); + } + + /** + * Verifies that chunkSimilarityThreshold is applied on the normalized RRF score + * AFTER fusion, not on raw vector scores. + * + * Setup: + * vector: [highScore(rank1), lowScore(rank2)] + * lexical: [highScore(rank1)] ??? only highScore has lexical support + * + * Expected RRF normalized scores: + * highScore = 1.0 (top score in both lists ??? normalized to 1.0) + * lowScore ??? 0.50 (only vector rank2 contribution, no lexical ??? ~50% of top) + * + * With threshold = 0.85: lowScore (???0.50) should be filtered out. + * With threshold = 0.85: highScore (1.0) should be kept. + */ + @Test + public void testHybridSearchSimilarityThresholdFiltersVectorHits() throws Exception { + File caseDir = tempFolder.newFolder("case-thresh"); + RAGConfig cfg = cloneConfig(config); + setField(cfg, "chunkSimilarityThreshold", 0.85f); + setField(cfg, "maxRetrievedChunks", 10); + RAGService.initialize(caseDir, cfg); + + // highScore appears in both lists ??? RRF normalized score = 1.0 + // lowScore appears only in vector list at rank 2 ??? RRF normalized score ??? 0.50 + RAGSourceDoc highScore = doc("1", "text", "t1", 0.9f); + RAGSourceDoc lowScore = doc("2", "text", "t2", 0.5f); + + RAGService service = new RAGService(caseDir, cfg) { + @Override + public List performVectorSearch(float[] v, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(highScore, lowScore); + } + @Override + public List performLexicalSearch(String q, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(highScore); // only highScore has lexical support + } + }; + + List result = service.performHybridSearch("q", new float[4]); + + // highScore must be present (normalized RRF = 1.0 >= threshold 0.85) + boolean foundHigh = result.stream().anyMatch(d -> "t1".equals(d.contenttrackID)); + assertTrue("High-relevance fragment must be kept", foundHigh); + + // lowScore must be filtered (normalized RRF ??? 0.50 < threshold 0.85) + boolean foundLow = result.stream().anyMatch(d -> "t2".equals(d.contenttrackID)); + assertFalse("Low combined-relevance fragment must be filtered by post-RRF threshold", foundLow); + } + + /** + * Verifies the key benefit of the new design: a fragment that is weak on the + * vector side (raw score below old threshold) but strong on the lexical side + * is NOT lost ??? the RRF fusion rescues it. + * + * Setup: + * vector: [docA(rank1), docB(rank2)] ??? docB has raw vector score < old threshold + * lexical: [docB(rank1), docA(rank2)] ??? docB is the best lexical hit + * + * With threshold = 0.0 (disabled), both docs pass. + * The important assertion is that docB is present despite a weak vector score, + * because the lexical signal compensated it via RRF. + */ + @Test + public void testHybridSearchLexicalRescuesWeakVectorHit() throws Exception { + File caseDir = tempFolder.newFolder("case-rescue"); + RAGConfig cfg = cloneConfig(config); + setField(cfg, "chunkSimilarityThreshold", 0.0f); // no post-filter + setField(cfg, "maxRetrievedChunks", 10); + RAGService.initialize(caseDir, cfg); + + RAGSourceDoc docA = doc("1", "textA", "tA", 0.9f); + RAGSourceDoc docB = doc("2", "textB", "tB", 0.55f); // weak vector, strong lexical + + RAGService service = new RAGService(caseDir, cfg) { + @Override + public List performVectorSearch(float[] v, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(docA, docB); + } + @Override + public List performLexicalSearch(String q, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(docB, docA); // docB is best lexical hit + } + }; + + List result = service.performHybridSearch("q", new float[4]); + + boolean foundB = result.stream().anyMatch(d -> "tB".equals(d.contenttrackID)); + assertTrue("docB must be kept: weak vector score rescued by strong lexical rank", foundB); + } + + // ---------------------------------------------------------------- deduplication + + @Test + public void testFindSimilarDocumentIdsDeduplicatesParents() throws Exception { + File caseDir = tempFolder.newFolder("case-dedup"); + RAGService.initialize(caseDir, config); + + RAGSourceDoc frag1 = doc("42", "text1", "42-frag1", 0.9f); + RAGSourceDoc frag2 = doc("42", "text2", "42-frag2", 0.8f); // same itemId + RAGSourceDoc frag3 = doc("99", "text3", "99-frag1", 0.7f); + + RAGService service = new RAGService(caseDir, config) { + @Override + public List performVectorSearch(float[] v, int k, java.util.Collection allowedItemIds) { + return Arrays.asList(frag1, frag2, frag3); + } + }; + + List ids = service.findSimilarDocumentIds(new float[4], 10); + assertTrue("Item 42 must be in results", ids.contains("42")); + assertTrue("Item 99 must be in results", ids.contains("99")); + assertEquals("Duplicate parent must be deduped ??? 2 unique parents", 2, ids.size()); + } + + // ---------------------------------------------------------------- provider wiring + + @Test + public void testLocalEmbeddingProviderSelectedByConfig() throws Exception { + File caseDir = tempFolder.newFolder("prov-local-emb"); + setField(config, "embeddingProvider", "local"); + RAGService.initialize(caseDir, config); + assertTrue("Expected LocalEmbeddingProvider", + RAGService.getInstance().getEmbeddingProvider() instanceof LocalEmbeddingProvider); + } + + @Test + public void testRemoteEmbeddingProviderSelectedByConfig() throws Exception { + File caseDir = tempFolder.newFolder("prov-remote-emb"); + setField(config, "embeddingProvider", "remote"); + RAGService.initialize(caseDir, config); + assertTrue("Expected RemoteEmbeddingProvider", + RAGService.getInstance().getEmbeddingProvider() instanceof RemoteEmbeddingProvider); + } + + @Test + public void testLocalLLMProviderSelectedByConfig() throws Exception { + File caseDir = tempFolder.newFolder("prov-local-llm"); + setField(config, "llmProvider", "local"); + RAGService.initialize(caseDir, config); + assertTrue("Expected LocalLLMProvider", + RAGService.getInstance().getLlmProvider() instanceof LocalLLMProvider); + } + + @Test + public void testGeminiLLMProviderSelectedByConfig() throws Exception { + File caseDir = tempFolder.newFolder("prov-gemini"); + setField(config, "llmProvider", "gemini"); + RAGService.initialize(caseDir, config); + assertTrue("Expected GeminiLLMProvider", + RAGService.getInstance().getLlmProvider() instanceof GeminiLLMProvider); + } + + @Test + public void testClaudeLLMProviderSelectedByConfig() throws Exception { + File caseDir = tempFolder.newFolder("prov-claude"); + setField(config, "llmProvider", "claude"); + RAGService.initialize(caseDir, config); + assertTrue("Expected ClaudeLLMProvider", + RAGService.getInstance().getLlmProvider() instanceof ClaudeLLMProvider); + } + + @Test + public void testShouldEmbedFilterModes() { + RAGConfig cfg = new RAGConfig(); + + // 1. Test default Blacklist mode + setField(cfg, "embeddingFilterMode", "blacklist"); + setField(cfg, "embeddingCategoryBlacklist", "^(Programs and Libraries|Image Disks)$"); + setField(cfg, "embeddingMimeTypeBlacklist", "^(application/octet-stream)$"); + + // Allowed: Documents, PDF + assertTrue(cfg.shouldEmbed(Arrays.asList("Documents"), "application/pdf")); + // Blacklisted category: Programs and Libraries + assertFalse(cfg.shouldEmbed(Arrays.asList("Programs and Libraries"), "application/pdf")); + // Blacklisted mime: octet-stream + assertFalse(cfg.shouldEmbed(Arrays.asList("Documents"), "application/octet-stream")); + + // 2. Test Whitelist mode + setField(cfg, "embeddingFilterMode", "whitelist"); + setField(cfg, "embeddingCategoryWhitelist", "^(Documents|Chats)$"); + setField(cfg, "embeddingMimeTypeWhitelist", "^(text/plain)$"); + + // Whitelisted category: Documents + assertTrue(cfg.shouldEmbed(Arrays.asList("Documents"), "application/pdf")); + // Whitelisted category: Chats + assertTrue(cfg.shouldEmbed(Arrays.asList("Chats"), "application/octet-stream")); + // Whitelisted mime: text/plain + assertTrue(cfg.shouldEmbed(Arrays.asList("NonWhitelistedCategory"), "text/plain")); + // Neither whitelisted: NonWhitelistedCategory + application/pdf + assertFalse(cfg.shouldEmbed(Arrays.asList("NonWhitelistedCategory"), "application/pdf")); + } + + // ---------------------------------------------------------------- helpers + + private static RAGSourceDoc doc(String itemId, String content, String trackId, float score) { + RAGSourceDoc d = new RAGSourceDoc(); + d.itemId = itemId; + d.content = content; + d.contenttrackID = trackId; + d.score = score; + return d; + } + + /** Shallow-clone a RAGConfig by copying fields reflectively. */ + private static RAGConfig cloneConfig(RAGConfig src) { + RAGConfig dst = new RAGConfig(); + for (java.lang.reflect.Field f : src.getClass().getDeclaredFields()) { + if (java.lang.reflect.Modifier.isStatic(f.getModifiers())) continue; + f.setAccessible(true); + try { f.set(dst, f.get(src)); } catch (IllegalAccessException ignore) {} + } + return dst; + } + + private static void setField(Object target, String fieldName, Object value) { + try { + java.lang.reflect.Field f = target.getClass().getDeclaredField(fieldName); + f.setAccessible(true); + f.set(target, value); + } catch (NoSuchFieldException | IllegalAccessException e) { + throw new RuntimeException("Could not set field " + fieldName, e); + } + } +}