Add Uptrain Callback Handler (run-llama#10432)

* Add UpTrain Callback Handler

* Add logging

* Add note for project_name prefix

* Minor fixes

* Add explanations

* Add showcase

* Add callback to uptrain.md

* update import

* linter

* Update .gitignore file to ignore .venv and __pycache__ directories

* Fix imports

* create llama-index-callback-uptrain

* Update UpTrainCallback.ipynb

* Update README.md

* Run pants tailor

docs/examples/evaluation/UpTrain.ipynb

@@ -5,7 +5,7 @@
    "id": "ec212316",
    "metadata": {},
    "source": [
-    "<a href=\"https://colab.research.google.com/github/uptrain-ai/llama_index/blob/uptrain_integration/docs/examples/evaluation/UpTrain.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/evaluation/UpTrain.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
    ]
   },
   {

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/BUILD

@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/Makefile

@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help:	## Show all Makefile targets.
+	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format:	## Run code autoformatters (black).
+	pre-commit install
+	git ls-files | xargs pre-commit run black --files
+
+lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
+	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test:	## Run tests via pytest.
+	pytest tests
+
+watch-docs:	## Build and watch documentation.
+	sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/README.md

@@ -0,0 +1,35 @@
+# LlamaIndex Callbacks Integration: UpTrain
+
+UpTrain is an open-source tool to evaluate and monitor the performance of language models. It provides a set of pre-built evaluations to assess the quality of responses generated by the model. Once you add UpTrainCallbackHandler to your existing LlamaIndex pipeline, it will take care of sending the generated responses to the UpTrain Managed Service for evaluations and display the results in the output.
+
+Three additional evaluations for Llamaindex have been introduced, complementing existing ones. These evaluations run automatically, with results displayed in the output. More details on UpTrain's evaluations can be found [here](https://github.com/uptrain-ai/uptrain?tab=readme-ov-file#pre-built-evaluations-we-offer-).
+
+Selected operators from the LlamaIndex pipeline are highlighted for demonstration:
+
+## 1. **RAG Query Engine Evaluations**:
+
+The RAG query engine plays a crucial role in retrieving context and generating responses. To ensure its performance and response quality, we conduct the following evaluations:
+
+- **Context Relevance**: Determines if the context extracted from the query is relevant to the response.
+- **Factual Accuracy**: Assesses if the LLM is hallcuinating or providing incorrect information.
+- **Response Completeness**: Checks if the response contains all the information requested by the query.
+
+## 2. **Sub-Question Query Generation Evaluation**:
+
+The SubQuestionQueryGeneration operator decomposes a question into sub-questions, generating responses for each using a RAG query engine. Given the complexity, we include the previous evaluations and add:
+
+- **Sub Query Completeness**: Assures that the sub-questions accurately and comprehensively cover the original query.
+
+## 3. **Re-Ranking Evaluations**:
+
+Re-ranking involves reordering nodes based on relevance to the query and choosing top n nodes. Different evaluations are performed based on the number of nodes returned after re-ranking.
+
+a. Same Number of Nodes
+
+- **Context Reranking**: Checks if the order of re-ranked nodes is more relevant to the query than the original order.
+
+b. Different Number of Nodes:
+
+- **Context Conciseness**: Examines whether the reduced number of nodes still provides all the required information.
+
+These evaluations collectively ensure the robustness and effectiveness of the RAG query engine, SubQuestionQueryGeneration operator, and the re-ranking process in the LlamaIndex pipeline.

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/llama_index/callbacks/uptrain/BUILD

@@ -0,0 +1 @@
+python_sources()

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/llama_index/callbacks/uptrain/__init__.py

@@ -0,0 +1,3 @@
+from llama_index.callbacks.uptrain.base import UpTrainCallbackHandler
+
+__all__ = ["UpTrainCallbackHandler"]

llama-index-integrations/callbacks/llama-index-callbacks-uptrain/llama_index/callbacks/uptrain/base.py

@@ -0,0 +1,320 @@
+from collections import defaultdict
+from typing import Any, DefaultDict, Dict, List, Literal, Optional, Set
+
+import nest_asyncio
+
+from llama_index.core.callbacks.base_handler import BaseCallbackHandler
+from llama_index.core.callbacks.schema import (
+    CBEvent,
+    CBEventType,
+)
+
+
+class UpTrainDataSchema:
+    """UpTrain data schema."""
+
+    def __init__(self, project_name_prefix: str) -> None:
+        """Initialize the UpTrain data schema."""
+        # For tracking project name and results
+        self.project_name_prefix: str = project_name_prefix
+        self.uptrain_results: DefaultDict[str, Any] = defaultdict(list)
+
+        # For tracking event types - reranking, sub_question
+        self.eval_types: Set[str] = set()
+
+        ## SYNTHESIZE
+        self.question: str = ""
+        self.context: str = ""
+        self.response: str = ""
+
+        ## RERANKING
+        self.old_context: List[str] = []
+        self.new_context: List[str] = []
+
+        ## SUB_QUESTION
+        # Map of sub question ID to question, context, and response
+        self.sub_question_map: DefaultDict[str, dict] = defaultdict(dict)
+        # Parent ID of sub questions
+        self.sub_question_parent_id: str = ""
+        # Parent question
+        self.parent_question: str = ""
+
+
+class UpTrainCallbackHandler(BaseCallbackHandler):
+    """
+    UpTrain callback handler.
+
+    This class is responsible for handling the UpTrain API and logging events to UpTrain.
+
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        key_type: Literal["uptrain", "openai"],
+        project_name_prefix: str = "llama",
+    ) -> None:
+        """Initialize the UpTrain callback handler."""
+        try:
+            from uptrain import APIClient, EvalLLM, Settings
+        except ImportError:
+            raise ImportError(
+                "UpTrainCallbackHandler requires the 'uptrain' package. "
+                "Please install it using 'pip install uptrain'."
+            )
+        nest_asyncio.apply()
+        super().__init__(
+            event_starts_to_ignore=[],
+            event_ends_to_ignore=[],
+        )
+        self.schema = UpTrainDataSchema(project_name_prefix=project_name_prefix)
+        self._event_pairs_by_id: Dict[str, List[CBEvent]] = defaultdict(list)
+        self._trace_map: Dict[str, List[str]] = defaultdict(list)
+
+        # Based on whether the user enters an UpTrain API key or an OpenAI API key, the client is initialized
+        # If both are entered, the UpTrain API key is used
+        if key_type == "uptrain":
+            settings = Settings(uptrain_access_token=api_key)
+            self.uptrain_client = APIClient(settings=settings)
+        elif key_type == "openai":
+            settings = Settings(openai_api_key=api_key)
+            self.uptrain_client = EvalLLM(settings=settings)
+        else:
+            raise ValueError("Invalid key type: Must be 'uptrain' or 'openai'")
+
+    def uptrain_evaluate(
+        self,
+        project_name: str,
+        data: List[Dict[str, str]],
+        checks: List[str],
+    ) -> None:
+        """Run an evaluation on the UpTrain server using UpTrain client."""
+        if self.uptrain_client.__class__.__name__ == "APIClient":
+            uptrain_result = self.uptrain_client.log_and_evaluate(
+                project_name=project_name,
+                data=data,
+                checks=checks,
+            )
+        else:
+            uptrain_result = self.uptrain_client.evaluate(
+                data=data,
+                checks=checks,
+            )
+        self.schema.uptrain_results[project_name].append(uptrain_result)
+
+        score_name_map = {
+            "score_context_relevance": "Context Relevance Score",
+            "score_factual_accuracy": "Factual Accuracy Score",
+            "score_response_completeness": "Response Completeness Score",
+            "score_sub_query_completeness": "Sub Query Completeness Score",
+            "score_context_reranking": "Context Reranking Score",
+            "score_context_conciseness": "Context Conciseness Score",
+        }
+
+        # Print the results
+        for row in uptrain_result:
+            columns = list(row.keys())
+            for column in columns:
+                if column == "question":
+                    print(f"\nQuestion: {row[column]}")
+                elif column == "response":
+                    print(f"Response: {row[column]}")
+                elif column.startswith("score"):
+                    if column in score_name_map:
+                        print(f"{score_name_map[column]}: {row[column]}")
+                    else:
+                        print(f"{column}: {row[column]}")
+            print()
+
+    def on_event_start(
+        self,
+        event_type: CBEventType,
+        payload: Any = None,
+        event_id: str = "",
+        parent_id: str = "",
+        **kwargs: Any,
+    ) -> str:
+        """Run when an event starts and return id of event."""
+        event = CBEvent(event_type, payload=payload, id_=event_id)
+        self._event_pairs_by_id[event.id_].append(event)
+
+        if event_type is CBEventType.QUERY:
+            self.schema.question = payload["query_str"]
+        if event_type is CBEventType.TEMPLATING and "template_vars" in payload:
+            template_vars = payload["template_vars"]
+            self.schema.context = template_vars.get("context_str", "")
+        elif event_type is CBEventType.RERANKING and "nodes" in payload:
+            self.schema.eval_types.add("reranking")
+            # Store old context data
+            self.schema.old_context = [node.text for node in payload["nodes"]]
+        elif event_type is CBEventType.SUB_QUESTION:
+            # For the first sub question, store parent question and parent id
+            if "sub_question" not in self.schema.eval_types:
+                self.schema.parent_question = self.schema.question
+                self.schema.eval_types.add("sub_question")
+            # Store sub question data - question and parent id
+            self.schema.sub_question_parent_id = parent_id
+        return event_id
+
+    def on_event_end(
+        self,
+        event_type: CBEventType,
+        payload: Any = None,
+        event_id: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Run when an event ends."""
+        try:
+            from uptrain import Evals
+        except ImportError:
+            raise ImportError(
+                "UpTrainCallbackHandler requires the 'uptrain' package. "
+                "Please install it using 'pip install uptrain'."
+            )
+        event = CBEvent(event_type, payload=payload, id_=event_id)
+        self._event_pairs_by_id[event.id_].append(event)
+        self._trace_map = defaultdict(list)
+        if event_id == self.schema.sub_question_parent_id:
+            # Perform individual evaluations for sub questions (but send all sub questions at once)
+            self.uptrain_evaluate(
+                project_name=f"{self.schema.project_name_prefix}_sub_question_answering",
+                data=list(self.schema.sub_question_map.values()),
+                checks=[
+                    Evals.CONTEXT_RELEVANCE,
+                    Evals.FACTUAL_ACCURACY,
+                    Evals.RESPONSE_COMPLETENESS,
+                ],
+            )
+            # Perform evaluation for question and all sub questions (as a whole)
+            sub_questions = [
+                sub_question["question"]
+                for sub_question in self.schema.sub_question_map.values()
+            ]
+            sub_questions_formatted = "\n".join(
+                [
+                    f"{index}. {string}"
+                    for index, string in enumerate(sub_questions, start=1)
+                ]
+            )
+            self.uptrain_evaluate(
+                project_name=f"{self.schema.project_name_prefix}_sub_query_completeness",
+                data=[
+                    {
+                        "question": self.schema.parent_question,
+                        "sub_questions": sub_questions_formatted,
+                    }
+                ],
+                checks=[Evals.SUB_QUERY_COMPLETENESS],
+            )
+        # Should not be called for sub questions
+        if (
+            event_type is CBEventType.SYNTHESIZE
+            and "sub_question" not in self.schema.eval_types
+        ):
+            self.schema.response = payload["response"].response
+            # Perform evaluation for synthesization
+            self.uptrain_evaluate(
+                project_name=f"{self.schema.project_name_prefix}_question_answering",
+                data=[
+                    {
+                        "question": self.schema.question,
+                        "context": self.schema.context,
+                        "response": self.schema.response,
+                    }
+                ],
+                checks=[
+                    Evals.CONTEXT_RELEVANCE,
+                    Evals.FACTUAL_ACCURACY,
+                    Evals.RESPONSE_COMPLETENESS,
+                ],
+            )
+
+        elif event_type is CBEventType.RERANKING:
+            # Store new context data
+            self.schema.new_context = [node.text for node in payload["nodes"]]
+            if len(self.schema.old_context) == len(self.schema.new_context):
+                context = "\n".join(
+                    [
+                        f"{index}. {string}"
+                        for index, string in enumerate(self.schema.old_context, start=1)
+                    ]
+                )
+                reranked_context = "\n".join(
+                    [
+                        f"{index}. {string}"
+                        for index, string in enumerate(self.schema.new_context, start=1)
+                    ]
+                )
+                # Perform evaluation for reranking
+                self.uptrain_evaluate(
+                    project_name=f"{self.schema.project_name_prefix}_context_reranking",
+                    data=[
+                        {
+                            "question": self.schema.question,
+                            "context": context,
+                            "reranked_context": reranked_context,
+                        }
+                    ],
+                    checks=[
+                        Evals.CONTEXT_RERANKING,
+                    ],
+                )
+            else:
+                context = "\n".join(self.schema.old_context)
+                concise_context = "\n".join(self.schema.new_context)
+                # Perform evaluation for resizing
+                self.uptrain_evaluate(
+                    project_name=f"{self.schema.project_name_prefix}_context_conciseness",
+                    data=[
+                        {
+                            "question": self.schema.question,
+                            "context": context,
+                            "concise_context": concise_context,
+                        }
+                    ],
+                    checks=[
+                        Evals.CONTEXT_CONCISENESS,
+                    ],
+                )
+        elif event_type is CBEventType.SUB_QUESTION:
+            # Store sub question data
+            self.schema.sub_question_map[event_id]["question"] = payload[
+                "sub_question"
+            ].sub_q.sub_question
+            self.schema.sub_question_map[event_id]["context"] = (
+                payload["sub_question"].sources[0].node.text
+            )
+            self.schema.sub_question_map[event_id]["response"] = payload[
+                "sub_question"
+            ].answer
+
+    def start_trace(self, trace_id: Optional[str] = None) -> None:
+        self._trace_map = defaultdict(list)
+        return super().start_trace(trace_id)
+
+    def end_trace(
+        self,
+        trace_id: Optional[str] = None,
+        trace_map: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        self._trace_map = trace_map or defaultdict(list)
+        return super().end_trace(trace_id, trace_map)
+
+    def build_trace_map(
+        self,
+        cur_event_id: str,
+        trace_map: Any,
+    ) -> Dict[str, Any]:
+        event_pair = self._event_pairs_by_id[cur_event_id]
+        if event_pair:
+            event_data = {
+                "event_type": event_pair[0].event_type,
+                "event_id": event_pair[0].id_,
+                "children": {},
+            }
+            trace_map[cur_event_id] = event_data
+
+        child_event_ids = self._trace_map[cur_event_id]
+        for child_event_id in child_event_ids:
+            self.build_trace_map(child_event_id, event_data["children"])
+        return trace_map