diff --git a/libs/kotaemon/kotaemon/indices/ingests/files.py b/libs/kotaemon/kotaemon/indices/ingests/files.py index b40dfbba..20ab5eec 100644 --- a/libs/kotaemon/kotaemon/indices/ingests/files.py +++ b/libs/kotaemon/kotaemon/indices/ingests/files.py @@ -4,7 +4,6 @@ from typing import Type from decouple import config from llama_index.core.readers.base import BaseReader from llama_index.readers.file import PDFReader -from theflow.settings import settings as flowsettings from kotaemon.base import BaseComponent, Document, Param from kotaemon.indices.extractors import BaseDocParser @@ -33,12 +32,8 @@ adobe_reader = AdobeReader() azure_reader = AzureAIDocumentIntelligenceLoader( endpoint=str(config("AZURE_DI_ENDPOINT", default="")), credential=str(config("AZURE_DI_CREDENTIAL", default="")), - cache_dir=getattr(flowsettings, "KH_MARKDOWN_OUTPUT_DIR", None), ) docling_reader = DoclingReader() -adobe_reader.vlm_endpoint = ( - azure_reader.vlm_endpoint -) = docling_reader.vlm_endpoint = getattr(flowsettings, "KH_VLM_ENDPOINT", "") paddle_device = str(config("PADDLE_DEVICE", default="gpu")) paddle_struct_reader = PPStructureV3Reader(device=paddle_device) diff --git a/libs/kotaemon/kotaemon/indices/qa/citation_qa.py b/libs/kotaemon/kotaemon/indices/qa/citation_qa.py index b5656328..1b014910 100644 --- a/libs/kotaemon/kotaemon/indices/qa/citation_qa.py +++ b/libs/kotaemon/kotaemon/indices/qa/citation_qa.py @@ -4,7 +4,6 @@ from typing import Generator import numpy as np from decouple import config -from theflow.settings import settings as flowsettings from kotaemon.base import ( AIMessage, @@ -24,13 +23,6 @@ from .format_context import ( ) from .utils import find_text -try: - from ktem.llms.manager import llms - from ktem.reasoning.prompt_optimization.mindmap import CreateMindmapPipeline - from ktem.utils.render import Render -except ImportError: - raise ImportError("Please install `ktem` to use this component") - MAX_IMAGES = 10 CITATION_TIMEOUT = 5.0 CONTEXT_RELEVANT_WARNING_SCORE = config( @@ -95,15 +87,11 @@ class AnswerWithContextPipeline(BaseComponent): lang: the language of the answer. Currently support English and Japanese """ - llm: ChatLLM = Node(default_callback=lambda _: llms.get_default()) - vlm_endpoint: str = getattr(flowsettings, "KH_VLM_ENDPOINT", "") - use_multimodal: bool = getattr(flowsettings, "KH_REASONINGS_USE_MULTIMODAL", True) - citation_pipeline: CitationPipeline = Node( - default_callback=lambda _: CitationPipeline(llm=llms.get_default()) - ) - create_mindmap_pipeline: CreateMindmapPipeline = Node( - default_callback=lambda _: CreateMindmapPipeline(llm=llms.get_default()) - ) + llm: ChatLLM = Node() + vlm_endpoint: str = "" + use_multimodal: bool = True + citation_pipeline: CitationPipeline = Node() + create_mindmap_pipeline: BaseComponent | None = Node(default=None) qa_template: str = DEFAULT_QA_TEXT_PROMPT qa_table_template: str = DEFAULT_QA_TABLE_PROMPT @@ -224,7 +212,7 @@ class AnswerWithContextPipeline(BaseComponent): citation_thread = threading.Thread(target=citation_call) citation_thread.start() - if self.enable_mindmap: + if self.enable_mindmap and self.create_mindmap_pipeline is not None: mindmap_thread = threading.Thread(target=mindmap_call) mindmap_thread.start() @@ -319,85 +307,3 @@ class AnswerWithContextPipeline(BaseComponent): # print("Matched citation:", quote, matched_excerpts), return spans - def prepare_citations(self, answer, docs) -> tuple[list[Document], list[Document]]: - """Prepare the citations to show on the UI""" - with_citation, without_citation = [], [] - has_llm_score = any("llm_trulens_score" in doc.metadata for doc in docs) - - spans = self.match_evidence_with_context(answer, docs) - id2docs = {doc.doc_id: doc for doc in docs} - not_detected = set(id2docs.keys()) - set(spans.keys()) - - # render highlight spans - for _id, ss in spans.items(): - if not ss: - not_detected.add(_id) - continue - cur_doc = id2docs[_id] - highlight_text = "" - - ss = sorted(ss, key=lambda x: x["start"]) - last_end = 0 - text = cur_doc.text[: ss[0]["start"]] - - for idx, span in enumerate(ss): - # prevent overlapping between span - span_start = max(last_end, span["start"]) - span_end = max(last_end, span["end"]) - - to_highlight = cur_doc.text[span_start:span_end] - last_end = span_end - - # append to highlight on PDF viewer - highlight_text += (" " if highlight_text else "") + to_highlight - - span_idx = span.get("idx", None) - if span_idx is not None: - to_highlight = f"【{span_idx}】" + to_highlight - - text += Render.highlight( - to_highlight, - elem_id=str(span_idx) if span_idx is not None else None, - ) - if idx < len(ss) - 1: - text += cur_doc.text[span["end"] : ss[idx + 1]["start"]] - - text += cur_doc.text[ss[-1]["end"] :] - # add to display list - with_citation.append( - Document( - channel="info", - content=Render.collapsible_with_header_score( - cur_doc, - override_text=text, - highlight_text=highlight_text, - open_collapsible=True, - ), - ) - ) - - print("Got {} cited docs".format(len(with_citation))) - - sorted_not_detected_items_with_scores = [ - (id_, id2docs[id_].metadata.get("llm_trulens_score", 0.0)) - for id_ in not_detected - ] - sorted_not_detected_items_with_scores.sort(key=lambda x: x[1], reverse=True) - - for id_, _ in sorted_not_detected_items_with_scores: - doc = id2docs[id_] - doc_score = doc.metadata.get("llm_trulens_score", 0.0) - is_open = not has_llm_score or ( - doc_score - > CONTEXT_RELEVANT_WARNING_SCORE - # and len(with_citation) == 0 - ) - without_citation.append( - Document( - channel="info", - content=Render.collapsible_with_header_score( - doc, open_collapsible=is_open - ), - ) - ) - return with_citation, without_citation diff --git a/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py b/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py index aa470652..77ff8d30 100644 --- a/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py +++ b/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py @@ -220,7 +220,7 @@ class AnswerWithInlineCitation(AnswerWithContextPipeline): # execute function call in thread if evidence: - if self.enable_mindmap: + if self.enable_mindmap and self.create_mindmap_pipeline is not None: mindmap_thread = threading.Thread(target=mindmap_call) mindmap_thread.start() diff --git a/libs/kotaemon/kotaemon/indices/rankings/cohere.py b/libs/kotaemon/kotaemon/indices/rankings/cohere.py index 039a00db..70c18b37 100644 --- a/libs/kotaemon/kotaemon/indices/rankings/cohere.py +++ b/libs/kotaemon/kotaemon/indices/rankings/cohere.py @@ -10,7 +10,6 @@ from .base import BaseReranking class CohereReranking(BaseReranking): model_name: str = "rerank-v4.0-fast" cohere_api_key: str = config("COHERE_API_KEY", "") - use_key_from_ktem: bool = False def run(self, documents: list[Document], query: str) -> list[Document]: """Use Cohere Reranker model to re-order documents @@ -19,25 +18,10 @@ class CohereReranking(BaseReranking): import cohere except ImportError: raise ImportError( - "Please install Cohere `pip install cohere` to use Cohere Reranking" + "Please install Cohere " + "`pip install cohere` to use Cohere Reranking" ) - # try to get COHERE_API_KEY from embeddings - if not self.cohere_api_key and self.use_key_from_ktem: - try: - from ktem.embeddings.manager import ( - embedding_models_manager as embeddings, - ) - - cohere_model = embeddings.get("cohere") - ktem_cohere_api_key = cohere_model._kwargs.get( # type: ignore - "cohere_api_key" - ) - if ktem_cohere_api_key != "your-key": - self.cohere_api_key = ktem_cohere_api_key - except Exception as e: - print("Cannot get Cohere API key from `ktem`", e) - if not self.cohere_api_key: print("Cohere API key not found. Skipping rerankings.") return documents diff --git a/libs/kotaemon/kotaemon/indices/vectorindex.py b/libs/kotaemon/kotaemon/indices/vectorindex.py index d882d6bf..1db2072a 100644 --- a/libs/kotaemon/kotaemon/indices/vectorindex.py +++ b/libs/kotaemon/kotaemon/indices/vectorindex.py @@ -5,8 +5,6 @@ import uuid from pathlib import Path from typing import Optional, Sequence, cast -from theflow.settings import settings as flowsettings - from kotaemon.base import BaseComponent, Document, RetrievedDocument from kotaemon.embeddings import BaseEmbeddings from kotaemon.storages import BaseDocumentStore, BaseVectorStore @@ -27,7 +25,7 @@ class VectorIndexing(BaseIndexing): - List of texts """ - cache_dir: Optional[str] = getattr(flowsettings, "KH_CHUNKS_OUTPUT_DIR", None) + cache_dir: Optional[str] = None vector_store: BaseVectorStore doc_store: Optional[BaseDocumentStore] = None embedding: BaseEmbeddings diff --git a/libs/kotaemon/kotaemon/loaders/html_loader.py b/libs/kotaemon/kotaemon/loaders/html_loader.py index 555efe50..d2df1975 100644 --- a/libs/kotaemon/kotaemon/loaders/html_loader.py +++ b/libs/kotaemon/kotaemon/loaders/html_loader.py @@ -3,7 +3,6 @@ from pathlib import Path from typing import Optional from llama_index.core.readers.base import BaseReader -from theflow.settings import settings as flowsettings from kotaemon.base import Document @@ -79,9 +78,7 @@ class MhtmlReader(BaseReader): def __init__( self, - cache_dir: Optional[str] = getattr( - flowsettings, "KH_MARKDOWN_OUTPUT_DIR", None - ), + cache_dir: Optional[str] = None, open_encoding: Optional[str] = None, bs_kwargs: Optional[dict] = None, get_text_separator: str = "", diff --git a/libs/ktem/ktem/index/file/pipelines.py b/libs/ktem/ktem/index/file/pipelines.py index 7d05e4d1..f5320600 100644 --- a/libs/ktem/ktem/index/file/pipelines.py +++ b/libs/ktem/ktem/index/file/pipelines.py @@ -53,6 +53,19 @@ from .base import BaseFileIndexIndexing, BaseFileIndexRetriever logger = logging.getLogger(__name__) +# Wire app-level config into framework reader singletons. +# These were previously read by kotaemon directly from flowsettings; +# after decoupling, ktem (app layer) is responsible for setting them. +_vlm_endpoint = getattr(settings, "KH_VLM_ENDPOINT", "") +_markdown_output_dir = getattr(settings, "KH_MARKDOWN_OUTPUT_DIR", None) +adobe_reader.vlm_endpoint = _vlm_endpoint +azure_reader.vlm_endpoint = _vlm_endpoint +docling_reader.vlm_endpoint = _vlm_endpoint +azure_reader.cache_dir = _markdown_output_dir +_mhtml_reader = KH_DEFAULT_FILE_EXTRACTORS.get(".mhtml") +if _mhtml_reader is not None and hasattr(_mhtml_reader, "cache_dir"): + _mhtml_reader.cache_dir = _markdown_output_dir + @lru_cache def dev_settings(): @@ -348,7 +361,12 @@ class IndexPipeline(BaseComponent): @Node.auto(depends_on=["Source", "Index", "embedding"]) def vector_indexing(self) -> VectorIndexing: return VectorIndexing( - vector_store=self.VS, doc_store=self.DS, embedding=self.embedding + vector_store=self.VS, + doc_store=self.DS, + embedding=self.embedding, + cache_dir=getattr( + settings, "KH_CHUNKS_OUTPUT_DIR", None + ), ) def handle_docs(self, docs, file_id, file_name) -> Generator[Document, None, int]: diff --git a/libs/ktem/ktem/reasoning/citation_display.py b/libs/ktem/ktem/reasoning/citation_display.py new file mode 100644 index 00000000..ba44bc6d --- /dev/null +++ b/libs/ktem/ktem/reasoning/citation_display.py @@ -0,0 +1,119 @@ +"""Citation display helpers. + +Renders citation evidence into HTML for the Gradio UI. +This logic lives in ktem (app layer) because it depends on +``ktem.utils.render.Render`` which is UI-specific and must +not be imported by the ``kotaemon`` framework layer. +""" + +from __future__ import annotations + +import logging + +from kotaemon.base import Document +from kotaemon.indices.qa.citation_qa import ( + CONTEXT_RELEVANT_WARNING_SCORE, + AnswerWithContextPipeline, +) + +from ktem.utils.render import Render + +logger = logging.getLogger(__name__) + + +def prepare_citations( + pipeline: AnswerWithContextPipeline, + answer: Document, + docs: list[Document], +) -> tuple[list[Document], list[Document]]: + """Prepare citation documents for UI display. + + Delegates evidence-matching to the framework-level + ``pipeline.match_evidence_with_context``, then + renders the results with ``Render``. + """ + with_citation: list[Document] = [] + without_citation: list[Document] = [] + has_llm_score = any( + "llm_trulens_score" in doc.metadata for doc in docs + ) + + spans = pipeline.match_evidence_with_context(answer, docs) + id2docs = {doc.doc_id: doc for doc in docs} + not_detected = set(id2docs.keys()) - set(spans.keys()) + + for _id, ss in spans.items(): + if not ss: + not_detected.add(_id) + continue + cur_doc = id2docs[_id] + highlight_text = "" + + ss = sorted(ss, key=lambda x: x["start"]) + last_end = 0 + text = cur_doc.text[: ss[0]["start"]] + + for idx, span in enumerate(ss): + span_start = max(last_end, span["start"]) + span_end = max(last_end, span["end"]) + + to_highlight = cur_doc.text[span_start:span_end] + last_end = span_end + + highlight_text += ( + (" " if highlight_text else "") + to_highlight + ) + + span_idx = span.get("idx", None) + if span_idx is not None: + to_highlight = f"\u3010{span_idx}\u3011" + to_highlight + + text += Render.highlight( + to_highlight, + elem_id=( + str(span_idx) if span_idx is not None else None + ), + ) + if idx < len(ss) - 1: + text += cur_doc.text[ + span["end"] : ss[idx + 1]["start"] + ] + + text += cur_doc.text[ss[-1]["end"] :] + with_citation.append( + Document( + channel="info", + content=Render.collapsible_with_header_score( + cur_doc, + override_text=text, + highlight_text=highlight_text, + open_collapsible=True, + ), + ) + ) + + logger.info("Got %d cited docs", len(with_citation)) + + sorted_not_detected = sorted( + not_detected, + key=lambda id_: id2docs[id_].metadata.get( + "llm_trulens_score", 0.0 + ), + reverse=True, + ) + + for id_ in sorted_not_detected: + doc = id2docs[id_] + doc_score = doc.metadata.get("llm_trulens_score", 0.0) + is_open = not has_llm_score or ( + doc_score > CONTEXT_RELEVANT_WARNING_SCORE + ) + without_citation.append( + Document( + channel="info", + content=Render.collapsible_with_header_score( + doc, open_collapsible=is_open + ), + ) + ) + return with_citation, without_citation diff --git a/libs/ktem/ktem/reasoning/simple.py b/libs/ktem/ktem/reasoning/simple.py index c40c9659..80ba13f6 100644 --- a/libs/ktem/ktem/reasoning/simple.py +++ b/libs/ktem/ktem/reasoning/simple.py @@ -6,6 +6,7 @@ from typing import Generator from decouple import config from ktem.embeddings.manager import embedding_models_manager as embeddings from ktem.llms.manager import llms +from theflow.settings import settings as flowsettings from ktem.reasoning.prompt_optimization import ( DecomposeQuestionPipeline, RewriteQuestionPipeline, @@ -23,11 +24,17 @@ from kotaemon.base import ( RetrievedDocument, SystemMessage, ) +from kotaemon.indices.qa.citation import CitationPipeline from kotaemon.indices.qa.citation_qa import ( CONTEXT_RELEVANT_WARNING_SCORE, DEFAULT_QA_TEXT_PROMPT, AnswerWithContextPipeline, ) + +from ktem.reasoning.citation_display import prepare_citations +from ktem.reasoning.prompt_optimization.mindmap import ( + CreateMindmapPipeline, +) from kotaemon.indices.qa.citation_qa_inline import AnswerWithInlineCitation from kotaemon.indices.qa.format_context import PrepareEvidencePipeline from kotaemon.indices.qa.utils import replace_think_tag_with_details @@ -222,8 +229,8 @@ class FullQAPipeline(BaseReasoning): def show_citations_and_addons(self, answer, docs, question): # show the evidence - with_citation, without_citation = self.answering_pipeline.prepare_citations( - answer, docs + with_citation, without_citation = prepare_citations( + self.answering_pipeline, answer, docs ) mindmap_output = self.prepare_mindmap(answer) citation_plot_output = self.prepare_citation_viz(answer, question, docs) @@ -366,7 +373,10 @@ class FullQAPipeline(BaseReasoning): answer_pipeline = pipeline.answering_pipeline = AnswerWithContextPipeline() answer_pipeline.llm = llm - answer_pipeline.citation_pipeline.llm = llm + answer_pipeline.citation_pipeline = CitationPipeline(llm=llm) + answer_pipeline.create_mindmap_pipeline = CreateMindmapPipeline( + llm=llm + ) answer_pipeline.n_last_interactions = settings[f"{prefix}.n_last_interactions"] answer_pipeline.enable_citation = ( settings[f"{prefix}.highlight_citation"] != "off" @@ -374,6 +384,9 @@ class FullQAPipeline(BaseReasoning): answer_pipeline.enable_mindmap = settings[f"{prefix}.create_mindmap"] answer_pipeline.enable_citation_viz = settings[f"{prefix}.create_citation_viz"] answer_pipeline.use_multimodal = settings[f"{prefix}.use_multimodal"] + answer_pipeline.vlm_endpoint = getattr( + flowsettings, "KH_VLM_ENDPOINT", "" + ) answer_pipeline.system_prompt = settings[f"{prefix}.system_prompt"] answer_pipeline.qa_template = settings[f"{prefix}.qa_prompt"] answer_pipeline.lang = SUPPORTED_LANGUAGE_MAP.get( @@ -563,8 +576,8 @@ class FullDecomposeQAPipeline(FullQAPipeline): ) # show the evidence - with_citation, without_citation = self.answering_pipeline.prepare_citations( - answer, docs + with_citation, without_citation = prepare_citations( + self.answering_pipeline, answer, docs ) if not with_citation and not without_citation: yield Document(channel="info", content="