diff --git a/libs/kotaemon/kotaemon/indices/qa/citation_qa.py b/libs/kotaemon/kotaemon/indices/qa/citation_qa.py index 1b014910..50894b86 100644 --- a/libs/kotaemon/kotaemon/indices/qa/citation_qa.py +++ b/libs/kotaemon/kotaemon/indices/qa/citation_qa.py @@ -201,7 +201,11 @@ class AnswerWithContextPipeline(BaseComponent): def mindmap_call(): nonlocal mindmap - mindmap = self.create_mindmap_pipeline(context=evidence, question=question) + + if self.create_mindmap_pipeline is not None: + mindmap = self.create_mindmap_pipeline( + context=evidence, question=question + ) citation_thread = None mindmap_thread = None @@ -306,4 +310,3 @@ class AnswerWithContextPipeline(BaseComponent): # print("Matched citation:", quote, matched_excerpts), return spans - diff --git a/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py b/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py index 77ff8d30..356e267e 100644 --- a/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py +++ b/libs/kotaemon/kotaemon/indices/qa/citation_qa_inline.py @@ -214,7 +214,10 @@ class AnswerWithInlineCitation(AnswerWithContextPipeline): def mindmap_call(): nonlocal mindmap - mindmap = self.create_mindmap_pipeline(context=evidence, question=question) + if self.create_mindmap_pipeline is not None: + mindmap = self.create_mindmap_pipeline( + context=evidence, question=question + ) mindmap_thread = None diff --git a/libs/kotaemon/kotaemon/indices/rankings/cohere.py b/libs/kotaemon/kotaemon/indices/rankings/cohere.py index 70c18b37..1cd369e4 100644 --- a/libs/kotaemon/kotaemon/indices/rankings/cohere.py +++ b/libs/kotaemon/kotaemon/indices/rankings/cohere.py @@ -18,8 +18,7 @@ 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" ) if not self.cohere_api_key: diff --git a/libs/ktem/ktem/index/file/pipelines.py b/libs/ktem/ktem/index/file/pipelines.py index f5320600..4e81fa5a 100644 --- a/libs/ktem/ktem/index/file/pipelines.py +++ b/libs/ktem/ktem/index/file/pipelines.py @@ -364,9 +364,7 @@ class IndexPipeline(BaseComponent): vector_store=self.VS, doc_store=self.DS, embedding=self.embedding, - cache_dir=getattr( - settings, "KH_CHUNKS_OUTPUT_DIR", None - ), + 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 index ba44bc6d..be8c9e37 100644 --- a/libs/ktem/ktem/reasoning/citation_display.py +++ b/libs/ktem/ktem/reasoning/citation_display.py @@ -10,14 +10,14 @@ from __future__ import annotations import logging +from ktem.utils.render import Render + 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__) @@ -34,9 +34,7 @@ def prepare_citations( """ with_citation: list[Document] = [] without_citation: list[Document] = [] - has_llm_score = any( - "llm_trulens_score" in doc.metadata for doc in docs - ) + 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} @@ -60,9 +58,7 @@ def prepare_citations( to_highlight = cur_doc.text[span_start:span_end] last_end = span_end - highlight_text += ( - (" " if highlight_text else "") + to_highlight - ) + highlight_text += (" " if highlight_text else "") + to_highlight span_idx = span.get("idx", None) if span_idx is not None: @@ -70,14 +66,10 @@ def prepare_citations( text += Render.highlight( to_highlight, - elem_id=( - str(span_idx) if span_idx is not None else None - ), + 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[span["end"] : ss[idx + 1]["start"]] text += cur_doc.text[ss[-1]["end"] :] with_citation.append( @@ -96,18 +88,14 @@ def prepare_citations( sorted_not_detected = sorted( not_detected, - key=lambda id_: id2docs[id_].metadata.get( - "llm_trulens_score", 0.0 - ), + 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 - ) + is_open = not has_llm_score or (doc_score > CONTEXT_RELEVANT_WARNING_SCORE) without_citation.append( Document( channel="info", diff --git a/libs/ktem/ktem/reasoning/simple.py b/libs/ktem/ktem/reasoning/simple.py index 80ba13f6..7e197319 100644 --- a/libs/ktem/ktem/reasoning/simple.py +++ b/libs/ktem/ktem/reasoning/simple.py @@ -6,14 +6,16 @@ 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.citation_display import prepare_citations from ktem.reasoning.prompt_optimization import ( DecomposeQuestionPipeline, RewriteQuestionPipeline, ) +from ktem.reasoning.prompt_optimization.mindmap import CreateMindmapPipeline from ktem.utils.render import Render from ktem.utils.visualize_cited import CreateCitationVizPipeline from plotly.io import to_json +from theflow.settings import settings as flowsettings from kotaemon.base import ( AIMessage, @@ -30,11 +32,6 @@ from kotaemon.indices.qa.citation_qa import ( 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 @@ -374,9 +371,7 @@ class FullQAPipeline(BaseReasoning): answer_pipeline.llm = llm answer_pipeline.citation_pipeline = CitationPipeline(llm=llm) - answer_pipeline.create_mindmap_pipeline = CreateMindmapPipeline( - 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" @@ -384,9 +379,7 @@ 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.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(