Merge branch 'main' into fix/remote-storage-v2

.gitignore

@@ -632,7 +632,6 @@ unstract/tool-registry/tool_registry_config/registry.yaml
 # Docker related
 # End of https://www.toptal.com/developers/gitignore/api/windows,macos,linux,pycharm,pycharm+all,pycharm+iml,python,visualstudiocode,react,django
 docker/temp/*
-docker/init.sql/*
 docker/*.env
 !docker/sample*.env
 docker/public_tools.json

backend/entrypoint.sh

@@ -4,11 +4,6 @@ cmd=$1
 if [ "$cmd" = "migrate" ]; then
     echo "Migration initiated"
     .venv/bin/python manage.py migrate
-elif [ "$cmd" = "prepare_and_migrate" ]; then
-    echo "Creating schema in database"
-    .venv/bin/python manage.py create_schema
-    echo "Migration initiated"
-    .venv/bin/python manage.py migrate
 fi
 
 # NOTE: Leaving below for reference incase required in the future

backend/prompt_studio/prompt_studio_core_v2/serializers.py

@@ -8,6 +8,9 @@
 from prompt_studio.prompt_profile_manager_v2.models import ProfileManager
 from prompt_studio.prompt_studio_core_v2.constants import ToolStudioKeys as TSKeys
 from prompt_studio.prompt_studio_core_v2.exceptions import DefaultProfileError
+from prompt_studio.prompt_studio_output_manager_v2.output_manager_util import (
+    OutputManagerUtils,
+)
 from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt
 from prompt_studio.prompt_studio_v2.serializers import ToolStudioPromptSerializer
 from rest_framework import serializers
@@ -55,11 +58,10 @@ def to_representation(self, instance):  # type: ignore
             profile_manager = ProfileManager.get_default_llm_profile(instance)
             data[TSKeys.DEFAULT_PROFILE] = profile_manager.profile_id
         except DefaultProfileError:
-            logger.info(
+            logger.warning(
                 "Default LLM profile doesnt exist for prompt tool %s",
                 str(instance.tool_id),
             )
-        try:
             prompt_instance: ToolStudioPrompt = ToolStudioPrompt.objects.filter(
                 tool_id=data.get(TSKeys.TOOL_ID)
             ).order_by("sequence_number")
@@ -69,11 +71,15 @@ def to_representation(self, instance):  # type: ignore
             if prompt_instance.count() != 0:
                 for prompt in prompt_instance:
                     prompt_serializer = ToolStudioPromptSerializer(prompt)
-                    output.append(prompt_serializer.data)
+                    coverage = OutputManagerUtils.get_coverage(
+                        data.get(TSKeys.TOOL_ID),
+                        profile_manager.profile_id,
+                        prompt.prompt_id,
+                    )
+                    serialized_data = prompt_serializer.data
+                    serialized_data["coverage"] = coverage
+                    output.append(serialized_data)
                 data[TSKeys.PROMPTS] = output
-        except Exception as e:
-            logger.error(f"Error occured while appending prompts {e}")
-            return data
 
         data["created_by_email"] = instance.created_by.email
 

backend/prompt_studio/prompt_studio_output_manager/serializers.py

@@ -1,6 +1,7 @@
 import json
 import logging
 
+from django.db.models import Count
 from usage.helper import UsageHelper
 
 from backend.serializers import AuditSerializer
@@ -26,6 +27,35 @@ def to_representation(self, instance):
             )
             token_usage = {}
         data["token_usage"] = token_usage
+        # Get the coverage for the current tool_id and profile_manager_id
+        try:
+            # Fetch all relevant outputs for the current tool and profile
+            related_outputs = (
+                PromptStudioOutputManager.objects.filter(
+                    tool_id=instance.tool_id,
+                    profile_manager_id=instance.profile_manager_id,
+                    prompt_id=instance.prompt_id,
+                )
+                .values("prompt_id", "profile_manager_id")
+                .annotate(document_count=Count("document_manager_id"))
+            )
+
+            coverage = {}
+            for output in related_outputs:
+                prompt_key = str(output["prompt_id"])
+                profile_key = str(output["profile_manager_id"])
+                coverage[f"coverage_{profile_key}_{prompt_key}"] = output[
+                    "document_count"
+                ]
+
+            data["coverage"] = coverage
+        except Exception as e:
+            logger.error(
+                "Error occurred while fetching "
+                f"coverage for tool_id {instance.tool_id} "
+                f"and profile_manager_id {instance.profile_manager_id}: {e}"
+            )
+            data["coverage"] = {}
         # Convert string to list
         try:
             data["context"] = json.loads(data["context"])

backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_util.py

@@ -0,0 +1,46 @@
+from django.db.models import Count
+from prompt_studio.prompt_studio_output_manager_v2.models import (
+    PromptStudioOutputManager,
+)
+
+
+class OutputManagerUtils:
+    @staticmethod
+    def get_coverage(
+        tool_id: str, profile_manager_id: str, prompt_id: str = None
+    ) -> dict[str, int]:
+        """
+        Method to fetch coverage data for given tool and profile manager.
+
+        Args:
+            tool (CustomTool): The tool instance or ID for which coverage is fetched.
+            profile_manager_id (str): The ID of the profile manager
+            for which coverage is calculated.
+            prompt_id (Optional[str]): The ID of the prompt (optional).
+            If provided, coverage is fetched for the specific prompt.
+
+        Returns:
+            dict[str, int]: A dictionary containing coverage information.
+                Keys are formatted as "coverage_<prompt_id>_<profile_manager_id>".
+                Values are the count of documents associated with each prompt
+                and profile combination.
+        """
+
+        prompt_outputs = (
+            PromptStudioOutputManager.objects.filter(
+                tool_id=tool_id,
+                profile_manager_id=profile_manager_id,
+                **({"prompt_id": prompt_id} if prompt_id else {}),
+            )
+            .values("prompt_id", "profile_manager_id")
+            .annotate(document_count=Count("document_manager_id"))
+        )
+
+        coverage = {}
+        for prompt_output in prompt_outputs:
+            prompt_key = str(prompt_output["prompt_id"])
+            profile_key = str(prompt_output["profile_manager_id"])
+            coverage[f"coverage_{prompt_key}_{profile_key}"] = prompt_output[
+                "document_count"
+            ]
+        return coverage

backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py

@@ -6,6 +6,7 @@
 from backend.serializers import AuditSerializer
 
 from .models import PromptStudioOutputManager
+from .output_manager_util import OutputManagerUtils
 
 logger = logging.getLogger(__name__)
 
@@ -16,16 +17,34 @@ class Meta:
         fields = "__all__"
 
     def to_representation(self, instance):
+
         data = super().to_representation(instance)
         try:
             token_usage = UsageHelper.get_aggregated_token_count(instance.run_id)
         except Exception as e:
-            logger.error(
+            logger.warning(
                 "Error occured while fetching token usage for run_id"
                 f"{instance.run_id}: {e}"
+                " | Process continued"
             )
             token_usage = {}
         data["token_usage"] = token_usage
+        # Get the coverage for the current tool_id and profile_manager_id
+        try:
+            # Fetch all relevant outputs for the current tool and profile
+            coverage = OutputManagerUtils.get_coverage(
+                instance.tool_id, instance.profile_manager_id, instance.prompt_id
+            )
+            data["coverage"] = coverage
+
+        except Exception as e:
+            logger.error(
+                "Error occurred while fetching "
+                f"coverage for tool_id {instance.tool_id} "
+                f"and profile_manager_id {instance.profile_manager_id}: {e}"
+                " | Process continued"
+            )
+            data["coverage"] = {}
         # Convert string to list
         try:
             context = data["context"]

backend/sample.env

@@ -146,6 +146,9 @@ TOOL_REGISTRY_CONFIG_PATH="/data/tool_registry_config"
 # Flipt Service
 FLIPT_SERVICE_AVAILABLE=False
 
+# Remote storage config for tool registry
+TOOL_REGISTRY_STORAGE_CREDENTIALS='{"provider":"local"}'
+
 
 # File System Configuration for Workflow and API Execution
 

docker/docker-compose-dev-essentials.yaml

@@ -9,7 +9,7 @@ services:
       - "5432:5432"
     volumes:
       - postgres_data:/var/lib/postgresql/data/
-      - ../backend/init.sql:/docker-entrypoint-initdb.d/init.sql
+      - ./scripts/db-setup/db_setup.sh:/docker-entrypoint-initdb.d/db_setup.sh
     env_file:
       - ./essentials.env
     labels:

docker/dockerfiles/platform.Dockerfile

@@ -12,7 +12,14 @@ ENV BUILD_CONTEXT_PATH platform-service
 ENV BUILD_PACKAGES_PATH unstract
 ENV PDM_VERSION 2.16.1
 
-RUN pip install --no-cache-dir -U pip pdm~=${PDM_VERSION}; \
+RUN apt-get update; \
+    apt-get --no-install-recommends install -y  \
+    # unstract sdk
+    build-essential libmagic-dev; \
+    \
+    apt-get clean && rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/*; \
+    \
+    pip install --no-cache-dir -U pip pdm~=${PDM_VERSION}; \
     \
     # Creates a non-root user with an explicit UID and adds permission to access the /app folder
     # For more info, please refer to https://aka.ms/vscode-docker-python-configure-containers

docker/sample.essentials.env

@@ -1,7 +1,9 @@
-
+# Refer https://hub.docker.com/_/postgres#:~:text=How%20to%20extend%20this%20image
 POSTGRES_USER=unstract_dev
 POSTGRES_PASSWORD=unstract_pass
 POSTGRES_DB=unstract_db
+# Used by db setup script
+POSTGRES_SCHEMA=unstract
 
 MINIO_ROOT_USER=minio
 MINIO_ROOT_PASSWORD=minio123

docker/scripts/db-setup/README.md

@@ -0,0 +1,9 @@
+# Unstract DB Setup Script
+
+[The db_setup.sh](/docker/scripts/db-setup/db_setup.sh) script helps setup the postgres database by making use of environment variables defined in the `.essentials.env` (user copy of the [sample.essentials.env](/docker/sample.essentials.env))
+
+- POSTGRES_USER
+- POSTGRES_DB
+- POSTGRES_SCHEMA
+
+This script helps setup the DB user and creates a new schema as well.

docker/scripts/db-setup/db_setup.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+echo "Creating DB '$POSTGRES_DB' and schema '$POSTGRES_SCHEMA' with user '$POSTGRES_USER'"
+
+psql -U ${POSTGRES_USER} -d ${POSTGRES_DB}<<-END
+    ALTER ROLE ${POSTGRES_USER} SET client_encoding TO 'utf8';
+    ALTER ROLE ${POSTGRES_USER} SET default_transaction_isolation TO 'read committed';
+    ALTER ROLE ${POSTGRES_USER} SET timezone TO 'UTC';
+    ALTER USER ${POSTGRES_USER} CREATEDB;
+    CREATE DATABASE ${POSTGRES_DB};
+    GRANT ALL PRIVILEGES ON DATABASE ${POSTGRES_DB} TO ${POSTGRES_USER};
+    CREATE SCHEMA IF NOT EXISTS ${POSTGRES_SCHEMA}
+END

frontend/src/components/custom-tools/document-parser/DocumentParser.jsx

@@ -30,9 +30,15 @@ function DocumentParser({
 }) {
   const [enforceTypeList, setEnforceTypeList] = useState([]);
   const [updatedPromptsCopy, setUpdatedPromptsCopy] = useState({});
+  const [isChallenge, setIsChallenge] = useState(false);
   const bottomRef = useRef(null);
-  const { details, isSimplePromptStudio, updateCustomTool, getDropdownItems } =
-    useCustomToolStore();
+  const {
+    details,
+    isSimplePromptStudio,
+    updateCustomTool,
+    getDropdownItems,
+    isChallengeEnabled,
+  } = useCustomToolStore();
   const { sessionDetails } = useSessionStore();
   const { setAlertDetails } = useAlertStore();
   const axiosPrivate = useAxiosPrivate();
@@ -45,6 +51,7 @@ function DocumentParser({
       return { value: outputTypeData[item] };
     });
     setEnforceTypeList(dropdownList1);
+    setIsChallenge(isChallengeEnabled);
 
     return () => {
       // Set the prompts with updated changes when the component is unmounted
@@ -63,6 +70,10 @@ function DocumentParser({
     };
   }, []);
 
+  useEffect(() => {
+    setIsChallenge(details.enable_challenge);
+  }, [details.enable_challenge]);
+
   useEffect(() => {
     if (scrollToBottom) {
       // Scroll down to the lastest chat.
@@ -169,6 +180,29 @@ function DocumentParser({
     return outputs;
   };
 
+  const getPromptCoverageCount = (promptId) => {
+    const keys = Object.keys(promptOutputs || {});
+    const coverageKey = `coverage_${promptId}`;
+    const outputs = {};
+    if (!keys?.length) {
+      details?.prompts?.forEach((prompt) => {
+        if (prompt?.coverage) {
+          const key = Object.keys(prompt?.coverage)[0];
+          if (key?.startsWith(coverageKey)) {
+            outputs[key] = prompt?.coverage[key];
+          }
+        }
+      });
+      return outputs;
+    }
+    keys?.forEach((key) => {
+      if (key?.startsWith(coverageKey)) {
+        outputs[key] = promptOutputs[key];
+      }
+    });
+    return outputs;
+  };
+
   if (!details?.prompts?.length) {
     if (isSimplePromptStudio && SpsPromptsEmptyState) {
       return <SpsPromptsEmptyState />;
@@ -196,6 +230,8 @@ function DocumentParser({
               outputs={getPromptOutputs(item?.prompt_id)}
               enforceTypeList={enforceTypeList}
               setUpdatedPromptsCopy={setUpdatedPromptsCopy}
+              coverageCountData={getPromptCoverageCount(item?.prompt_id)}
+              isChallenge={isChallenge}
             />
             <div ref={bottomRef} className="doc-parser-pad-bottom" />
           </div>

frontend/src/components/custom-tools/output-for-doc-modal/OutputForDocModal.jsx

@@ -21,7 +21,6 @@ import {
 import { SpinnerLoader } from "../../widgets/spinner-loader/SpinnerLoader";
 import { useAlertStore } from "../../../store/alert-store";
 import { useExceptionHandler } from "../../../hooks/useExceptionHandler";
-import { TokenUsage } from "../token-usage/TokenUsage";
 import { useTokenUsageStore } from "../../../store/token-usage-store";
 import { ProfileInfoBar } from "../profile-info-bar/ProfileInfoBar";
 
@@ -42,7 +41,7 @@ const outputStatus = {
   fail: "FAIL",
 };
 
-const errorTypes = ["null", "undefined", "false"];
+const errorTypes = [null, undefined, false];
 
 function OutputForDocModal({
   open,
@@ -243,6 +242,7 @@ function OutputForDocModal({
         status = outputStatus.success;
         message = displayPromptResult(output?.output, true);
       }
+      const promptTokenUsage = output?.token_usage?.total_tokens;
 
       if (output?.output === undefined) {
         status = outputStatus.yet_to_process;
@@ -253,17 +253,7 @@ function OutputForDocModal({
       const result = {
         key: item?.document_id,
         document: item?.document_name,
-        token_count: !singlePassExtractMode && (
-          <TokenUsage
-            tokenUsageId={
-              promptId +
-              "__" +
-              item?.document_id +
-              "__" +
-              (selectedProfile || profileManagerId)
-            }
-          />
-        ),
+        token_count: !singlePassExtractMode && (promptTokenUsage || "NA"),
         value: (
           <>
             {isLoading ? (