datarobot-genai 0.2.42__tar.gz → 0.3.0__tar.gz
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.
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/PKG-INFO +3 -1
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/pyproject.toml +3 -1
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/agents/__init__.py +1 -1
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/crewai/__init__.py +1 -4
- datarobot_genai-0.2.42/src/datarobot_genai/crewai/base.py → datarobot_genai-0.3.0/src/datarobot_genai/crewai/agent.py +14 -1
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/config.py +21 -1
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/tool_config.py +16 -0
- datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/clients/perplexity.py +173 -0
- datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/clients/tavily.py +199 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/gdrive/tools.py +2 -4
- datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/perplexity/tools.py +121 -0
- datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/tavily/__init__.py +13 -0
- datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/tavily/tools.py +148 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/llama_index/__init__.py +1 -1
- datarobot_genai-0.2.42/src/datarobot_genai/llama_index/base.py → datarobot_genai-0.3.0/src/datarobot_genai/llama_index/agent.py +37 -10
- datarobot_genai-0.3.0/src/datarobot_genai/py.typed +0 -0
- datarobot_genai-0.2.42/src/datarobot_genai/crewai/agent.py +0 -50
- datarobot_genai-0.2.42/src/datarobot_genai/llama_index/agent.py +0 -57
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/.gitignore +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/AUTHORS +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/LICENSE +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/README.md +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/agents/base.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/chat/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/chat/auth.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/chat/client.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/chat/responses.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/cli/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/cli/agent_kernel.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/custom_model.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/mcp/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/mcp/common.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/telemetry_agent.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/utils/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/utils/auth.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/utils/urls.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/crewai/events.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/crewai/mcp.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/auth.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/clients.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/constants.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/exceptions.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/logging.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/routes.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/utils.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/server.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/clients/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/clients/anthropic.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/clients/base.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/clients/dr_gateway.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/clients/openai.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/elicitation_test_tool.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/test_interactive.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/test_utils/utils.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/atlassian.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/confluence.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/gdrive.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/jira.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/clients/s3.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/confluence/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/confluence/tools.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/gdrive/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/jira/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/jira/tools.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/microsoft_graph/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/microsoft_graph/tools.py +0 -0
- {datarobot_genai-0.2.42/src/datarobot_genai/langgraph → datarobot_genai-0.3.0/src/datarobot_genai/drmcp/tools/perplexity}/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/data.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/predict.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/project.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/predictive/training.py +0 -0
- {datarobot_genai-0.2.42/src/datarobot_genai/nat → datarobot_genai-0.3.0/src/datarobot_genai/langgraph}/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/langgraph/agent.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/langgraph/mcp.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/llama_index/mcp.py +0 -0
- /datarobot_genai-0.2.42/src/datarobot_genai/py.typed → /datarobot_genai-0.3.0/src/datarobot_genai/nat/__init__.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/agent.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/datarobot_auth_provider.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/datarobot_llm_clients.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/datarobot_mcp_client.py +0 -0
- {datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/nat/helpers.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: datarobot-genai
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Generic helpers for GenAI
|
|
5
5
|
Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
|
|
6
6
|
Author: DataRobot, Inc.
|
|
@@ -43,9 +43,11 @@ Requires-Dist: opentelemetry-api<2.0.0,>=1.22.0; extra == 'drmcp'
|
|
|
43
43
|
Requires-Dist: opentelemetry-exporter-otlp-proto-http<2.0.0,>=1.22.0; extra == 'drmcp'
|
|
44
44
|
Requires-Dist: opentelemetry-exporter-otlp<2.0.0,>=1.22.0; extra == 'drmcp'
|
|
45
45
|
Requires-Dist: opentelemetry-sdk<2.0.0,>=1.22.0; extra == 'drmcp'
|
|
46
|
+
Requires-Dist: perplexityai<1.0,>=0.27; extra == 'drmcp'
|
|
46
47
|
Requires-Dist: pydantic-settings<3.0.0,>=2.1.0; extra == 'drmcp'
|
|
47
48
|
Requires-Dist: pydantic<3.0.0,>=2.6.1; extra == 'drmcp'
|
|
48
49
|
Requires-Dist: python-dotenv<2.0.0,>=1.1.0; extra == 'drmcp'
|
|
50
|
+
Requires-Dist: tavily-python<1.0.0,>=0.7.20; extra == 'drmcp'
|
|
49
51
|
Provides-Extra: langgraph
|
|
50
52
|
Requires-Dist: langchain-mcp-adapters<0.2.0,>=0.1.12; extra == 'langgraph'
|
|
51
53
|
Requires-Dist: langgraph-prebuilt<0.7.0,>=0.2.3; extra == 'langgraph'
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "datarobot-genai"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.3.0"
|
|
8
8
|
description = "Generic helpers for GenAI"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10, <3.13"
|
|
@@ -84,6 +84,8 @@ drmcp = [
|
|
|
84
84
|
"python-dotenv>=1.1.0,<2.0.0",
|
|
85
85
|
"boto3>=1.34.0,<2.0.0",
|
|
86
86
|
"httpx>=0.28.1,<1.0.0",
|
|
87
|
+
"tavily-python>=0.7.20,<1.0.0",
|
|
88
|
+
"perplexityai>=0.27,<1.0",
|
|
87
89
|
"pydantic>=2.6.1,<3.0.0",
|
|
88
90
|
"pydantic-settings>=2.1.0,<3.0.0",
|
|
89
91
|
"opentelemetry-api>=1.22.0,<2.0.0",
|
{datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/core/agents/__init__.py
RENAMED
|
@@ -17,7 +17,7 @@ This package provides:
|
|
|
17
17
|
- BaseAgent: common initialization for agent env/config fields
|
|
18
18
|
- Common helpers: make_system_prompt, extract_user_prompt_content
|
|
19
19
|
- Framework utilities (optional extras):
|
|
20
|
-
- crewai:
|
|
20
|
+
- crewai: create_pipeline_interactions_from_messages
|
|
21
21
|
- langgraph: create_pipeline_interactions_from_events
|
|
22
22
|
- llamaindex: DataRobotLiteLLM, create_pipeline_interactions_from_events
|
|
23
23
|
"""
|
|
@@ -2,22 +2,19 @@
|
|
|
2
2
|
|
|
3
3
|
Public API:
|
|
4
4
|
- mcp_tools_context: Context manager returning available MCP tools for CrewAI.
|
|
5
|
-
- build_llm: Construct a CrewAI LLM configured for DataRobot endpoints.
|
|
6
5
|
- create_pipeline_interactions_from_messages: Convert messages to MultiTurnSample.
|
|
7
6
|
"""
|
|
8
7
|
|
|
9
8
|
from datarobot_genai.core.mcp.common import MCPConfig
|
|
10
9
|
|
|
11
|
-
from .agent import
|
|
10
|
+
from .agent import CrewAIAgent
|
|
12
11
|
from .agent import create_pipeline_interactions_from_messages
|
|
13
|
-
from .base import CrewAIAgent
|
|
14
12
|
from .events import CrewAIEventListener
|
|
15
13
|
from .mcp import mcp_tools_context
|
|
16
14
|
|
|
17
15
|
__all__ = [
|
|
18
16
|
"mcp_tools_context",
|
|
19
17
|
"CrewAIAgent",
|
|
20
|
-
"build_llm",
|
|
21
18
|
"create_pipeline_interactions_from_messages",
|
|
22
19
|
"CrewAIEventListener",
|
|
23
20
|
"MCPConfig",
|
|
@@ -41,11 +41,24 @@ from datarobot_genai.core.agents.base import default_usage_metrics
|
|
|
41
41
|
from datarobot_genai.core.agents.base import extract_user_prompt_content
|
|
42
42
|
from datarobot_genai.core.agents.base import is_streaming
|
|
43
43
|
|
|
44
|
-
from .agent import create_pipeline_interactions_from_messages
|
|
45
44
|
from .mcp import mcp_tools_context
|
|
46
45
|
|
|
47
46
|
if TYPE_CHECKING:
|
|
48
47
|
from ragas import MultiTurnSample
|
|
48
|
+
from ragas.messages import AIMessage
|
|
49
|
+
from ragas.messages import HumanMessage
|
|
50
|
+
from ragas.messages import ToolMessage
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def create_pipeline_interactions_from_messages(
|
|
54
|
+
messages: list[HumanMessage | AIMessage | ToolMessage] | None,
|
|
55
|
+
) -> MultiTurnSample | None:
|
|
56
|
+
if not messages:
|
|
57
|
+
return None
|
|
58
|
+
# Lazy import to reduce memory overhead when ragas is not used
|
|
59
|
+
from ragas import MultiTurnSample
|
|
60
|
+
|
|
61
|
+
return MultiTurnSample(user_input=messages)
|
|
49
62
|
|
|
50
63
|
|
|
51
64
|
class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
|
|
@@ -76,7 +76,25 @@ class MCPToolConfig(BaseSettings):
|
|
|
76
76
|
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
|
|
77
77
|
"ENABLE_MICROSOFT_GRAPH_TOOLS",
|
|
78
78
|
),
|
|
79
|
-
description="Enable/disable Sharepoint tools",
|
|
79
|
+
description="Enable/disable Microsoft Graph (Sharepoint/OneDrive) tools",
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
enable_perplexity_tools: bool = Field(
|
|
83
|
+
default=False,
|
|
84
|
+
validation_alias=AliasChoices(
|
|
85
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PERPLEXITY_TOOLS",
|
|
86
|
+
"ENABLE_PERPLEXITY_TOOLS",
|
|
87
|
+
),
|
|
88
|
+
description="Enable/disable Perplexity tools",
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
enable_tavily_tools: bool = Field(
|
|
92
|
+
default=False,
|
|
93
|
+
validation_alias=AliasChoices(
|
|
94
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_TAVILY_TOOLS",
|
|
95
|
+
"ENABLE_TAVILY_TOOLS",
|
|
96
|
+
),
|
|
97
|
+
description="Enable/disable Tavily search tools",
|
|
80
98
|
)
|
|
81
99
|
|
|
82
100
|
is_atlassian_oauth_provider_configured: bool = Field(
|
|
@@ -131,6 +149,8 @@ class MCPToolConfig(BaseSettings):
|
|
|
131
149
|
"enable_confluence_tools",
|
|
132
150
|
"enable_gdrive_tools",
|
|
133
151
|
"enable_microsoft_graph_tools",
|
|
152
|
+
"enable_perplexity_tools",
|
|
153
|
+
"enable_tavily_tools",
|
|
134
154
|
"is_atlassian_oauth_provider_configured",
|
|
135
155
|
"is_google_oauth_provider_configured",
|
|
136
156
|
"is_microsoft_oauth_provider_configured",
|
{datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/core/tool_config.py
RENAMED
|
@@ -31,6 +31,8 @@ class ToolType(str, Enum):
|
|
|
31
31
|
CONFLUENCE = "confluence"
|
|
32
32
|
GDRIVE = "gdrive"
|
|
33
33
|
MICROSOFT_GRAPH = "microsoft_graph"
|
|
34
|
+
PERPLEXITY = "perplexity"
|
|
35
|
+
TAVILY = "tavily"
|
|
34
36
|
|
|
35
37
|
|
|
36
38
|
class ToolConfig(TypedDict):
|
|
@@ -80,6 +82,20 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
|
|
|
80
82
|
package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
|
|
81
83
|
config_field_name="enable_microsoft_graph_tools",
|
|
82
84
|
),
|
|
85
|
+
ToolType.PERPLEXITY: ToolConfig(
|
|
86
|
+
name="perplexity",
|
|
87
|
+
oauth_check=None, # OAuth for Perplexity is not supported
|
|
88
|
+
directory="perplexity",
|
|
89
|
+
package_prefix="datarobot_genai.drmcp.tools.perplexity",
|
|
90
|
+
config_field_name="enable_perplexity_tools",
|
|
91
|
+
),
|
|
92
|
+
ToolType.TAVILY: ToolConfig(
|
|
93
|
+
name="tavily",
|
|
94
|
+
oauth_check=None,
|
|
95
|
+
directory="tavily",
|
|
96
|
+
package_prefix="datarobot_genai.drmcp.tools.tavily",
|
|
97
|
+
config_field_name="enable_tavily_tools",
|
|
98
|
+
),
|
|
83
99
|
}
|
|
84
100
|
|
|
85
101
|
|
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
import logging
|
|
16
|
+
from typing import Any
|
|
17
|
+
from typing import Literal
|
|
18
|
+
|
|
19
|
+
from fastmcp.exceptions import ToolError
|
|
20
|
+
from fastmcp.server.dependencies import get_http_headers
|
|
21
|
+
from perplexity import AsyncPerplexity
|
|
22
|
+
from perplexity.types import search_create_response
|
|
23
|
+
from pydantic import BaseModel
|
|
24
|
+
from pydantic import ConfigDict
|
|
25
|
+
|
|
26
|
+
logger = logging.getLogger(__name__)
|
|
27
|
+
|
|
28
|
+
MAX_QUERIES: int = 5
|
|
29
|
+
MAX_RESULTS: int = 20
|
|
30
|
+
MAX_TOKENS_PER_PAGE: int = 8192
|
|
31
|
+
MAX_SEARCH_DOMAIN_FILTER: int = 20
|
|
32
|
+
|
|
33
|
+
MAX_RESULTS_DEFAULT: int = 10
|
|
34
|
+
MAX_TOKENS_PER_PAGE_DEFAULT: int = 2048
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
async def get_perplexity_access_token() -> str | ToolError:
|
|
38
|
+
"""
|
|
39
|
+
Get Perplexity API key from HTTP headers.
|
|
40
|
+
|
|
41
|
+
At the moment of creating this fn. Perplexity does not support OAuth.
|
|
42
|
+
It allows only API-KEY authorized flow.
|
|
43
|
+
|
|
44
|
+
Returns
|
|
45
|
+
-------
|
|
46
|
+
Access token string on success, ToolError on failure
|
|
47
|
+
|
|
48
|
+
Example:
|
|
49
|
+
```python
|
|
50
|
+
token = await get_perplexity_access_token()
|
|
51
|
+
if isinstance(token, ToolError):
|
|
52
|
+
# Handle error
|
|
53
|
+
return token
|
|
54
|
+
# Use token
|
|
55
|
+
```
|
|
56
|
+
"""
|
|
57
|
+
try:
|
|
58
|
+
headers = get_http_headers()
|
|
59
|
+
|
|
60
|
+
if api_key := headers.get("x-perplexity-api-key"):
|
|
61
|
+
return api_key
|
|
62
|
+
|
|
63
|
+
logger.warning("Perplexity API key not found in headers.")
|
|
64
|
+
return ToolError(
|
|
65
|
+
"Perplexity API key not found in headers. "
|
|
66
|
+
"Please provide it via 'x-perplexity-api-key' header."
|
|
67
|
+
)
|
|
68
|
+
except Exception as e:
|
|
69
|
+
logger.error(f"Unexpected error obtaining Perplexity API key: {e}.", exc_info=e)
|
|
70
|
+
return ToolError("An unexpected error occured while obtaining Perplexity API key.")
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class PerplexityError(Exception):
|
|
74
|
+
"""Exception for Perplexity API errors."""
|
|
75
|
+
|
|
76
|
+
def __init__(self, message: str) -> None:
|
|
77
|
+
super().__init__(message)
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
class PerplexitySearchResult(BaseModel):
|
|
81
|
+
snippet: str
|
|
82
|
+
title: str
|
|
83
|
+
url: str
|
|
84
|
+
date: str | None = None
|
|
85
|
+
last_updated: str | None = None
|
|
86
|
+
|
|
87
|
+
model_config = ConfigDict(populate_by_name=True)
|
|
88
|
+
|
|
89
|
+
@classmethod
|
|
90
|
+
def from_perplexity_sdk(cls, result: search_create_response.Result) -> "PerplexitySearchResult":
|
|
91
|
+
"""Create a PerplexitySearchResult from perplexity sdk response data."""
|
|
92
|
+
return cls(**result.model_dump())
|
|
93
|
+
|
|
94
|
+
def as_flat_dict(self) -> dict[str, Any]:
|
|
95
|
+
"""Return a flat dictionary representation of the search result."""
|
|
96
|
+
return self.model_dump(by_alias=True)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
class PerplexityClient:
|
|
100
|
+
"""Client for interacting with Perplexity API.
|
|
101
|
+
Its simple wrapper around perplexity python sdk.
|
|
102
|
+
"""
|
|
103
|
+
|
|
104
|
+
def __init__(self, access_token: str) -> None:
|
|
105
|
+
self._client = AsyncPerplexity(api_key=access_token)
|
|
106
|
+
|
|
107
|
+
async def search(
|
|
108
|
+
self,
|
|
109
|
+
query: str | list[str],
|
|
110
|
+
search_domain_filter: list[str] | None = None,
|
|
111
|
+
recency: Literal["hour", "day", "week", "month", "year"] | None = None,
|
|
112
|
+
max_results: int = MAX_RESULTS_DEFAULT,
|
|
113
|
+
max_tokens_per_page: int = MAX_TOKENS_PER_PAGE_DEFAULT,
|
|
114
|
+
) -> list[PerplexitySearchResult]:
|
|
115
|
+
"""
|
|
116
|
+
Search using Perplexity.
|
|
117
|
+
|
|
118
|
+
Args:
|
|
119
|
+
query: Query to filter results.
|
|
120
|
+
search_domain_filter: Up to 20 domains/URLs to allowlist or denylist.
|
|
121
|
+
recency: Filter results by time period.
|
|
122
|
+
max_results: Number of ranked results to return.
|
|
123
|
+
max_tokens_per_page: Context extraction cap per page.
|
|
124
|
+
|
|
125
|
+
Returns
|
|
126
|
+
-------
|
|
127
|
+
List of Perplexity search results.
|
|
128
|
+
"""
|
|
129
|
+
if not query:
|
|
130
|
+
raise PerplexityError("Error: query cannot be empty.")
|
|
131
|
+
if query and isinstance(query, str) and not query.strip():
|
|
132
|
+
raise PerplexityError("Error: query cannot be empty.")
|
|
133
|
+
if query and isinstance(query, list) and len(query) > MAX_QUERIES:
|
|
134
|
+
raise PerplexityError(f"Error: query list cannot be bigger than {MAX_QUERIES}.")
|
|
135
|
+
if query and isinstance(query, list) and not all(q.strip() for q in query):
|
|
136
|
+
raise PerplexityError("Error: query cannot contain empty str.")
|
|
137
|
+
if search_domain_filter and len(search_domain_filter) > MAX_SEARCH_DOMAIN_FILTER:
|
|
138
|
+
raise PerplexityError("Error: maximum number of search domain filters is 20.")
|
|
139
|
+
if max_results <= 0:
|
|
140
|
+
raise PerplexityError("Error: max_results must be greater than 0.")
|
|
141
|
+
if max_results > MAX_RESULTS:
|
|
142
|
+
raise PerplexityError("Error: max_results must be smaller than or equal to 20.")
|
|
143
|
+
if max_tokens_per_page <= 0:
|
|
144
|
+
raise PerplexityError("Error: max_tokens_per_page must be greater than 0.")
|
|
145
|
+
if max_tokens_per_page > MAX_TOKENS_PER_PAGE:
|
|
146
|
+
raise PerplexityError(
|
|
147
|
+
"Error: max_tokens_per_page must be smaller than or equal to 8192."
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
max_results = min(max_results, MAX_RESULTS)
|
|
151
|
+
max_tokens_per_page = min(max_tokens_per_page, MAX_TOKENS_PER_PAGE)
|
|
152
|
+
|
|
153
|
+
search_result = await self._client.search.create(
|
|
154
|
+
query=query,
|
|
155
|
+
search_domain_filter=search_domain_filter,
|
|
156
|
+
search_recency_filter=recency,
|
|
157
|
+
max_results=max_results,
|
|
158
|
+
max_tokens_per_page=max_tokens_per_page,
|
|
159
|
+
)
|
|
160
|
+
|
|
161
|
+
return [
|
|
162
|
+
PerplexitySearchResult.from_perplexity_sdk(result) for result in search_result.results
|
|
163
|
+
]
|
|
164
|
+
|
|
165
|
+
async def __aenter__(self) -> "PerplexityClient":
|
|
166
|
+
"""Async context manager entry."""
|
|
167
|
+
return self
|
|
168
|
+
|
|
169
|
+
async def __aexit__(
|
|
170
|
+
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
|
|
171
|
+
) -> None:
|
|
172
|
+
"""Async context manager exit."""
|
|
173
|
+
await self._client.close()
|
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Tavily API Client and utilities for API key authentication."""
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Any
|
|
19
|
+
from typing import Literal
|
|
20
|
+
|
|
21
|
+
from fastmcp.exceptions import ToolError
|
|
22
|
+
from fastmcp.server.dependencies import get_http_headers
|
|
23
|
+
from pydantic import BaseModel
|
|
24
|
+
from pydantic import ConfigDict
|
|
25
|
+
from tavily import AsyncTavilyClient
|
|
26
|
+
|
|
27
|
+
logger = logging.getLogger(__name__)
|
|
28
|
+
|
|
29
|
+
MAX_RESULTS: int = 20
|
|
30
|
+
MAX_CHUNKS_PER_SOURCE: int = 3
|
|
31
|
+
|
|
32
|
+
MAX_RESULTS_DEFAULT: int = 5
|
|
33
|
+
CHUNKS_PER_SOURCE_DEFAULT: int = 1
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
async def get_tavily_access_token() -> str:
|
|
37
|
+
"""
|
|
38
|
+
Get Tavily API key from HTTP headers.
|
|
39
|
+
|
|
40
|
+
Returns
|
|
41
|
+
-------
|
|
42
|
+
API key string
|
|
43
|
+
|
|
44
|
+
Raises
|
|
45
|
+
------
|
|
46
|
+
ToolError: If API key is not found in headers
|
|
47
|
+
"""
|
|
48
|
+
headers = get_http_headers()
|
|
49
|
+
|
|
50
|
+
api_key = headers.get("x-tavily-api-key")
|
|
51
|
+
if api_key:
|
|
52
|
+
return api_key
|
|
53
|
+
|
|
54
|
+
logger.warning("Tavily API key not found in headers")
|
|
55
|
+
raise ToolError(
|
|
56
|
+
"Tavily API key not found in headers. Please provide it via 'x-tavily-api-key' header."
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class TavilySearchResult(BaseModel):
|
|
61
|
+
"""A single search result from Tavily API."""
|
|
62
|
+
|
|
63
|
+
title: str
|
|
64
|
+
url: str
|
|
65
|
+
content: str
|
|
66
|
+
score: float
|
|
67
|
+
|
|
68
|
+
model_config = ConfigDict(populate_by_name=True)
|
|
69
|
+
|
|
70
|
+
@classmethod
|
|
71
|
+
def from_tavily_sdk(cls, result: dict[str, Any]) -> "TavilySearchResult":
|
|
72
|
+
"""Create a TavilySearchResult from Tavily SDK response data."""
|
|
73
|
+
return cls(
|
|
74
|
+
title=result.get("title", ""),
|
|
75
|
+
url=result.get("url", ""),
|
|
76
|
+
content=result.get("content", ""),
|
|
77
|
+
score=result.get("score", 0.0),
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
def as_flat_dict(self) -> dict[str, Any]:
|
|
81
|
+
"""Return a flat dictionary representation of the search result."""
|
|
82
|
+
return self.model_dump(by_alias=True)
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
class TavilyImage(BaseModel):
|
|
86
|
+
"""An image result from Tavily API."""
|
|
87
|
+
|
|
88
|
+
url: str
|
|
89
|
+
description: str | None = None
|
|
90
|
+
|
|
91
|
+
model_config = ConfigDict(populate_by_name=True)
|
|
92
|
+
|
|
93
|
+
@classmethod
|
|
94
|
+
def from_tavily_sdk(cls, image: dict[str, Any] | str) -> "TavilyImage":
|
|
95
|
+
"""Create a TavilyImage from Tavily SDK response data."""
|
|
96
|
+
if isinstance(image, str):
|
|
97
|
+
return cls(url=image)
|
|
98
|
+
return cls(
|
|
99
|
+
url=image.get("url", ""),
|
|
100
|
+
description=image.get("description"),
|
|
101
|
+
)
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class TavilyClient:
|
|
105
|
+
"""Client for interacting with Tavily Search API.
|
|
106
|
+
|
|
107
|
+
This is a wrapper around the official tavily-python SDK.
|
|
108
|
+
"""
|
|
109
|
+
|
|
110
|
+
def __init__(self, api_key: str) -> None:
|
|
111
|
+
self._client = AsyncTavilyClient(api_key=api_key)
|
|
112
|
+
|
|
113
|
+
async def search(
|
|
114
|
+
self,
|
|
115
|
+
query: str,
|
|
116
|
+
*,
|
|
117
|
+
topic: Literal["general", "news", "finance"] = "general",
|
|
118
|
+
search_depth: Literal["basic", "advanced"] = "basic",
|
|
119
|
+
max_results: int = MAX_RESULTS_DEFAULT,
|
|
120
|
+
time_range: Literal["day", "week", "month", "year"] | None = None,
|
|
121
|
+
include_images: bool = False,
|
|
122
|
+
include_image_descriptions: bool = False,
|
|
123
|
+
chunks_per_source: int = CHUNKS_PER_SOURCE_DEFAULT,
|
|
124
|
+
include_answer: bool = False,
|
|
125
|
+
) -> dict[str, Any]:
|
|
126
|
+
"""
|
|
127
|
+
Perform a web search using Tavily API.
|
|
128
|
+
|
|
129
|
+
Args:
|
|
130
|
+
query: The search query to execute.
|
|
131
|
+
topic: The category of search ("general", "news", or "finance").
|
|
132
|
+
search_depth: The depth of search ("basic" or "advanced").
|
|
133
|
+
max_results: Maximum number of results to return (1-20).
|
|
134
|
+
time_range: Time range filter ("day", "week", "month", "year").
|
|
135
|
+
include_images: Whether to include images in results.
|
|
136
|
+
include_image_descriptions: Whether to include image descriptions.
|
|
137
|
+
chunks_per_source: Maximum content snippets per URL (1-3).
|
|
138
|
+
include_answer: Whether to include an AI-generated answer.
|
|
139
|
+
|
|
140
|
+
Returns
|
|
141
|
+
-------
|
|
142
|
+
Dict with search results from Tavily API.
|
|
143
|
+
|
|
144
|
+
Raises
|
|
145
|
+
------
|
|
146
|
+
ValueError: If validation fails.
|
|
147
|
+
TavilyInvalidAPIKeyError: If the API key is invalid.
|
|
148
|
+
TavilyUsageLimitExceededError: If usage limit is exceeded.
|
|
149
|
+
TavilyForbiddenError: If access is forbidden.
|
|
150
|
+
TavilyBadRequestError: If the request is malformed.
|
|
151
|
+
"""
|
|
152
|
+
# Validate inputs
|
|
153
|
+
if not query:
|
|
154
|
+
raise ValueError("query cannot be empty.")
|
|
155
|
+
if isinstance(query, str) and not query.strip():
|
|
156
|
+
raise ValueError("query cannot be empty.")
|
|
157
|
+
if max_results <= 0:
|
|
158
|
+
raise ValueError("max_results must be greater than 0.")
|
|
159
|
+
if max_results > MAX_RESULTS:
|
|
160
|
+
raise ValueError(f"max_results must be smaller than or equal to {MAX_RESULTS}.")
|
|
161
|
+
if chunks_per_source <= 0:
|
|
162
|
+
raise ValueError("chunks_per_source must be greater than 0.")
|
|
163
|
+
if chunks_per_source > MAX_CHUNKS_PER_SOURCE:
|
|
164
|
+
raise ValueError(
|
|
165
|
+
f"chunks_per_source must be smaller than or equal to {MAX_CHUNKS_PER_SOURCE}."
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
# Clamp values to valid ranges
|
|
169
|
+
max_results = min(max_results, MAX_RESULTS)
|
|
170
|
+
chunks_per_source = min(chunks_per_source, MAX_CHUNKS_PER_SOURCE)
|
|
171
|
+
|
|
172
|
+
# Build search parameters
|
|
173
|
+
search_kwargs: dict[str, Any] = {
|
|
174
|
+
"query": query,
|
|
175
|
+
"topic": topic,
|
|
176
|
+
"search_depth": search_depth,
|
|
177
|
+
"max_results": max_results,
|
|
178
|
+
"include_images": include_images,
|
|
179
|
+
"include_image_descriptions": include_image_descriptions,
|
|
180
|
+
"chunks_per_source": chunks_per_source,
|
|
181
|
+
"include_answer": include_answer,
|
|
182
|
+
}
|
|
183
|
+
|
|
184
|
+
if time_range:
|
|
185
|
+
search_kwargs["time_range"] = time_range
|
|
186
|
+
|
|
187
|
+
return await self._client.search(**search_kwargs)
|
|
188
|
+
|
|
189
|
+
async def __aenter__(self) -> "TavilyClient":
|
|
190
|
+
"""Async context manager entry."""
|
|
191
|
+
return self
|
|
192
|
+
|
|
193
|
+
async def __aexit__(
|
|
194
|
+
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
|
|
195
|
+
) -> None:
|
|
196
|
+
"""Async context manager exit."""
|
|
197
|
+
# AsyncTavilyClient doesn't have a close method, but we keep the context manager
|
|
198
|
+
# pattern for consistency with other clients
|
|
199
|
+
pass
|
{datarobot_genai-0.2.42 → datarobot_genai-0.3.0}/src/datarobot_genai/drmcp/tools/gdrive/tools.py
RENAMED
|
@@ -33,9 +33,7 @@ from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
|
|
|
33
33
|
logger = logging.getLogger(__name__)
|
|
34
34
|
|
|
35
35
|
|
|
36
|
-
@dr_mcp_tool(
|
|
37
|
-
tags={"google", "gdrive", "list", "search", "files", "find", "contents"}, enabled=False
|
|
38
|
-
)
|
|
36
|
+
@dr_mcp_tool(tags={"google", "gdrive", "list", "search", "files", "find", "contents"})
|
|
39
37
|
async def gdrive_find_contents(
|
|
40
38
|
*,
|
|
41
39
|
page_size: Annotated[
|
|
@@ -317,7 +315,7 @@ async def gdrive_update_metadata(
|
|
|
317
315
|
)
|
|
318
316
|
|
|
319
317
|
|
|
320
|
-
@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"})
|
|
318
|
+
@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"}, enabled=False)
|
|
321
319
|
async def gdrive_manage_access(
|
|
322
320
|
*,
|
|
323
321
|
file_id: Annotated[str, "The ID of the file or folder."],
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
# Copyright 2026 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Perplexity MCP tools."""
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Annotated
|
|
19
|
+
from typing import Literal
|
|
20
|
+
|
|
21
|
+
from fastmcp.exceptions import ToolError
|
|
22
|
+
from fastmcp.tools.tool import ToolResult
|
|
23
|
+
|
|
24
|
+
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
25
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_QUERIES
|
|
26
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_RESULTS
|
|
27
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_RESULTS_DEFAULT
|
|
28
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_SEARCH_DOMAIN_FILTER
|
|
29
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_TOKENS_PER_PAGE
|
|
30
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import MAX_TOKENS_PER_PAGE_DEFAULT
|
|
31
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import PerplexityClient
|
|
32
|
+
from datarobot_genai.drmcp.tools.clients.perplexity import get_perplexity_access_token
|
|
33
|
+
|
|
34
|
+
logger = logging.getLogger(__name__)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@dr_mcp_tool(tags={"perplexity", "web", "search", "websearch"})
|
|
38
|
+
async def perplexity_search(
|
|
39
|
+
*,
|
|
40
|
+
query: Annotated[
|
|
41
|
+
str,
|
|
42
|
+
list[str],
|
|
43
|
+
f"The search query string OR "
|
|
44
|
+
f"a list of up to {MAX_QUERIES} sub-queries for multi-query research.",
|
|
45
|
+
],
|
|
46
|
+
search_domain_filter: Annotated[
|
|
47
|
+
list[str] | None,
|
|
48
|
+
f"Up to {MAX_SEARCH_DOMAIN_FILTER} domains/URLs "
|
|
49
|
+
f"to allowlist or denylist (prefix with '-').",
|
|
50
|
+
] = None,
|
|
51
|
+
recency: Annotated[
|
|
52
|
+
Literal["day", "week", "month", "year"] | None, "Filter results by time period."
|
|
53
|
+
] = None,
|
|
54
|
+
max_results: Annotated[
|
|
55
|
+
int, f"Number of ranked results to return (1-{MAX_RESULTS})."
|
|
56
|
+
] = MAX_RESULTS_DEFAULT,
|
|
57
|
+
max_tokens_per_page: Annotated[
|
|
58
|
+
int,
|
|
59
|
+
f"Content extraction cap per page (1-{MAX_TOKENS_PER_PAGE}) "
|
|
60
|
+
f"(default {MAX_TOKENS_PER_PAGE_DEFAULT}).",
|
|
61
|
+
] = MAX_TOKENS_PER_PAGE_DEFAULT,
|
|
62
|
+
) -> ToolResult:
|
|
63
|
+
"""Perplexity web search tool combining multi-query research and content extraction control."""
|
|
64
|
+
if not query:
|
|
65
|
+
raise ToolError("Argument validation error: query cannot be empty.")
|
|
66
|
+
if query and isinstance(query, str) and not query.strip():
|
|
67
|
+
raise ToolError("Argument validation error: query cannot be empty.")
|
|
68
|
+
if query and isinstance(query, list) and len(query) > MAX_QUERIES:
|
|
69
|
+
raise ToolError(
|
|
70
|
+
f"Argument validation error: query list cannot be bigger than {MAX_QUERIES}."
|
|
71
|
+
)
|
|
72
|
+
if query and isinstance(query, list) and not all(q.strip() for q in query):
|
|
73
|
+
raise ToolError("Argument validation error: query cannot contain empty str.")
|
|
74
|
+
if search_domain_filter and len(search_domain_filter) > MAX_SEARCH_DOMAIN_FILTER:
|
|
75
|
+
raise ToolError(
|
|
76
|
+
f"Argument validation error: "
|
|
77
|
+
f"maximum number of search domain filters is {MAX_SEARCH_DOMAIN_FILTER}."
|
|
78
|
+
)
|
|
79
|
+
if max_results <= 0:
|
|
80
|
+
raise ToolError("Argument validation error: max_results must be greater than 0.")
|
|
81
|
+
if max_results > MAX_RESULTS:
|
|
82
|
+
raise ToolError(
|
|
83
|
+
f"Argument validation error: "
|
|
84
|
+
f"max_results must be smaller than or equal to {MAX_RESULTS}."
|
|
85
|
+
)
|
|
86
|
+
if max_tokens_per_page <= 0:
|
|
87
|
+
raise ToolError("Argument validation error: max_tokens_per_page must be greater than 0.")
|
|
88
|
+
if max_tokens_per_page > MAX_TOKENS_PER_PAGE:
|
|
89
|
+
raise ToolError(
|
|
90
|
+
f"Argument validation error: "
|
|
91
|
+
f"max_tokens_per_page must be smaller than or equal to {MAX_TOKENS_PER_PAGE}."
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
access_token = await get_perplexity_access_token()
|
|
95
|
+
if isinstance(access_token, ToolError):
|
|
96
|
+
raise access_token
|
|
97
|
+
|
|
98
|
+
async with PerplexityClient(access_token=access_token) as perplexity_client:
|
|
99
|
+
results = await perplexity_client.search(
|
|
100
|
+
query=query,
|
|
101
|
+
search_domain_filter=search_domain_filter,
|
|
102
|
+
recency=recency,
|
|
103
|
+
max_results=max_results,
|
|
104
|
+
max_tokens_per_page=max_tokens_per_page,
|
|
105
|
+
)
|
|
106
|
+
|
|
107
|
+
query_txt = f"query '{query}'" if isinstance(query, str) else f"queries '{', '.join(query)}'"
|
|
108
|
+
n = len(results)
|
|
109
|
+
|
|
110
|
+
return ToolResult(
|
|
111
|
+
content=f"Successfully executed search for {query_txt}. Found {n} result(s).",
|
|
112
|
+
structured_content={
|
|
113
|
+
"results": results,
|
|
114
|
+
"count": n,
|
|
115
|
+
"metadata": {
|
|
116
|
+
"queriesExecuted": len(query) if isinstance(query, list) else 1,
|
|
117
|
+
"filtersApplied": {"domains": search_domain_filter, "recency": recency},
|
|
118
|
+
"extractionLimit": max_tokens_per_page,
|
|
119
|
+
},
|
|
120
|
+
},
|
|
121
|
+
)
|