aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of aiqtoolkit might be problematic. Click here for more details.
- aiq/agent/base.py +170 -8
- aiq/agent/dual_node.py +1 -1
- aiq/agent/react_agent/agent.py +146 -112
- aiq/agent/react_agent/prompt.py +1 -6
- aiq/agent/react_agent/register.py +36 -35
- aiq/agent/rewoo_agent/agent.py +36 -35
- aiq/agent/rewoo_agent/register.py +2 -2
- aiq/agent/tool_calling_agent/agent.py +3 -7
- aiq/agent/tool_calling_agent/register.py +1 -1
- aiq/authentication/__init__.py +14 -0
- aiq/authentication/api_key/__init__.py +14 -0
- aiq/authentication/api_key/api_key_auth_provider.py +92 -0
- aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
- aiq/authentication/api_key/register.py +26 -0
- aiq/authentication/exceptions/__init__.py +14 -0
- aiq/authentication/exceptions/api_key_exceptions.py +38 -0
- aiq/authentication/exceptions/auth_code_grant_exceptions.py +86 -0
- aiq/authentication/exceptions/call_back_exceptions.py +38 -0
- aiq/authentication/exceptions/request_exceptions.py +54 -0
- aiq/authentication/http_basic_auth/__init__.py +0 -0
- aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
- aiq/authentication/http_basic_auth/register.py +30 -0
- aiq/authentication/interfaces.py +93 -0
- aiq/authentication/oauth2/__init__.py +14 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
- aiq/authentication/oauth2/register.py +25 -0
- aiq/authentication/register.py +21 -0
- aiq/builder/builder.py +64 -2
- aiq/builder/component_utils.py +16 -3
- aiq/builder/context.py +37 -0
- aiq/builder/eval_builder.py +43 -2
- aiq/builder/function.py +44 -12
- aiq/builder/function_base.py +1 -1
- aiq/builder/intermediate_step_manager.py +6 -8
- aiq/builder/user_interaction_manager.py +3 -0
- aiq/builder/workflow.py +23 -18
- aiq/builder/workflow_builder.py +421 -61
- aiq/cli/commands/info/list_mcp.py +103 -16
- aiq/cli/commands/sizing/__init__.py +14 -0
- aiq/cli/commands/sizing/calc.py +294 -0
- aiq/cli/commands/sizing/sizing.py +27 -0
- aiq/cli/commands/start.py +2 -1
- aiq/cli/entrypoint.py +2 -0
- aiq/cli/register_workflow.py +80 -0
- aiq/cli/type_registry.py +151 -30
- aiq/data_models/api_server.py +124 -12
- aiq/data_models/authentication.py +231 -0
- aiq/data_models/common.py +35 -7
- aiq/data_models/component.py +17 -9
- aiq/data_models/component_ref.py +33 -0
- aiq/data_models/config.py +60 -3
- aiq/data_models/dataset_handler.py +2 -1
- aiq/data_models/embedder.py +1 -0
- aiq/data_models/evaluate.py +23 -0
- aiq/data_models/function_dependencies.py +8 -0
- aiq/data_models/interactive.py +10 -1
- aiq/data_models/intermediate_step.py +38 -5
- aiq/data_models/its_strategy.py +30 -0
- aiq/data_models/llm.py +1 -0
- aiq/data_models/memory.py +1 -0
- aiq/data_models/object_store.py +44 -0
- aiq/data_models/profiler.py +1 -0
- aiq/data_models/retry_mixin.py +35 -0
- aiq/data_models/span.py +187 -0
- aiq/data_models/telemetry_exporter.py +2 -2
- aiq/embedder/nim_embedder.py +2 -1
- aiq/embedder/openai_embedder.py +2 -1
- aiq/eval/config.py +19 -1
- aiq/eval/dataset_handler/dataset_handler.py +87 -2
- aiq/eval/evaluate.py +208 -27
- aiq/eval/evaluator/base_evaluator.py +73 -0
- aiq/eval/evaluator/evaluator_model.py +1 -0
- aiq/eval/intermediate_step_adapter.py +11 -5
- aiq/eval/rag_evaluator/evaluate.py +55 -15
- aiq/eval/rag_evaluator/register.py +6 -1
- aiq/eval/remote_workflow.py +7 -2
- aiq/eval/runners/__init__.py +14 -0
- aiq/eval/runners/config.py +39 -0
- aiq/eval/runners/multi_eval_runner.py +54 -0
- aiq/eval/trajectory_evaluator/evaluate.py +22 -65
- aiq/eval/tunable_rag_evaluator/evaluate.py +150 -168
- aiq/eval/tunable_rag_evaluator/register.py +2 -0
- aiq/eval/usage_stats.py +41 -0
- aiq/eval/utils/output_uploader.py +10 -1
- aiq/eval/utils/weave_eval.py +184 -0
- aiq/experimental/__init__.py +0 -0
- aiq/experimental/decorators/__init__.py +0 -0
- aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
- aiq/experimental/inference_time_scaling/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/editing/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/editing/iterative_plan_refinement_editor.py +147 -0
- aiq/experimental/inference_time_scaling/editing/llm_as_a_judge_editor.py +204 -0
- aiq/experimental/inference_time_scaling/editing/motivation_aware_summarization.py +107 -0
- aiq/experimental/inference_time_scaling/functions/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/functions/execute_score_select_function.py +105 -0
- aiq/experimental/inference_time_scaling/functions/its_tool_orchestration_function.py +205 -0
- aiq/experimental/inference_time_scaling/functions/its_tool_wrapper_function.py +146 -0
- aiq/experimental/inference_time_scaling/functions/plan_select_execute_function.py +224 -0
- aiq/experimental/inference_time_scaling/models/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/models/editor_config.py +132 -0
- aiq/experimental/inference_time_scaling/models/its_item.py +48 -0
- aiq/experimental/inference_time_scaling/models/scoring_config.py +112 -0
- aiq/experimental/inference_time_scaling/models/search_config.py +120 -0
- aiq/experimental/inference_time_scaling/models/selection_config.py +154 -0
- aiq/experimental/inference_time_scaling/models/stage_enums.py +43 -0
- aiq/experimental/inference_time_scaling/models/strategy_base.py +66 -0
- aiq/experimental/inference_time_scaling/models/tool_use_config.py +41 -0
- aiq/experimental/inference_time_scaling/register.py +36 -0
- aiq/experimental/inference_time_scaling/scoring/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/scoring/llm_based_agent_scorer.py +168 -0
- aiq/experimental/inference_time_scaling/scoring/llm_based_plan_scorer.py +168 -0
- aiq/experimental/inference_time_scaling/scoring/motivation_aware_scorer.py +111 -0
- aiq/experimental/inference_time_scaling/search/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/search/multi_llm_planner.py +128 -0
- aiq/experimental/inference_time_scaling/search/multi_query_retrieval_search.py +122 -0
- aiq/experimental/inference_time_scaling/search/single_shot_multi_plan_planner.py +128 -0
- aiq/experimental/inference_time_scaling/selection/__init__.py +0 -0
- aiq/experimental/inference_time_scaling/selection/best_of_n_selector.py +63 -0
- aiq/experimental/inference_time_scaling/selection/llm_based_agent_output_selector.py +131 -0
- aiq/experimental/inference_time_scaling/selection/llm_based_output_merging_selector.py +159 -0
- aiq/experimental/inference_time_scaling/selection/llm_based_plan_selector.py +128 -0
- aiq/experimental/inference_time_scaling/selection/threshold_selector.py +58 -0
- aiq/front_ends/console/authentication_flow_handler.py +233 -0
- aiq/front_ends/console/console_front_end_plugin.py +11 -2
- aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
- aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
- aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
- aiq/front_ends/fastapi/fastapi_front_end_config.py +93 -9
- aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
- aiq/front_ends/fastapi/fastapi_front_end_plugin.py +14 -1
- aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +537 -52
- aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
- aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
- aiq/front_ends/fastapi/job_store.py +47 -25
- aiq/front_ends/fastapi/main.py +2 -0
- aiq/front_ends/fastapi/message_handler.py +108 -89
- aiq/front_ends/fastapi/step_adaptor.py +2 -1
- aiq/llm/aws_bedrock_llm.py +57 -0
- aiq/llm/nim_llm.py +2 -1
- aiq/llm/openai_llm.py +3 -2
- aiq/llm/register.py +1 -0
- aiq/meta/pypi.md +12 -12
- aiq/object_store/__init__.py +20 -0
- aiq/object_store/in_memory_object_store.py +74 -0
- aiq/object_store/interfaces.py +84 -0
- aiq/object_store/models.py +36 -0
- aiq/object_store/register.py +20 -0
- aiq/observability/__init__.py +14 -0
- aiq/observability/exporter/__init__.py +14 -0
- aiq/observability/exporter/base_exporter.py +449 -0
- aiq/observability/exporter/exporter.py +78 -0
- aiq/observability/exporter/file_exporter.py +33 -0
- aiq/observability/exporter/processing_exporter.py +269 -0
- aiq/observability/exporter/raw_exporter.py +52 -0
- aiq/observability/exporter/span_exporter.py +264 -0
- aiq/observability/exporter_manager.py +335 -0
- aiq/observability/mixin/__init__.py +14 -0
- aiq/observability/mixin/batch_config_mixin.py +26 -0
- aiq/observability/mixin/collector_config_mixin.py +23 -0
- aiq/observability/mixin/file_mixin.py +288 -0
- aiq/observability/mixin/file_mode.py +23 -0
- aiq/observability/mixin/resource_conflict_mixin.py +134 -0
- aiq/observability/mixin/serialize_mixin.py +61 -0
- aiq/observability/mixin/type_introspection_mixin.py +183 -0
- aiq/observability/processor/__init__.py +14 -0
- aiq/observability/processor/batching_processor.py +316 -0
- aiq/observability/processor/intermediate_step_serializer.py +28 -0
- aiq/observability/processor/processor.py +68 -0
- aiq/observability/register.py +36 -39
- aiq/observability/utils/__init__.py +14 -0
- aiq/observability/utils/dict_utils.py +236 -0
- aiq/observability/utils/time_utils.py +31 -0
- aiq/profiler/calc/__init__.py +14 -0
- aiq/profiler/calc/calc_runner.py +623 -0
- aiq/profiler/calc/calculations.py +288 -0
- aiq/profiler/calc/data_models.py +176 -0
- aiq/profiler/calc/plot.py +345 -0
- aiq/profiler/callbacks/langchain_callback_handler.py +22 -10
- aiq/profiler/data_models.py +24 -0
- aiq/profiler/inference_metrics_model.py +3 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +8 -0
- aiq/profiler/inference_optimization/data_models.py +2 -2
- aiq/profiler/inference_optimization/llm_metrics.py +2 -2
- aiq/profiler/profile_runner.py +61 -21
- aiq/runtime/loader.py +9 -3
- aiq/runtime/runner.py +23 -9
- aiq/runtime/session.py +25 -7
- aiq/runtime/user_metadata.py +2 -3
- aiq/tool/chat_completion.py +74 -0
- aiq/tool/code_execution/README.md +152 -0
- aiq/tool/code_execution/code_sandbox.py +151 -72
- aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
- aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +139 -24
- aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +3 -1
- aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +27 -2
- aiq/tool/code_execution/register.py +7 -3
- aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
- aiq/tool/mcp/exceptions.py +142 -0
- aiq/tool/mcp/mcp_client.py +41 -6
- aiq/tool/mcp/mcp_tool.py +3 -2
- aiq/tool/register.py +1 -0
- aiq/tool/server_tools.py +6 -3
- aiq/utils/exception_handlers/automatic_retries.py +289 -0
- aiq/utils/exception_handlers/mcp.py +211 -0
- aiq/utils/io/model_processing.py +28 -0
- aiq/utils/log_utils.py +37 -0
- aiq/utils/string_utils.py +38 -0
- aiq/utils/type_converter.py +18 -2
- aiq/utils/type_utils.py +87 -0
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/METADATA +53 -21
- aiqtoolkit-1.2.0rc1.dist-info/RECORD +436 -0
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/WHEEL +1 -1
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/entry_points.txt +3 -0
- aiq/front_ends/fastapi/websocket.py +0 -148
- aiq/observability/async_otel_listener.py +0 -429
- aiqtoolkit-1.2.0.dev0.dist-info/RECORD +0 -316
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/licenses/LICENSE.md +0 -0
- {aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import typing
|
|
17
|
+
|
|
18
|
+
from pydantic import Field
|
|
19
|
+
from pydantic import model_validator
|
|
20
|
+
|
|
21
|
+
from aiq.data_models.component_ref import LLMRef
|
|
22
|
+
from aiq.data_models.its_strategy import ITSStrategyBaseConfig
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class SingleShotMultiPlanConfig(ITSStrategyBaseConfig, name="single_shot_multi_plan"):
|
|
26
|
+
num_plans: int = Field(default=4, description="Number of plans to generate.")
|
|
27
|
+
max_temperature: float = Field(default=1.0,
|
|
28
|
+
description="Maximum temperature to use for sampling when generating plans. "
|
|
29
|
+
"This can help control the randomness of the generated plans.")
|
|
30
|
+
min_temperature: float = Field(default=0.5,
|
|
31
|
+
description="Minimum temperature to use for sampling when generating plans. "
|
|
32
|
+
"This can help control the randomness of the generated plans.")
|
|
33
|
+
# If strategy is provided, LLM must be
|
|
34
|
+
planning_llm: LLMRef | typing.Any | None = Field(
|
|
35
|
+
default=None,
|
|
36
|
+
description="The LLM to use for planning. This can be a callable or an "
|
|
37
|
+
"instance of an LLM client.")
|
|
38
|
+
|
|
39
|
+
planning_template: str = Field(
|
|
40
|
+
default=("You are an expert reasoning model task with creating a detailed execution plan"
|
|
41
|
+
" for a system that has the following information to get the result of a given input:\n\n"
|
|
42
|
+
"**System Information:**\n {context}"
|
|
43
|
+
"**Input:** \n{prompt}\n\n"
|
|
44
|
+
"An example plan could look like this:\n\n"
|
|
45
|
+
"1. Call tool A with input X\n"
|
|
46
|
+
"2. Call tool B with input Y\n"
|
|
47
|
+
"3. Interpret the output of tool A and B\n"
|
|
48
|
+
"4. Return the final result"
|
|
49
|
+
"\n\nBegin the final plan with PLAN:\n"),
|
|
50
|
+
description="The template to use for generating plans.")
|
|
51
|
+
|
|
52
|
+
@model_validator(mode="before")
|
|
53
|
+
def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
|
|
54
|
+
"""
|
|
55
|
+
Ensure that the required LLMs are provided based on the selected strategies.
|
|
56
|
+
"""
|
|
57
|
+
# Validate planning strategy: planning_llm must be provided if planning_strategy is set
|
|
58
|
+
if values.get('planning_llm') is None:
|
|
59
|
+
raise ValueError('planning_llm must be provided when planning_strategy is set.')
|
|
60
|
+
|
|
61
|
+
return values
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class MultiLLMPlanConfig(ITSStrategyBaseConfig, name="multi_llm_plan"):
|
|
65
|
+
"""Configuration for a 'multi LLM plan generation' strategy."""
|
|
66
|
+
llms: list[LLMRef] = Field(
|
|
67
|
+
default_factory=list,
|
|
68
|
+
description="list of LLMs to use for plan generation. Each LLM can generate one or more plans.")
|
|
69
|
+
plans_per_llm: int = Field(default=2, description="Number of plans each LLM should generate.")
|
|
70
|
+
max_temperature: float = Field(default=1.0,
|
|
71
|
+
description="Maximum temperature to use for sampling when generating plans. "
|
|
72
|
+
"This can help control the randomness of the generated plans.")
|
|
73
|
+
min_temperature: float = Field(default=0.5,
|
|
74
|
+
description="Minimum temperature to use for sampling when generating plans. "
|
|
75
|
+
"This can help control the randomness of the generated plans.")
|
|
76
|
+
planning_template: str = Field(
|
|
77
|
+
default=("You are an expert reasoning model task with creating a detailed execution plan"
|
|
78
|
+
" for a system that has the following information to get the result of a given input:\n\n"
|
|
79
|
+
"**System Information:**\n {context}"
|
|
80
|
+
"**Input:** \n{prompt}\n\n"
|
|
81
|
+
"An example plan could look like this:\n\n"
|
|
82
|
+
"1. Call tool A with input X\n"
|
|
83
|
+
"2. Call tool B with input Y\n"
|
|
84
|
+
"3. Interpret the output of tool A and B\n"
|
|
85
|
+
"4. Return the final result"
|
|
86
|
+
"\n\nBegin the final plan with PLAN:\n"),
|
|
87
|
+
description="The template to use for generating plans.")
|
|
88
|
+
|
|
89
|
+
@model_validator(mode="before")
|
|
90
|
+
def validate_multi_llm_strategies(cls, values: dict) -> dict:
|
|
91
|
+
if not values.get('llms'):
|
|
92
|
+
raise ValueError('Must provide at least one LLMRef in `llms` for multi-LLM strategy.')
|
|
93
|
+
return values
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class MultiQueryRetrievalSearchConfig(ITSStrategyBaseConfig, name="multi_query_retrieval_search"):
|
|
97
|
+
"""
|
|
98
|
+
Configuration for the MultiQueryRetrievalSearch strategy.
|
|
99
|
+
This strategy generates multiple new 'ITSItem's per original item,
|
|
100
|
+
each containing a differently phrased or re-focused version of the original task.
|
|
101
|
+
"""
|
|
102
|
+
llms: list[LLMRef] = Field(default_factory=list,
|
|
103
|
+
description="list of LLM references to use for generating diverse queries.")
|
|
104
|
+
|
|
105
|
+
query_generation_template: str = Field(
|
|
106
|
+
default=("You are an expert at re-framing a user's query to encourage new solution paths. "
|
|
107
|
+
"Given the task description and an optional motivation, produce a short alternative query "
|
|
108
|
+
"that addresses the same task from a different angle. By generating multiple "
|
|
109
|
+
"perspectives on the task, your goal is to help "
|
|
110
|
+
"the user overcome some of the limitations of distance-based similarity search.\n\n"
|
|
111
|
+
"Task: {task}\n"
|
|
112
|
+
"Motivation: {motivation}\n\n"
|
|
113
|
+
"Output a concise new query statement below. Only output the revised query and nothing else.\n"),
|
|
114
|
+
description="Prompt template for rewriting the task from a different perspective.")
|
|
115
|
+
|
|
116
|
+
@model_validator(mode="before")
|
|
117
|
+
def validate_llms(cls, values):
|
|
118
|
+
if not values.get('llms'):
|
|
119
|
+
raise ValueError("At least one LLMRef must be provided for multi_query_retrieval_search.")
|
|
120
|
+
return values
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License..
|
|
15
|
+
|
|
16
|
+
import typing
|
|
17
|
+
|
|
18
|
+
from pydantic import Field
|
|
19
|
+
from pydantic import model_validator
|
|
20
|
+
|
|
21
|
+
from aiq.data_models.component_ref import LLMRef
|
|
22
|
+
from aiq.data_models.its_strategy import ITSStrategyBaseConfig
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class LLMBasedPlanSelectionConfig(ITSStrategyBaseConfig, name="llm_based_plan_selection"):
|
|
26
|
+
"""
|
|
27
|
+
Configuration for LLMBasedSelection.
|
|
28
|
+
"""
|
|
29
|
+
selection_llm: LLMRef | typing.Any | None = Field(
|
|
30
|
+
default=None,
|
|
31
|
+
description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
|
|
32
|
+
|
|
33
|
+
selection_template: str = Field(
|
|
34
|
+
default=("You are tasked with selecting the best plan from several alternative plans."
|
|
35
|
+
" Review the following plans and their feedback carefully to select the most "
|
|
36
|
+
"comprehensive, efficient, and effective one."
|
|
37
|
+
"The plan is for an agent system with the following objective and context:\n\n"
|
|
38
|
+
"{context}\n\n"
|
|
39
|
+
"The system is asked to achieve the following goal:\n\n"
|
|
40
|
+
"{original_prompt}\n\n"
|
|
41
|
+
"The generated plans are as follows."
|
|
42
|
+
"\n\n{plans}"
|
|
43
|
+
"\n\nBased on your analysis, which plan (numbered 1 and onwards) is the best? "
|
|
44
|
+
"Provide a thorough explanation of your choice,"
|
|
45
|
+
" referencing specific strengths from the feedback and how they outweigh any weaknesses."
|
|
46
|
+
"Make sure you begin your choice of selected plan with the words 'SELECTED PLAN:' "
|
|
47
|
+
"followed by the plan number."),
|
|
48
|
+
description="The template to use for selecting the best plan. This should guide the LLM on how to evaluate "
|
|
49
|
+
"the plans and select the best one. Ensure it is clear and concise.")
|
|
50
|
+
|
|
51
|
+
@model_validator(mode="before")
|
|
52
|
+
def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
|
|
53
|
+
"""
|
|
54
|
+
Ensure that the selection_llm is provided when using LLMBasedSelection.
|
|
55
|
+
"""
|
|
56
|
+
if values.get('selection_llm') is None:
|
|
57
|
+
raise ValueError('selection_llm must be provided when'
|
|
58
|
+
' selection_strategy is set to LLM_BASED_PLAN_SELECTION.')
|
|
59
|
+
|
|
60
|
+
return values
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
class LLMBasedAgentOutputSelectionConfig(ITSStrategyBaseConfig, name="llm_based_agent_output_selection"):
|
|
64
|
+
"""
|
|
65
|
+
Configuration for LLMBasedSelection.
|
|
66
|
+
"""
|
|
67
|
+
selection_llm: LLMRef | typing.Any | None = Field(
|
|
68
|
+
default=None,
|
|
69
|
+
description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
|
|
70
|
+
|
|
71
|
+
selection_template: str = Field(
|
|
72
|
+
default=("You are tasked with selecting the best output from several output."
|
|
73
|
+
"The outputs are from an agent system whose object and input will be provided below.\n "
|
|
74
|
+
"Review all the outputs and select one that fits the best. You will do this by "
|
|
75
|
+
"looking at how many outputs have the same classification. Chose the one that has the most. "
|
|
76
|
+
"Of the ones that have the same classification, choose the one that is the most complete, "
|
|
77
|
+
"clear, and comprehensive. The objective of the agent is: \n"
|
|
78
|
+
"{objective}\n\n"
|
|
79
|
+
"\n\nThe agent is asked to achieve the following goal:\n\n"
|
|
80
|
+
"{input}\n\n"
|
|
81
|
+
"The generated outputs are as follows."
|
|
82
|
+
"\n\n{results}"
|
|
83
|
+
"\n\nBased on your analysis, which plan (numbered 1 and onwards) is the best? "
|
|
84
|
+
"Provide a thorough explanation of your choice,"
|
|
85
|
+
" referencing specific strengths from the feedback and how they outweigh any weaknesses."
|
|
86
|
+
"You must ALWAYS select an option, even if the options are identical or similar. "
|
|
87
|
+
"Make sure you begin your choice of selected plan with the words 'SELECTED ITEM:' "
|
|
88
|
+
"followed by the plan number."),
|
|
89
|
+
description="The template to use for selecting the best output. This should guide the LLM on how to evaluate "
|
|
90
|
+
"the outputs and select the best one. Ensure it is clear and concise. Must contain {objective}, "
|
|
91
|
+
"{input}, and {results} ")
|
|
92
|
+
|
|
93
|
+
@model_validator(mode="before")
|
|
94
|
+
def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
|
|
95
|
+
"""
|
|
96
|
+
Ensure that the selection_llm is provided when using LLMBasedSelection.
|
|
97
|
+
"""
|
|
98
|
+
if values.get('selection_llm') is None:
|
|
99
|
+
raise ValueError('selection_llm must be provided when '
|
|
100
|
+
'selection_strategy is set to LLM_BASED_AGENT_OUTPUT_SELECTION.')
|
|
101
|
+
|
|
102
|
+
return values
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
class LLMBasedOutputMergingConfig(ITSStrategyBaseConfig, name="llm_based_agent_output_merging"):
|
|
106
|
+
"""
|
|
107
|
+
Configuration for LLMBasedSelection.
|
|
108
|
+
"""
|
|
109
|
+
selection_llm: LLMRef | typing.Any | None = Field(
|
|
110
|
+
default=None,
|
|
111
|
+
description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
|
|
112
|
+
|
|
113
|
+
selection_template: str = Field(
|
|
114
|
+
default=("You are tasked with merging the output of an agent systems that produces {pipeline_type}."
|
|
115
|
+
"The outputs are from an agent system whose objective and input will be provided below.\n "
|
|
116
|
+
"Review all the outputs, please combine them all into one output, keeping with the intended structure "
|
|
117
|
+
"generated by the outputs and general tone. Capture the important pieces of each of the outputs "
|
|
118
|
+
"to create comprehensive output that achieves the input and objective. "
|
|
119
|
+
"The objective of the agent is: \n"
|
|
120
|
+
"{objective}\n\n"
|
|
121
|
+
"\n\nThe agent is asked to achieve the following goal:\n\n"
|
|
122
|
+
"{input}\n\n"
|
|
123
|
+
"The generated outputs are as follows."
|
|
124
|
+
"\n\n{results}"
|
|
125
|
+
"\n\n Make sure you begin your updated output with the words 'MERGED OUTPUT:' "),
|
|
126
|
+
description="The template to use for selecting the best output. This should guide the LLM on how to evaluate "
|
|
127
|
+
"the outputs and select the best one. Ensure it is clear and concise. Must contain {objective}, "
|
|
128
|
+
"{input}, and {results} ")
|
|
129
|
+
|
|
130
|
+
@model_validator(mode="before")
|
|
131
|
+
def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
|
|
132
|
+
"""
|
|
133
|
+
Ensure that the selection_llm is provided when using LLMBasedSelection.
|
|
134
|
+
"""
|
|
135
|
+
if values.get('selection_llm') is None:
|
|
136
|
+
raise ValueError('selection_llm must be provided when '
|
|
137
|
+
'selection_strategy is set to LLM_BASED_AGENT_OUTPUT_SELECTION.')
|
|
138
|
+
|
|
139
|
+
return values
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
class ThresholdSelectionConfig(ITSStrategyBaseConfig, name="threshold_selection"):
|
|
143
|
+
"""
|
|
144
|
+
Configuration for a selection strategy that keeps only the items
|
|
145
|
+
whose scores exceed a specified threshold.
|
|
146
|
+
"""
|
|
147
|
+
threshold: float = Field(default=5.0, description="Only keep ITSItems with score >= this value.")
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
class BestOfNSelectionConfig(ITSStrategyBaseConfig, name="best_of_n_selection"):
|
|
151
|
+
"""
|
|
152
|
+
Configuration for Best of N Selection
|
|
153
|
+
"""
|
|
154
|
+
pass
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from enum import Enum
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class PipelineTypeEnum(str, Enum):
|
|
20
|
+
"""
|
|
21
|
+
Enum to represent the type of pipeline used in Inference Time Scaling.
|
|
22
|
+
"""
|
|
23
|
+
PLANNING = "planning"
|
|
24
|
+
TOOL_USE = "tool_use"
|
|
25
|
+
AGENT_EXECUTION = "agent_execution"
|
|
26
|
+
CUSTOM = "custom"
|
|
27
|
+
|
|
28
|
+
def __str__(self) -> str:
|
|
29
|
+
return self.value
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class StageTypeEnum(str, Enum):
|
|
33
|
+
"""
|
|
34
|
+
Enum to represent the type of stage in a pipeline.
|
|
35
|
+
"""
|
|
36
|
+
SEARCH = "search"
|
|
37
|
+
EDITING = "editing"
|
|
38
|
+
SCORING = "scoring"
|
|
39
|
+
SELECTION = "selection"
|
|
40
|
+
CUSTOM = "custom"
|
|
41
|
+
|
|
42
|
+
def __str__(self) -> str:
|
|
43
|
+
return self.value
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from abc import ABC
|
|
17
|
+
from abc import abstractmethod
|
|
18
|
+
|
|
19
|
+
from aiq.builder.builder import Builder
|
|
20
|
+
from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
|
|
21
|
+
from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum, PipelineTypeEnum
|
|
22
|
+
from aiq.data_models.its_strategy import ITSStrategyBaseConfig
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class StrategyBase(ABC):
|
|
26
|
+
"""
|
|
27
|
+
Abstract base class for strategy implementations.
|
|
28
|
+
|
|
29
|
+
This class defines the interface for strategies that can be used in the
|
|
30
|
+
Inference Time Scaling (ITS) framework. Concrete strategy classes should
|
|
31
|
+
implement the methods defined in this class.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(self, config: ITSStrategyBaseConfig) -> None:
|
|
35
|
+
self.config: ITSStrategyBaseConfig = config
|
|
36
|
+
self.pipeline_type: PipelineTypeEnum | None = None
|
|
37
|
+
|
|
38
|
+
@abstractmethod
|
|
39
|
+
async def build_components(self, builder: Builder) -> None:
|
|
40
|
+
"""Build the components required for the selector."""
|
|
41
|
+
pass
|
|
42
|
+
|
|
43
|
+
@abstractmethod
|
|
44
|
+
async def ainvoke(self,
|
|
45
|
+
items: list[ITSItem],
|
|
46
|
+
original_prompt: str | None = None,
|
|
47
|
+
agent_context: str | None = None,
|
|
48
|
+
**kwargs) -> [ITSItem]:
|
|
49
|
+
pass
|
|
50
|
+
|
|
51
|
+
@abstractmethod
|
|
52
|
+
def supported_pipeline_types(self) -> [PipelineTypeEnum]:
|
|
53
|
+
"""Return the stage types supported by this selector."""
|
|
54
|
+
pass
|
|
55
|
+
|
|
56
|
+
@abstractmethod
|
|
57
|
+
def stage_type(self) -> StageTypeEnum:
|
|
58
|
+
"""Return the stage type of this strategy."""
|
|
59
|
+
pass
|
|
60
|
+
|
|
61
|
+
def set_pipeline_type(self, pipeline_type: PipelineTypeEnum) -> None:
|
|
62
|
+
"""Set the pipeline type for this strategy."""
|
|
63
|
+
if pipeline_type in self.supported_pipeline_types():
|
|
64
|
+
self.pipeline_type = pipeline_type
|
|
65
|
+
else:
|
|
66
|
+
raise ValueError(f"Pipeline type {pipeline_type} is not supported by this strategy.")
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from pydantic import BaseModel
|
|
17
|
+
from pydantic import Field
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class ToolUseInputSchema(BaseModel):
|
|
21
|
+
"""
|
|
22
|
+
Input schema for the tool use function.
|
|
23
|
+
"""
|
|
24
|
+
tool_name: str = Field(description="The name of the tool to use. Must be registered in the system.", )
|
|
25
|
+
task_description: str = Field(description="The description of the task to perform with the tool.", )
|
|
26
|
+
motivation: str | None = Field(
|
|
27
|
+
default=None,
|
|
28
|
+
description="An optional motivation for the tool use, providing additional context or reasoning.",
|
|
29
|
+
)
|
|
30
|
+
output: str | None = Field(
|
|
31
|
+
default=None,
|
|
32
|
+
description="The output of the tool use. This can be used to store the result of the tool execution.",
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class ToolUselist(BaseModel):
|
|
37
|
+
"""
|
|
38
|
+
A list of tools to use.
|
|
39
|
+
"""
|
|
40
|
+
tools: list[ToolUseInputSchema] = Field(
|
|
41
|
+
description="A list of tool use inputs, each containing the tool name and task description.", )
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
# pylint: disable=unused-import
|
|
17
|
+
# flake8: noqa
|
|
18
|
+
|
|
19
|
+
from .editing import iterative_plan_refinement_editor
|
|
20
|
+
from .editing import llm_as_a_judge_editor
|
|
21
|
+
from .editing import motivation_aware_summarization
|
|
22
|
+
from .functions import execute_score_select_function
|
|
23
|
+
from .functions import its_tool_orchestration_function
|
|
24
|
+
from .functions import its_tool_wrapper_function
|
|
25
|
+
from .functions import plan_select_execute_function
|
|
26
|
+
from .scoring import llm_based_agent_scorer
|
|
27
|
+
from .scoring import llm_based_plan_scorer
|
|
28
|
+
from .scoring import motivation_aware_scorer
|
|
29
|
+
from .search import multi_llm_planner
|
|
30
|
+
from .search import multi_query_retrieval_search
|
|
31
|
+
from .search import single_shot_multi_plan_planner
|
|
32
|
+
from .selection import best_of_n_selector
|
|
33
|
+
from .selection import llm_based_agent_output_selector
|
|
34
|
+
from .selection import llm_based_output_merging_selector
|
|
35
|
+
from .selection import llm_based_plan_selector
|
|
36
|
+
from .selection import threshold_selector
|
|
File without changes
|
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import asyncio
|
|
17
|
+
import logging
|
|
18
|
+
import re
|
|
19
|
+
|
|
20
|
+
from aiq.builder.builder import Builder
|
|
21
|
+
from aiq.builder.framework_enum import LLMFrameworkEnum
|
|
22
|
+
from aiq.cli.register_workflow import register_its_strategy
|
|
23
|
+
from aiq.data_models.its_strategy import ITSStrategyBaseConfig
|
|
24
|
+
from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
|
|
25
|
+
from aiq.experimental.inference_time_scaling.models.scoring_config import LLMBasedAgentScoringConfig
|
|
26
|
+
from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
|
|
27
|
+
from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
|
|
28
|
+
from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
|
|
29
|
+
from aiq.utils.io.model_processing import remove_r1_think_tags
|
|
30
|
+
|
|
31
|
+
logger = logging.getLogger(__name__)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class LLMBasedAgentScorer(StrategyBase):
|
|
35
|
+
|
|
36
|
+
def __init__(self, config: ITSStrategyBaseConfig) -> None:
|
|
37
|
+
super().__init__(config)
|
|
38
|
+
self.llm_bound = None
|
|
39
|
+
|
|
40
|
+
async def build_components(self, builder: Builder) -> None:
|
|
41
|
+
"""
|
|
42
|
+
Build the components required for the planner.
|
|
43
|
+
"""
|
|
44
|
+
self.llm_bound = await builder.get_llm(self.config.scoring_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
45
|
+
|
|
46
|
+
def supported_pipeline_types(self) -> [PipelineTypeEnum]:
|
|
47
|
+
return [PipelineTypeEnum.AGENT_EXECUTION]
|
|
48
|
+
|
|
49
|
+
def stage_type(self) -> StageTypeEnum:
|
|
50
|
+
return StageTypeEnum.SCORING
|
|
51
|
+
|
|
52
|
+
async def score_single(self, original_prompt: str, agent_context: str, item: ITSItem) -> float:
|
|
53
|
+
"""
|
|
54
|
+
Score a single planning item using the LLM.
|
|
55
|
+
|
|
56
|
+
Args:
|
|
57
|
+
original_prompt (str): The original prompt.
|
|
58
|
+
agent_context (str): The agent context.
|
|
59
|
+
item (ITSItem): The item to score.
|
|
60
|
+
|
|
61
|
+
Returns:
|
|
62
|
+
float: The score of the item.
|
|
63
|
+
"""
|
|
64
|
+
|
|
65
|
+
try:
|
|
66
|
+
from langchain_core.language_models import BaseChatModel
|
|
67
|
+
from langchain_core.prompts import PromptTemplate
|
|
68
|
+
except ImportError:
|
|
69
|
+
raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
|
|
70
|
+
"This error can be resolved by installing aiqtoolkit-langchain.")
|
|
71
|
+
|
|
72
|
+
if not isinstance(self.llm_bound, BaseChatModel):
|
|
73
|
+
raise ValueError("The `scoring_llm` must be an instance of `BaseChatModel`.")
|
|
74
|
+
|
|
75
|
+
model: BaseChatModel = self.llm_bound
|
|
76
|
+
|
|
77
|
+
prompt_template = PromptTemplate(
|
|
78
|
+
template=self.config.scoring_template,
|
|
79
|
+
input_variables=["objective", "input", "output"],
|
|
80
|
+
validate_template=True,
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
prompt = (await prompt_template.ainvoke(
|
|
84
|
+
input={
|
|
85
|
+
"objective": agent_context,
|
|
86
|
+
"input": str(item.input) if not original_prompt else original_prompt,
|
|
87
|
+
"output": str(item.output)
|
|
88
|
+
}))
|
|
89
|
+
|
|
90
|
+
response = (await model.ainvoke(prompt)).content
|
|
91
|
+
response = remove_r1_think_tags(response)
|
|
92
|
+
|
|
93
|
+
# Score will following the format of `FINAL SCORE: <float>` in the response from the LLM
|
|
94
|
+
if not isinstance(response, str):
|
|
95
|
+
logger.warning(f"Invalid response from LLM for scoring: {response}.")
|
|
96
|
+
raise ValueError("Unable to parse the score from the LLM response.")
|
|
97
|
+
|
|
98
|
+
response = response.strip()
|
|
99
|
+
match = re.search(r'FINAL SCORE:\s*([\d.]+)', response)
|
|
100
|
+
if not match:
|
|
101
|
+
logger.warning(f"Could not parse the score from the response: {response}.")
|
|
102
|
+
score_str = '0.0'
|
|
103
|
+
else:
|
|
104
|
+
score_str = match.group(1)
|
|
105
|
+
|
|
106
|
+
try:
|
|
107
|
+
score = float(score_str)
|
|
108
|
+
except ValueError:
|
|
109
|
+
logger.warning(f"Could not convert the score string '{score_str}' to float.")
|
|
110
|
+
raise ValueError(f"Unable to convert the extracted score '{score_str}' to a float.")
|
|
111
|
+
|
|
112
|
+
return score
|
|
113
|
+
|
|
114
|
+
async def ainvoke(self,
|
|
115
|
+
items: list[ITSItem],
|
|
116
|
+
original_prompt: str | None = None,
|
|
117
|
+
agent_context: str | None = None,
|
|
118
|
+
**kwargs) -> list[ITSItem]:
|
|
119
|
+
"""
|
|
120
|
+
Score a list of planning items.
|
|
121
|
+
|
|
122
|
+
Args:
|
|
123
|
+
original_prompt (str): The original prompt.
|
|
124
|
+
agent_context (str): The agent context.
|
|
125
|
+
items (list[ITSItem]): The list of planning items to score.
|
|
126
|
+
|
|
127
|
+
Returns:
|
|
128
|
+
list[float]: A list of scores corresponding to each planning item.
|
|
129
|
+
"""
|
|
130
|
+
# Run score single concurrently for all planning items
|
|
131
|
+
# Then set the score attribute on each planning item
|
|
132
|
+
if not items:
|
|
133
|
+
return []
|
|
134
|
+
tasks = [
|
|
135
|
+
self.score_single(original_prompt=original_prompt, agent_context=agent_context, item=item) for item in items
|
|
136
|
+
]
|
|
137
|
+
|
|
138
|
+
# Gather all scores concurrently
|
|
139
|
+
scores = await asyncio.gather(*tasks)
|
|
140
|
+
|
|
141
|
+
if len(scores) != len(items):
|
|
142
|
+
logger.warning(f"Number of scores {len(scores)} does not match the number of items {len(items)}.")
|
|
143
|
+
raise ValueError("Mismatch in number of scores and planning items.")
|
|
144
|
+
|
|
145
|
+
logger.debug("Scores for planning items: %s", scores)
|
|
146
|
+
|
|
147
|
+
# Set the score on each planning item for reference
|
|
148
|
+
for idx, score in enumerate(scores):
|
|
149
|
+
items[idx].score = score
|
|
150
|
+
|
|
151
|
+
return items
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
@register_its_strategy(config_type=LLMBasedAgentScoringConfig)
|
|
155
|
+
async def register_llm_based_agent_scorer(config: LLMBasedAgentScoringConfig, builder: Builder):
|
|
156
|
+
"""
|
|
157
|
+
Register the LLM-based agent scorer with the provided configuration and builder.
|
|
158
|
+
|
|
159
|
+
Args:
|
|
160
|
+
config (LLMBasedAgentScoringConfig): The configuration for the LLM-based agent scorer.
|
|
161
|
+
builder (Builder): The builder instance to use for building components.
|
|
162
|
+
|
|
163
|
+
Returns:
|
|
164
|
+
LLMBasedAgentScorer: The registered LLM-based agent scorer.
|
|
165
|
+
"""
|
|
166
|
+
scorer = LLMBasedAgentScorer(config)
|
|
167
|
+
await scorer.build_components(builder)
|
|
168
|
+
yield scorer
|