aiqtoolkit 1.1.0__py3-none-any.whl

aiq/eval/tunable_rag_evaluator/register.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import Field
+
+from aiq.builder.builder import EvalBuilder
+from aiq.builder.evaluator import EvaluatorInfo
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_evaluator
+from aiq.data_models.component_ref import LLMRef
+from aiq.data_models.evaluator import EvaluatorBaseConfig
+
+
+class TunableRagEvaluatorConfig(EvaluatorBaseConfig, name="tunable_rag_evaluator"):
+    '''Configuration for tunable RAG evaluator'''
+    llm_name: LLMRef = Field(description="Name of the judge LLM")
+    judge_llm_prompt: str = Field(description="LLM prompt for the judge LLM")
+    default_scoring: bool = Field(description="Whether to use default scoring", default=False)
+    default_score_weights: dict = Field(
+        default={
+            "coverage": 0.5, "correctness": 0.3, "relevance": 0.2
+        },
+        description="Weights for the different scoring components when using default scoring")
+
+
+@register_evaluator(config_type=TunableRagEvaluatorConfig)
+async def register_tunable_rag_evaluator(config: TunableRagEvaluatorConfig, builder: EvalBuilder):
+    '''Register tunable RAG evaluator'''
+    from .evaluate import TunableRagEvaluator
+
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    evaluator = TunableRagEvaluator(llm,
+                                    config.judge_llm_prompt,
+                                    builder.get_max_concurrency(),
+                                    config.default_scoring,
+                                    config.default_score_weights)
+
+    yield EvaluatorInfo(config=config, evaluate_fn=evaluator.evaluate, description="Tunable RAG Evaluator")

aiq/eval/utils/output_uploader.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import aioboto3
+from botocore.exceptions import NoCredentialsError
+from tqdm import tqdm
+
+from aiq.data_models.evaluate import EvalOutputConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OutputUploader:
+    """
+    Run custom scripts and upload evaluation outputs using the configured s3
+    credentials.
+    """
+
+    def __init__(self, output_config: EvalOutputConfig, job_id: str | None = None):
+        self.output_config = output_config
+        self._s3_client = None
+        self.job_id = job_id
+
+    @property
+    def s3_config(self):
+        return self.output_config.s3
+
+    async def _upload_file(self, s3_client, bucket, s3_key, local_path, pbar):
+        try:
+            await s3_client.upload_file(str(local_path), bucket, s3_key)
+            logger.info("Uploaded %s to s3://%s/%s", local_path, bucket, s3_key)
+            pbar.update(1)
+        except Exception as e:
+            logger.error("Failed to upload %s to s3://%s/%s: %s", local_path, bucket, s3_key, e)
+            raise
+
+    async def upload_directory(self):
+        """
+        Upload the contents of the local output directory to the remote S3 bucket in parallel.
+        """
+        if not self.output_config.s3:
+            logger.info("No S3 config provided; skipping upload.")
+            return
+
+        local_dir = self.output_config.dir
+        bucket = self.s3_config.bucket
+        remote_prefix = self.output_config.remote_dir or ""
+        if self.job_id:
+            remote_prefix = str(Path(remote_prefix) / f"jobs/{self.job_id}")
+
+        file_entries = []
+        for root, _, files in os.walk(local_dir):
+            for file in files:
+                local_path = Path(root) / file
+                relative_path = local_path.relative_to(local_dir)
+                s3_path = Path(remote_prefix) / relative_path
+                s3_key = str(s3_path).replace("\\", "/")  # Normalize for S3
+                file_entries.append((local_path, s3_key))
+
+        session = aioboto3.Session()
+        try:
+            async with session.client(
+                    "s3",
+                    endpoint_url=self.s3_config.endpoint_url,
+                    aws_access_key_id=self.s3_config.access_key,
+                    aws_secret_access_key=self.s3_config.secret_key,
+            ) as s3_client:
+                with tqdm(total=len(file_entries), desc="Uploading files to S3") as pbar:
+                    upload_tasks = [
+                        self._upload_file(s3_client, bucket, s3_key, local_path, pbar)
+                        for local_path, s3_key in file_entries
+                    ]
+                    await asyncio.gather(*upload_tasks)
+
+        except NoCredentialsError as e:
+            logger.error("AWS credentials not available: %s", e)
+            raise
+        except Exception as e:
+            logger.error("Failed to upload files to S3: %s", e)
+            raise
+
+    def run_custom_scripts(self):
+        """
+        Run custom Python scripts defined in the EvalOutputConfig.
+        Each script is run with its kwargs passed as command-line arguments.
+        The output directory is passed as the first argument.
+        """
+        for _, script_config in self.output_config.custom_scripts.items():
+            script_path = script_config.script
+            if not script_path.exists():
+                logger.error("Custom script %s does not exist.", script_path)
+                continue
+
+            # use python interpreter
+            args = [sys.executable, str(script_path)]
+            # add output directory as first keyword argument
+            args.append("--output_dir")
+            args.append(str(self.output_config.dir))
+            if script_config.kwargs:
+                for key, value in script_config.kwargs.items():
+                    args.append(f"--{key}")
+                    args.append(str(value))
+
+            display_args = " ".join(f'"{arg}"' if " " in arg else arg for arg in args[1:])
+
+            try:
+                logger.info("Running custom script: %s %s", script_path, display_args)
+                subprocess.run(args, check=True, text=True)
+                logger.info("Custom script %s completed successfully.", script_path)
+            except subprocess.CalledProcessError as e:
+                logger.error("Custom script %s failed with return code %s", script_path, e.returncode)
+                raise

aiq/eval/utils/tqdm_position_registry.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class TqdmPositionRegistry:
+    """
+    A simple registry for tqdm positions.
+    """
+    _positions = set()
+    _max_positions = 100
+
+    @classmethod
+    def claim(cls) -> int:
+        """
+        Claim a tqdm position in the range of 0-99.
+        """
+        for i in range(cls._max_positions):
+            if i not in cls._positions:
+                cls._positions.add(i)
+                return i
+        raise RuntimeError("No available tqdm positions.")
+
+    @classmethod
+    def release(cls, pos: int):
+        """
+        Release a tqdm position.
+        """
+        cls._positions.discard(pos)

aiq/front_ends/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/front_ends/console/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/front_ends/console/console_front_end_config.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from pydantic import Field
+
+from aiq.data_models.front_end import FrontEndBaseConfig
+
+
+class ConsoleFrontEndConfig(FrontEndBaseConfig, name="console"):
+    """
+    A front end that allows an AIQ Toolkit workflow to be run from the console.
+    """
+
+    input_query: list[str] | None = Field(default=None,
+                                          alias="input",
+                                          description="A single input to submit the the workflow.")
+    input_file: Path | None = Field(default=None,
+                                    description="Path to a json file of inputs to submit to the workflow.")

aiq/front_ends/console/console_front_end_plugin.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from io import StringIO
+
+import click
+from colorama import Fore
+
+from aiq.builder.workflow_builder import WorkflowBuilder
+from aiq.data_models.interactive import HumanPromptModelType
+from aiq.data_models.interactive import HumanResponse
+from aiq.data_models.interactive import HumanResponseText
+from aiq.data_models.interactive import InteractionPrompt
+from aiq.front_ends.console.console_front_end_config import ConsoleFrontEndConfig
+from aiq.front_ends.simple_base.simple_front_end_plugin_base import SimpleFrontEndPluginBase
+from aiq.runtime.session import AIQSessionManager
+
+logger = logging.getLogger(__name__)
+
+
+async def prompt_for_input_cli(question: InteractionPrompt) -> HumanResponse:
+    """
+    A simple CLI-based callback.
+    Takes question as str, returns the typed line as str.
+    """
+
+    if question.content.input_type == HumanPromptModelType.TEXT:
+        user_response = click.prompt(text=question.content.text)
+
+        return HumanResponseText(text=user_response)
+
+    raise ValueError("Unsupported human propmt input type. The run command only supports the 'HumanPromptText' "
+                     "input type. Please use the 'serve' command to ensure full support for all input types.")
+
+
+class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
+
+    async def pre_run(self):
+
+        if (not self.front_end_config.input_query and not self.front_end_config.input_file):
+            raise click.UsageError("Must specify either --input_query or --input_file")
+
+    async def run(self):
+
+        # Must yield the workflow function otherwise it cleans up
+        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+
+            session_manager: AIQSessionManager = None
+
+            if logger.isEnabledFor(logging.INFO):
+                stream = StringIO()
+
+                self.full_config.print_summary(stream=stream)
+
+                click.echo(stream.getvalue())
+
+                workflow = builder.build()
+                session_manager = AIQSessionManager(workflow)
+
+            await self.run_workflow(session_manager)
+
+    async def run_workflow(self, session_manager: AIQSessionManager = None):
+
+        runner_outputs = None
+
+        if (self.front_end_config.input_query):
+
+            async def run_single_query(query):
+
+                async with session_manager.session(user_input_callback=prompt_for_input_cli) as session:
+                    async with session.run(query) as runner:
+                        base_output = await runner.result(to_type=str)
+
+                        return base_output
+
+            # Convert to a list
+            input_list = list(self.front_end_config.input_query)
+            logger.debug("Processing input: %s", self.front_end_config.input_query)
+
+            runner_outputs = await asyncio.gather(*[run_single_query(query) for query in input_list])
+
+        elif (self.front_end_config.input_file):
+
+            # Run the workflow
+            with open(self.front_end_config.input_file, "r", encoding="utf-8") as f:
+
+                async with session_manager.workflow.run(f) as runner:
+                    runner_outputs = await runner.result(to_type=str)
+        else:
+            assert False, "Should not reach here. Should have been caught by pre_run"
+
+        # Print result
+        logger.info(f"\n{'-' * 50}\n{Fore.GREEN}Workflow Result:\n%s{Fore.RESET}\n{'-' * 50}", runner_outputs)

aiq/front_ends/console/register.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from aiq.cli.register_workflow import register_front_end
+from aiq.data_models.config import AIQConfig
+from aiq.front_ends.console.console_front_end_config import ConsoleFrontEndConfig
+
+
+@register_front_end(config_type=ConsoleFrontEndConfig)
+async def register_fastapi_front_end(config: ConsoleFrontEndConfig, full_config: AIQConfig):
+    from aiq.front_ends.console.console_front_end_plugin import ConsoleFrontEndPlugin
+
+    yield ConsoleFrontEndPlugin(full_config=full_config)

aiq/front_ends/cron/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/front_ends/fastapi/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/front_ends/fastapi/fastapi_front_end_config.py

@@ -0,0 +1,150 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import typing
+from datetime import datetime
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from aiq.data_models.front_end import FrontEndBaseConfig
+from aiq.data_models.step_adaptor import StepAdaptorConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AIQEvaluateRequest(BaseModel):
+    """Request model for the evaluate endpoint."""
+    config_file: str = Field(description="Path to the configuration file for evaluation")
+    job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job")
+    reps: int = Field(default=1, description="Number of repetitions for the evaluation, defaults to 1")
+    expiry_seconds: int = Field(
+        default=3600,
+        description="Optional time (in seconds) before the job expires. Clamped between 600 (10 min) and 86400 (24h).")
+
+
+class AIQEvaluateResponse(BaseModel):
+    """Response model for the evaluate endpoint."""
+    job_id: str = Field(description="Unique identifier for the evaluation job")
+    status: str = Field(description="Current status of the evaluation job")
+
+
+class AIQEvaluateStatusResponse(BaseModel):
+    """Response model for the evaluate status endpoint."""
+    job_id: str = Field(description="Unique identifier for the evaluation job")
+    status: str = Field(description="Current status of the evaluation job")
+    config_file: str = Field(description="Path to the configuration file used for evaluation")
+    error: str | None = Field(default=None, description="Error message if the job failed")
+    output_path: str | None = Field(default=None,
+                                    description="Path to the output file if the job completed successfully")
+    created_at: datetime = Field(description="Timestamp when the job was created")
+    updated_at: datetime = Field(description="Timestamp when the job was last updated")
+    expires_at: datetime | None = Field(default=None, description="Timestamp when the job will expire")
+
+
+class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
+    """
+    A FastAPI based front end that allows an AIQ Toolkit workflow to be served as a microservice.
+    """
+
+    class EndpointBase(BaseModel):
+
+        method: typing.Literal["GET", "POST", "PUT", "DELETE"]
+        description: str
+        path: str | None = Field(
+            default=None,
+            description=("Path for the default workflow. If None, no workflow endpoint is created."),
+        )
+        websocket_path: str | None = Field(
+            default=None,
+            description=("Path for the websocket. If None, no websocket is created."),
+        )
+        openai_api_path: str | None = Field(
+            default=None,
+            description=("Path for the default workflow using the OpenAI API Specification. "
+                         "If None, no workflow endpoint with the OpenAI API Specification is created."),
+        )
+
+    class Endpoint(EndpointBase):
+        function_name: str = Field(description="The name of the function to call for this endpoint")
+
+    class CrossOriginResourceSharing(BaseModel):
+        allow_origins: list[str] | None = Field(
+            default=None, description=" A list of origins that should be permitted to make cross-origin requests.")
+        allow_origin_regex: str | None = Field(
+            default=None,
+            description="A permitted regex string to match against origins to make cross-origin requests",
+        )
+        allow_methods: list[str] | None = Field(
+            default_factory=lambda: ['GET'],
+            description="A list of HTTP methods that should be allowed for cross-origin requests.")
+        allow_headers: list[str] | None = Field(
+            default_factory=list,
+            description="A list of HTTP request headers that should be supported for cross-origin requests.")
+        allow_credentials: bool | None = Field(
+            default=False,
+            description="Indicate that cookies should be supported for cross-origin requests.",
+        )
+        expose_headers: list[str] | None = Field(
+            default_factory=list,
+            description="Indicate any response headers that should be made accessible to the browser.",
+        )
+        max_age: int | None = Field(
+            default=600,
+            description="Sets a maximum time in seconds for browsers to cache CORS responses.",
+        )
+
+    root_path: str = Field(default="", description="The root path for the API")
+    host: str = Field(default="localhost", description="Host to bind the server to")
+    port: int = Field(default=8000, description="Port to bind the server to", ge=0, le=65535)
+    reload: bool = Field(default=False, description="Enable auto-reload for development")
+    workers: int = Field(default=1, description="Number of workers to run", ge=1)
+    step_adaptor: StepAdaptorConfig = StepAdaptorConfig()
+
+    workflow: typing.Annotated[EndpointBase, Field(description="Endpoint for the default workflow.")] = EndpointBase(
+        method="POST",
+        path="/generate",
+        websocket_path="/websocket",
+        openai_api_path="/chat",
+        description="Executes the default AIQ Toolkit workflow from the loaded configuration ",
+    )
+
+    evaluate: typing.Annotated[EndpointBase, Field(description="Endpoint for evaluating workflows.")] = EndpointBase(
+        method="POST",
+        path="/evaluate",
+        description="Evaluates the performance and accuracy of the workflow on a dataset",
+    )
+
+    endpoints: list[Endpoint] = Field(
+        default_factory=list,
+        description=(
+            "Additional endpoints to add to the FastAPI app which run functions within the AIQ Toolkit configuration. "
+            "Each endpoint must have a unique path."))
+
+    cors: CrossOriginResourceSharing = Field(
+        default_factory=CrossOriginResourceSharing,
+        description="Cross origin resource sharing configuration for the FastAPI app")
+
+    use_gunicorn: bool = Field(
+        default=False,
+        description="Use Gunicorn to run the FastAPI app",
+    )
+    runner_class: str | None = Field(
+        default=None,
+        description=("The AIQ Toolkit runner class to use when launching the FastAPI app from multiple processes. "
+                     "Each runner is responsible for loading and running the AIQ Toolkit workflow. "
+                     "Note: This is different from the worker class used by Gunicorn."),
+    )