refactor: decouple kotaemon and ktem

This commit is contained in:
phv2312
2026-05-31 13:45:11 +07:00
parent d13d9e6165
commit 7e17d2a613
10 changed files with 171 additions and 137 deletions

View File

@@ -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)

View File

@@ -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

View File

@@ -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()

View File

@@ -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

View File

@@ -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

View File

@@ -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 = "",

View File

@@ -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]:

View File

@@ -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

View File

@@ -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="<h5><b>No evidence found.</b></h5>")

View File

@@ -19,6 +19,7 @@ dynamic = ["version"]
requires-python = ">= 3.10"
description = "RAG-based Question and Answering Application"
dependencies = [
"kotaemon",
"click>=8.1.7,<9",
"platformdirs>=4.2.1,<5",
"pluggy>=1.5.0,<2",
@@ -43,3 +44,6 @@ classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
]
[tool.uv.sources]
kotaemon = { workspace = true }