refactor: update DependencyContainer and ConfluenceExtractor for improved summarization handling

Author: a-klos

libs/admin-api-lib/src/admin_api_lib/dependency_container.py

@@ -138,14 +138,9 @@ class DependencyContainer(DeclarativeContainer):
         chunker=summary_text_splitter,
         semaphore=Singleton(AsyncThreadsafeSemaphore, summarizer_settings.maximum_concurrreny),
     )
-    traced_summarizer = Singleton(
-        LangfuseTracedRunnable,
-        inner_chain=summarizer,
-        settings=langfuse_settings,
-    )
 
     summary_enhancer = List(
-        Singleton(PageSummaryEnhancer, traced_summarizer, chunker_settings),
+        Singleton(PageSummaryEnhancer, summarizer, chunker_settings),
     )
     untraced_information_enhancer = Singleton(
         GeneralEnhancer,

libs/extractor-api-lib/src/extractor_api_lib/impl/extractors/confluence_extractor.py

@@ -1,72 +1,104 @@
-"""Module for the DefaultConfluenceExtractor class."""
+"""Module for the LangchainSummarizer class."""
 
 import logging
-from langchain_community.document_loaders import ConfluenceLoader
-
-from extractor_api_lib.impl.types.extractor_types import ExtractorTypes
-from extractor_api_lib.models.dataclasses.internal_information_piece import InternalInformationPiece
-from extractor_api_lib.models.extraction_parameters import ExtractionParameters
-from extractor_api_lib.extractors.information_extractor import InformationExtractor
-from extractor_api_lib.impl.mapper.confluence_langchain_document2information_piece import (
-    ConfluenceLangchainDocument2InformationPiece,
+import traceback
+from typing import Optional
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langchain_core.runnables import Runnable, RunnableConfig, ensure_config
+
+from admin_api_lib.summarizer.summarizer import (
+    Summarizer,
+    SummarizerInput,
+    SummarizerOutput,
 )
+from rag_core_lib.impl.langfuse_manager.langfuse_manager import LangfuseManager
+from rag_core_lib.impl.utils.async_threadsafe_semaphore import AsyncThreadsafeSemaphore
 
 logger = logging.getLogger(__name__)
 
 
-class ConfluenceExtractor(InformationExtractor):
-    """Implementation of the InformationExtractor interface for confluence."""
+class LangchainSummarizer(Summarizer):
+    """Is responsible for summarizing input data.
+
+    LangchainSummarizer is responsible for summarizing input data using the LangfuseManager,
+    RecursiveCharacterTextSplitter, and AsyncThreadsafeSemaphore. It handles chunking of the input
+    document and retries the summarization process if an error occurs.
+    """
 
     def __init__(
         self,
-        mapper: ConfluenceLangchainDocument2InformationPiece,
+        langfuse_manager: LangfuseManager,
+        chunker: RecursiveCharacterTextSplitter,
+        semaphore: AsyncThreadsafeSemaphore,
     ):
-        """
-        Initialize the ConfluenceExtractor.
-
-        Parameters
-        ----------
-        mapper : ConfluenceLangchainDocument2InformationPiece
-            An instance of ConfluenceLangchainDocument2InformationPiece used for mapping langchain documents
-            to information pieces.
-        """
-        self._mapper = mapper
+        self._chunker = chunker
+        self._langfuse_manager = langfuse_manager
+        self._semaphore = semaphore
 
-    @property
-    def extractor_type(self) -> ExtractorTypes:
-        return ExtractorTypes.CONFLUENCE
-
-    async def aextract_content(
-        self,
-        extraction_parameters: ExtractionParameters,
-    ) -> list[InternalInformationPiece]:
+    async def ainvoke(self, query: SummarizerInput, config: Optional[RunnableConfig] = None) -> SummarizerOutput:
         """
-        Asynchronously extracts information pieces from Confluence.
+        Asynchronously invokes the summarization process on the given query.
 
         Parameters
         ----------
-        extraction_parameters : ExtractionParameters
-            The parameters required to connect to and extract data from Confluence.
+        query : SummarizerInput
+            The input data to be summarized.
+        config : Optional[RunnableConfig], optional
+            Configuration options for the summarization process, by default None.
 
         Returns
         -------
-        list[InternalInformationPiece]
-            A list of information pieces extracted from Confluence.
+        SummarizerOutput
+            The summarized output.
+
+        Raises
+        ------
+        Exception
+            If the summary creation fails after the allowed number of tries.
+
+        Notes
+        -----
+        This method handles chunking of the input document and retries the summarization
+        process if an error occurs, up to the number of tries specified in the config.
         """
-        # Convert list of key value pairs to dict
-        confluence_loader_parameters = {
-            x.key: int(x.value) if x.value.isdigit() else x.value for x in extraction_parameters.kwargs
-        }
-        if not confluence_loader_parameters.get("max_pages") or isinstance(
-            confluence_loader_parameters.get("max_pages"), str
-        ):
-            logging.warning(
-                "max_pages parameter is not set or invalid discarding it. ConfluenceLoader will use default value."
-            )
-            confluence_loader_parameters.pop("max_pages")
-        # Drop the document_name parameter as it is not used by the ConfluenceLoader
-        if "document_name" in confluence_loader_parameters:
-            confluence_loader_parameters.pop("document_name", None)
-        document_loader = ConfluenceLoader(**confluence_loader_parameters)
-        documents = document_loader.load()
-        return [self._mapper.map_document2informationpiece(x, extraction_parameters.document_name) for x in documents]
+        assert query, "Query is empty: %s" % query  # noqa S101
+        config = ensure_config(config)
+        tries_remaining = config.get("configurable", {}).get("tries_remaining", 3)
+        logger.debug("Tries remaining %d" % tries_remaining)
+
+        if tries_remaining < 0:
+            raise Exception("Summary creation failed.")
+        document = Document(page_content=query)
+        langchain_documents = self._chunker.split_documents([document])
+
+        outputs = []
+        for langchain_document in langchain_documents:
+            async with self._semaphore:
+                try:
+                    result = await self._create_chain().ainvoke({"text": langchain_document.page_content}, config)
+                    # Extract content from AIMessage if it's not already a string
+                    content = result.content if hasattr(result, "content") else str(result)
+                    outputs.append(content)
+                except Exception as e:
+                    logger.error("Error in summarizing langchain doc: %s %s", e, traceback.format_exc())
+                    config["tries_remaining"] = tries_remaining - 1
+                    result = await self._create_chain().ainvoke({"text": langchain_document.page_content}, config)
+                    # Extract content from AIMessage if it's not already a string
+                    content = result.content if hasattr(result, "content") else str(result)
+                    outputs.append(content)
+
+        if len(outputs) == 1:
+            return outputs[0]
+        summary = " ".join(outputs)
+        logger.debug(
+            "Reduced number of chars from %d to %d"
+            % (len("".join([x.page_content for x in langchain_documents])), len(summary))
+        )
+        return await self.ainvoke(summary, config)
+
+    def _create_chain(self) -> Runnable:
+        return self._langfuse_manager.get_base_prompt(self.__class__.__name__) | self._langfuse_manager.get_base_llm(
+            self.__class__.__name__
+        )

libs/rag-core-api/src/rag_core_api/impl/evaluator/langfuse_ragas_evaluator.py

@@ -7,12 +7,14 @@
 from asyncio import gather
 from datetime import datetime
 from json import JSONDecodeError
+from time import sleep
 from uuid import uuid4
 
 import ragas
 from datasets import Dataset
 from langchain_core.runnables import RunnableConfig
 from langfuse import Langfuse
+from langfuse.api.core.api_error import ApiError
 from langfuse.api.resources.commons.errors.not_found_error import NotFoundError
 from langfuse._client.datasets import DatasetClient
 from ragas.llms import LangchainLLMWrapper
@@ -162,61 +164,67 @@ async def _aevaluate_question(self, item, experiment_name: str, generation_time:
         async with self._semaphore:
             chat_request = ChatRequest(message=item.input)
 
-            # Use item.run context manager for trace
-            with item.run(
-                run_name=experiment_name,
-                run_metadata={"model": self._settings.model},
-                run_description=f"Evaluation run for {experiment_name}",
-            ) as root_span:
-                # Use langfuse.start_as_current_generation for generation
-                try:
-                    response = await self._chat_endpoint.achat(config["metadata"]["session_id"], chat_request)
-                except Exception as e:
-                    logger.info("Error while answering question %s: %s", item.input, e)
-                    response = None
-                output = {
-                    "answer": response.answer if response else None,
-                    "documents": (
-                        [x.page_content for x in response.citations] if response and response.citations else None
-                    ),
-                }
-                with self._langfuse.start_as_current_generation(
-                    name="rag-eval-llm-call",
-                    input={"question": item.input, "context": output["documents"]},
-                    metadata={"item_id": item.id, "run": experiment_name},
-                    model=self._settings.model,
-                ) as generation:
-
-                    generation.update(output=output["answer"])
-                    generation.update_trace(
-                        input={"question": item.input, "context": output["documents"]},
-                        metadata={"item_id": item.id, "run": experiment_name},
-                        output=output["answer"],
-                    )
+            try:
+                response = await self._chat_endpoint.achat(config["metadata"]["session_id"], chat_request)
+            except Exception as e:
+                logger.info("Error while answering question %s: %s", item.input, e)
+                response = None
 
-                # Ragas metrics
-                if response and response.citations:
-                    eval_data = Dataset.from_dict(
-                        {
-                            "question": [item.input],
-                            "answer": [output["answer"]],
-                            "contexts": [output["documents"]],
-                            "ground_truth": [item.expected_output],
-                        }
-                    )
-                    result = ragas.evaluate(
-                        eval_data,
-                        metrics=self.METRICS,
-                        llm=self._llm_wrapped,
-                        embeddings=self._embedder,
+            if response and response.citations:
+                output = {"answer": response.answer, "documents": [x.page_content for x in response.citations]}
+            else:
+                output = {"answer": None, "documents": None}
+
+            langfuse_generation = self._langfuse.generation(
+                name=self._settings.evaluation_dataset_name,
+                input=item.input,
+                output=output,
+                start_time=generation_time,
+                end_time=datetime.now(),
+            )
+            self._link_item2generation(item, langfuse_generation, experiment_name)
+
+            if not (response and response.citations):
+                for metric in self.METRICS:
+                    langfuse_generation.score(
+                        name=metric.name,
+                        value=self.DEFAULT_SCORE_VALUE,
                     )
-                    for metric, score in result.scores[0].items():
-                        if math.isnan(score):
-                            score = self.DEFAULT_SCORE_VALUE
-                        root_span.score_trace(name=metric, value=score)
-                else:
-                    for metric in self.METRICS:
-                        root_span.score_trace(name=metric.name, value=self.DEFAULT_SCORE_VALUE)
+                return
+
+            eval_data = Dataset.from_dict(
+                {
+                    "question": [item.input],
+                    "answer": [output["answer"]],
+                    "contexts": [output["documents"]],
+                    "ground_truth": [item.expected_output],
+                }
+            )
+
+            result = ragas.evaluate(
+                eval_data,
+                metrics=self.METRICS,
+                llm=self._llm_wrapped,
+                embeddings=self._embedder,
+            )
+            for metric, score in result.scores[0].items():
+                if math.isnan(score):
+                    score = self.DEFAULT_SCORE_VALUE
+                langfuse_generation.score(
+                    name=metric,
+                    value=score,
+                )
+
+    def _link_item2generation(self, item, generation, experiment_name, retries: int = 0):
+        try:
+            item.link(generation, experiment_name)
+        except ApiError as e:
+            logger.warning("Failed to link item to generation: %s", e)
+            retries += 1
+            if retries > self.MAX_RETRIES:
+                raise e
+            sleep(1)
+            self._link_item2generation(item, generation, experiment_name, retries)
 
     def _get_dataset(self, dataset_name: str) -> DatasetClient:
         dataset = None
@@ -232,7 +240,7 @@ def _get_dataset(self, dataset_name: str) -> DatasetClient:
         return dataset
 
     def _create_dataset(self, dataset_name: str = None):
-        self._langfuse.create_dataset(name=dataset_name)
+        self._langfuse.create_dataset(dataset_name)
 
         data = self._load_dataset_items()
         self._store_items_in_dataset(data, dataset_name)

libs/rag-core-lib/src/rag_core_lib/tracers/traced_runnable.py

@@ -66,11 +66,9 @@ async def ainvoke(
         config = ensure_config(config)
         session_id = self._get_session_id(config)
         config_with_tracing = self._add_tracing_callback(config)
-        with self.langfuse_client.start_as_current_span(name=self._inner_chain.__class__.__name__) as span:
-            span.update_trace(session_id=session_id, input=chain_input)
-            output = await self._inner_chain.ainvoke(chain_input, config=config_with_tracing)
-            span.update_trace(output=output)
-            return output
+        with self.langfuse_client.start_as_current_span(name="traced_runnable") as span:
+            span.update_trace(session_id=session_id)
+            return await self._inner_chain.ainvoke(chain_input, config=config_with_tracing)
 
     @abstractmethod
     def _add_tracing_callback(self, config: Optional[RunnableConfig]) -> RunnableConfig: ...