datarobot-genai 0.2.39__py3-none-any.whl → 0.3.1__py3-none-any.whl

datarobot_genai/drmcp/tools/tavily/tools.py

@@ -0,0 +1,141 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# 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.
+
+"""Tavily MCP tools for web search."""
+
+import logging
+from typing import Annotated
+from typing import Literal
+
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.tavily import CHUNKS_PER_SOURCE_DEFAULT
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_CHUNKS_PER_SOURCE
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_RESULTS
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_RESULTS_DEFAULT
+from datarobot_genai.drmcp.tools.clients.tavily import TavilyClient
+from datarobot_genai.drmcp.tools.clients.tavily import TavilyImage
+from datarobot_genai.drmcp.tools.clients.tavily import TavilySearchResult
+from datarobot_genai.drmcp.tools.clients.tavily import get_tavily_access_token
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"tavily", "search", "web", "websearch"})
+async def tavily_search(
+    *,
+    query: Annotated[str, "The search query to execute."],
+    topic: Annotated[
+        Literal["general", "news", "finance"],
+        "The category of search. Use 'general' for broad web search, "
+        "'news' for recent news articles, or 'finance' for financial information.",
+    ] = "general",
+    search_depth: Annotated[
+        Literal["basic", "advanced"],
+        "The depth of search. 'basic' is faster and cheaper, "
+        "'advanced' provides more comprehensive results.",
+    ] = "basic",
+    max_results: Annotated[
+        int,
+        f"Maximum number of search results to return (1-{MAX_RESULTS}).",
+    ] = MAX_RESULTS_DEFAULT,
+    time_range: Annotated[
+        Literal["day", "week", "month", "year"] | None,
+        "Filter results by time range. Use 'day' for last 24 hours, "
+        "'week' for last 7 days, 'month' for last 30 days, or 'year' for last year.",
+    ] = None,
+    include_images: Annotated[
+        bool,
+        "Whether to include related images in the search results.",
+    ] = False,
+    include_image_descriptions: Annotated[
+        bool,
+        "Whether to include AI-generated descriptions for images. "
+        "Only applicable when include_images is True.",
+    ] = False,
+    chunks_per_source: Annotated[
+        int,
+        f"Maximum number of content snippets to return per source URL (1-{MAX_CHUNKS_PER_SOURCE}).",
+    ] = CHUNKS_PER_SOURCE_DEFAULT,
+    include_answer: Annotated[
+        bool,
+        "Whether to include an AI-generated answer summarizing the search results.",
+    ] = False,
+) -> ToolResult:
+    """
+    Perform a real-time web search using Tavily API.
+
+    Tavily is optimized for AI agents and provides clean, relevant search results
+    suitable for LLM consumption. Use this tool to search the web for current
+    information, news, financial data, or general knowledge.
+
+    Usage:
+        - Basic search: tavily_search(query="Python best practices 2026")
+        - News search: tavily_search(query="AI regulations", topic="news", time_range="week")
+        - Financial search: tavily_search(query="AAPL stock analysis", topic="finance")
+        - Comprehensive search: tavily_search(
+            query="climate change solutions",
+            search_depth="advanced",
+            max_results=10,
+            include_answer=True
+          )
+
+    Note:
+        - Advanced search depth consumes more API credits but provides better results
+        - Time range filtering is useful for finding recent information
+    """
+    api_key = await get_tavily_access_token()
+
+    async with TavilyClient(api_key) as client:
+        response = await client.search(
+            query=query,
+            topic=topic,
+            search_depth=search_depth,
+            max_results=max_results,
+            time_range=time_range,
+            include_images=include_images,
+            include_image_descriptions=include_image_descriptions,
+            chunks_per_source=chunks_per_source,
+            include_answer=include_answer,
+        )
+
+    results = [TavilySearchResult.from_tavily_sdk(r) for r in response.get("results", [])]
+
+    images: list[TavilyImage] | None = None
+    if include_images and response.get("images"):
+        images = [TavilyImage.from_tavily_sdk(img) for img in response.get("images", [])]
+
+    result_count = len(results)
+    answer = response.get("answer")
+    response_time = response.get("response_time", 0.0)
+
+    structured_content: dict = {
+        "query": response.get("query", query),
+        "results": [r.as_flat_dict() for r in results],
+        "resultCount": result_count,
+        "responseTime": response_time,
+    }
+
+    if answer:
+        structured_content["answer"] = answer
+
+    if images:
+        structured_content["images"] = [
+            {"url": img.url, "description": img.description} for img in images
+        ]
+
+    return ToolResult(
+        structured_content=structured_content,
+    )

datarobot_genai/langgraph/agent.py

@@ -11,9 +11,12 @@
 # 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 __future__ import annotations
+
 import abc
 import logging
 from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import cast
 
@@ -34,8 +37,6 @@ from langgraph.graph import MessagesState
 from langgraph.graph import StateGraph
 from langgraph.types import Command
 from openai.types.chat import CompletionCreateParams
-from ragas import MultiTurnSample
-from ragas.integrations.langgraph import convert_to_ragas_messages
 
 from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
@@ -44,6 +45,9 @@ from datarobot_genai.core.agents.base import extract_user_prompt_content
 from datarobot_genai.core.agents.base import is_streaming
 from datarobot_genai.langgraph.mcp import mcp_tools_context
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+
 logger = logging.getLogger(__name__)
 
 
@@ -337,5 +341,9 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
                 if v is not None:
                     messages.extend(v.get("messages", []))
         messages = [m for m in messages if not isinstance(m, ToolMessage)]
+        # Lazy import to reduce memory overhead when ragas is not used
+        from ragas import MultiTurnSample
+        from ragas.integrations.langgraph import convert_to_ragas_messages
+
         ragas_trace = convert_to_ragas_messages(messages)
         return MultiTurnSample(user_input=ragas_trace)

datarobot_genai/llama_index/__init__.py

@@ -3,8 +3,8 @@
 from datarobot_genai.core.mcp.common import MCPConfig
 
 from .agent import DataRobotLiteLLM
+from .agent import LlamaIndexAgent
 from .agent import create_pipeline_interactions_from_events
-from .base import LlamaIndexAgent
 from .mcp import load_mcp_tools
 
 __all__ = [

datarobot_genai/llama_index/agent.py

@@ -11,17 +11,32 @@
 # 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 __future__ import annotations
 
+import abc
+import inspect
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
+from typing import Any
 from typing import cast
 
 from llama_index.core.base.llms.types import LLMMetadata
+from llama_index.core.tools import BaseTool
 from llama_index.core.workflow import Event
 from llama_index.llms.litellm import LiteLLM
-from ragas import MultiTurnSample
-from ragas.integrations.llama_index import convert_to_ragas_messages
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolMessage
+from openai.types.chat import CompletionCreateParams
+
+from datarobot_genai.core.agents.base import BaseAgent
+from datarobot_genai.core.agents.base import InvokeReturn
+from datarobot_genai.core.agents.base import UsageMetrics
+from datarobot_genai.core.agents.base import default_usage_metrics
+from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
+
+from .mcp import load_mcp_tools
+
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
 
 
 class DataRobotLiteLLM(LiteLLM):
@@ -44,7 +59,271 @@ def create_pipeline_interactions_from_events(
 ) -> MultiTurnSample | None:
     if not events:
         return None
+    # Lazy import to reduce memory overhead when ragas is not used
+    from ragas import MultiTurnSample
+    from ragas.integrations.llama_index import convert_to_ragas_messages
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
+
     # convert_to_ragas_messages expects a list[Event]
     ragas_trace = convert_to_ragas_messages(list(events))
     ragas_messages = cast(list[HumanMessage | AIMessage | ToolMessage], ragas_trace)
     return MultiTurnSample(user_input=ragas_messages)
+
+
+class LlamaIndexAgent(BaseAgent[BaseTool], abc.ABC):
+    """Abstract base agent for LlamaIndex workflows."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._mcp_tools: list[Any] = []
+
+    def set_mcp_tools(self, tools: list[Any]) -> None:
+        """Set MCP tools for this agent."""
+        self._mcp_tools = tools
+
+    @property
+    def mcp_tools(self) -> list[Any]:
+        """Return the list of MCP tools available to this agent.
+
+        Subclasses can use this to wire tools into LlamaIndex agents during
+        workflow construction inside ``build_workflow``.
+        """
+        return self._mcp_tools
+
+    @abc.abstractmethod
+    def build_workflow(self) -> Any:
+        """Return an AgentWorkflow instance ready to run."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def extract_response_text(self, result_state: Any, events: list[Any]) -> str:
+        """Extract final response text from workflow state and/or events."""
+        raise NotImplementedError
+
+    def make_input_message(self, completion_create_params: CompletionCreateParams) -> str:
+        """Create an input string for the workflow from the user prompt."""
+        user_prompt_content = extract_user_prompt_content(completion_create_params)
+        return str(user_prompt_content)
+
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
+        """Run the LlamaIndex workflow with the provided completion parameters."""
+        input_message = self.make_input_message(completion_create_params)
+
+        # Load MCP tools (if configured) asynchronously before building workflow
+        mcp_tools = await load_mcp_tools(
+            authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
+        )
+        self.set_mcp_tools(mcp_tools)
+
+        # Preserve prior template startup print for CLI parity
+        try:
+            print(
+                "Running agent with user prompt:",
+                extract_user_prompt_content(completion_create_params),
+                flush=True,
+            )
+        except Exception:
+            # Printing is best-effort; proceed regardless
+            pass
+
+        workflow = self.build_workflow()
+        handler = workflow.run(user_msg=input_message)
+
+        usage_metrics: UsageMetrics = default_usage_metrics()
+
+        # Streaming parity with LangGraph: yield incremental deltas during event processing
+        if is_streaming(completion_create_params):
+
+            async def _gen() -> AsyncGenerator[tuple[str, MultiTurnSample | None, UsageMetrics]]:
+                events: list[Any] = []
+                current_agent_name: str | None = None
+                async for event in handler.stream_events():
+                    events.append(event)
+                    # Best-effort extraction of incremental text from LlamaIndex events
+                    delta: str | None = None
+                    # Agent switch banner if available on event
+                    try:
+                        if hasattr(event, "current_agent_name"):
+                            new_agent = getattr(event, "current_agent_name")
+                            if (
+                                isinstance(new_agent, str)
+                                and new_agent
+                                and new_agent != current_agent_name
+                            ):
+                                current_agent_name = new_agent
+                                # Print banner for agent switch (do not emit as streamed content)
+                                print("\n" + "=" * 50, flush=True)
+                                print(f"🤖 Agent: {current_agent_name}", flush=True)
+                                print("=" * 50 + "\n", flush=True)
+                    except Exception:
+                        pass
+
+                    try:
+                        if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
+                            delta = getattr(event, "delta")
+                        # Some event types may carry incremental text under "text" or similar
+                        elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
+                            delta = getattr(event, "text")
+                    except Exception:
+                        # Ignore malformed events and continue
+                        delta = None
+
+                    if delta:
+                        # Yield token/content delta with current (accumulated) usage metrics
+                        yield delta, None, usage_metrics
+
+                    # Best-effort debug/event messages printed to CLI (do not stream as content)
+                    try:
+                        event_type = type(event).__name__
+                        if event_type == "AgentInput" and hasattr(event, "input"):
+                            print("📥 Input:", getattr(event, "input"), flush=True)
+                        elif event_type == "AgentOutput":
+                            # Output content
+                            resp = getattr(event, "response", None)
+                            if (
+                                resp is not None
+                                and hasattr(resp, "content")
+                                and getattr(resp, "content")
+                            ):
+                                print("📤 Output:", getattr(resp, "content"), flush=True)
+                            # Planned tool calls
+                            tcalls = getattr(event, "tool_calls", None)
+                            if isinstance(tcalls, list) and tcalls:
+                                names = []
+                                for c in tcalls:
+                                    try:
+                                        nm = getattr(c, "tool_name", None) or (
+                                            c.get("tool_name") if isinstance(c, dict) else None
+                                        )
+                                        if nm:
+                                            names.append(str(nm))
+                                    except Exception:
+                                        pass
+                                if names:
+                                    print("🛠️  Planning to use tools:", names, flush=True)
+                        elif event_type == "ToolCallResult":
+                            tname = getattr(event, "tool_name", None)
+                            tkwargs = getattr(event, "tool_kwargs", None)
+                            tout = getattr(event, "tool_output", None)
+                            print(f"🔧 Tool Result ({tname}):", flush=True)
+                            print(f"  Arguments: {tkwargs}", flush=True)
+                            print(f"  Output: {tout}", flush=True)
+                        elif event_type == "ToolCall":
+                            tname = getattr(event, "tool_name", None)
+                            tkwargs = getattr(event, "tool_kwargs", None)
+                            print(f"🔨 Calling Tool: {tname}", flush=True)
+                            print(f"  With arguments: {tkwargs}", flush=True)
+                    except Exception:
+                        # Ignore best-effort debug rendering errors
+                        pass
+
+                # After streaming completes, build final interactions and finish chunk
+                # Extract state from workflow context (supports sync/async get or attribute)
+                state = None
+                ctx = getattr(handler, "ctx", None)
+                try:
+                    if ctx is not None:
+                        get = getattr(ctx, "get", None)
+                        if callable(get):
+                            result = get("state")
+                            state = await result if inspect.isawaitable(result) else result
+                        elif hasattr(ctx, "state"):
+                            state = getattr(ctx, "state")
+                except (AttributeError, TypeError):
+                    state = None
+
+                # Run subclass-defined response extraction (not streamed) for completeness
+                _ = self.extract_response_text(state, events)
+
+                pipeline_interactions = create_pipeline_interactions_from_events(events)
+                # Final empty chunk indicates end of stream, carrying interactions and usage
+                yield "", pipeline_interactions, usage_metrics
+
+            return _gen()
+
+        # Non-streaming path: run to completion, emit debug prints, then return final response
+        events: list[Any] = []
+        current_agent_name: str | None = None
+        async for event in handler.stream_events():
+            events.append(event)
+
+            # Replicate prior template CLI prints for non-streaming mode
+            try:
+                if hasattr(event, "current_agent_name"):
+                    new_agent = getattr(event, "current_agent_name")
+                    if isinstance(new_agent, str) and new_agent and new_agent != current_agent_name:
+                        current_agent_name = new_agent
+                        print(f"\n{'=' * 50}", flush=True)
+                        print(f"🤖 Agent: {current_agent_name}", flush=True)
+                        print(f"{'=' * 50}\n", flush=True)
+            except Exception:
+                pass
+
+            try:
+                if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
+                    print(getattr(event, "delta"), end="", flush=True)
+                elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
+                    print(getattr(event, "text"), end="", flush=True)
+                else:
+                    event_type = type(event).__name__
+                    if event_type == "AgentInput" and hasattr(event, "input"):
+                        print("📥 Input:", getattr(event, "input"), flush=True)
+                    elif event_type == "AgentOutput":
+                        resp = getattr(event, "response", None)
+                        if (
+                            resp is not None
+                            and hasattr(resp, "content")
+                            and getattr(resp, "content")
+                        ):
+                            print("📤 Output:", getattr(resp, "content"), flush=True)
+                        tcalls = getattr(event, "tool_calls", None)
+                        if isinstance(tcalls, list) and tcalls:
+                            names: list[str] = []
+                            for c in tcalls:
+                                try:
+                                    nm = getattr(c, "tool_name", None) or (
+                                        c.get("tool_name") if isinstance(c, dict) else None
+                                    )
+                                    if nm:
+                                        names.append(str(nm))
+                                except Exception:
+                                    pass
+                            if names:
+                                print("🛠️  Planning to use tools:", names, flush=True)
+                    elif event_type == "ToolCallResult":
+                        tname = getattr(event, "tool_name", None)
+                        tkwargs = getattr(event, "tool_kwargs", None)
+                        tout = getattr(event, "tool_output", None)
+                        print(f"🔧 Tool Result ({tname}):", flush=True)
+                        print(f"  Arguments: {tkwargs}", flush=True)
+                        print(f"  Output: {tout}", flush=True)
+                    elif event_type == "ToolCall":
+                        tname = getattr(event, "tool_name", None)
+                        tkwargs = getattr(event, "tool_kwargs", None)
+                        print(f"🔨 Calling Tool: {tname}", flush=True)
+                        print(f"  With arguments: {tkwargs}", flush=True)
+            except Exception:
+                # Best-effort debug printing; continue on errors
+                pass
+
+        # Extract state from workflow context (supports sync/async get or attribute)
+        state = None
+        ctx = getattr(handler, "ctx", None)
+        try:
+            if ctx is not None:
+                get = getattr(ctx, "get", None)
+                if callable(get):
+                    result = get("state")
+                    state = await result if inspect.isawaitable(result) else result
+                elif hasattr(ctx, "state"):
+                    state = getattr(ctx, "state")
+        except (AttributeError, TypeError):
+            state = None
+        response_text = self.extract_response_text(state, events)
+
+        pipeline_interactions = create_pipeline_interactions_from_events(events)
+
+        return response_text, pipeline_interactions, usage_metrics

datarobot_genai/nat/agent.py

@@ -11,9 +11,12 @@
 # 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 __future__ import annotations
+
 import asyncio
 import logging
 from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 from typing import Any
 
 from nat.builder.context import Context
@@ -23,10 +26,6 @@ from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.utils.type_utils import StrPath
 from openai.types.chat import CompletionCreateParams
-from ragas import MultiTurnSample
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolMessage
 
 from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
@@ -36,13 +35,22 @@ from datarobot_genai.core.agents.base import is_streaming
 from datarobot_genai.core.mcp.common import MCPConfig
 from datarobot_genai.nat.helpers import load_workflow
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+
 logger = logging.getLogger(__name__)
 
 
 def convert_to_ragas_messages(
     steps: list[IntermediateStep],
-) -> list[HumanMessage | AIMessage | ToolMessage]:
-    def _to_ragas(step: IntermediateStep) -> HumanMessage | AIMessage | ToolMessage:
+) -> list[HumanMessage | AIMessage]:
+    # Lazy import to reduce memory overhead when ragas is not used
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+
+    def _to_ragas(step: IntermediateStep) -> HumanMessage | AIMessage:
         if step.event_type == IntermediateStepType.LLM_START:
             return HumanMessage(content=_parse(step.data.input))
         elif step.event_type == IntermediateStepType.LLM_END:
@@ -78,6 +86,9 @@ def create_pipeline_interactions_from_steps(
 ) -> MultiTurnSample | None:
     if not steps:
         return None
+    # Lazy import to reduce memory overhead when ragas is not used
+    from ragas import MultiTurnSample
+
     ragas_trace = convert_to_ragas_messages(steps)
     return MultiTurnSample(user_input=ragas_trace)
 

{datarobot_genai-0.2.39.dist-info → datarobot_genai-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.39
+Version: 0.3.1
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -43,9 +43,11 @@ Requires-Dist: opentelemetry-api<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: opentelemetry-exporter-otlp-proto-http<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: opentelemetry-exporter-otlp<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: opentelemetry-sdk<2.0.0,>=1.22.0; extra == 'drmcp'
+Requires-Dist: perplexityai<1.0,>=0.27; extra == 'drmcp'
 Requires-Dist: pydantic-settings<3.0.0,>=2.1.0; extra == 'drmcp'
 Requires-Dist: pydantic<3.0.0,>=2.6.1; extra == 'drmcp'
 Requires-Dist: python-dotenv<2.0.0,>=1.1.0; extra == 'drmcp'
+Requires-Dist: tavily-python<1.0.0,>=0.7.20; extra == 'drmcp'
 Provides-Extra: langgraph
 Requires-Dist: langchain-mcp-adapters<0.2.0,>=0.1.12; extra == 'langgraph'
 Requires-Dist: langgraph-prebuilt<0.7.0,>=0.2.3; extra == 'langgraph'