vanna 0.7.9__py3-none-any.whl → 2.0.0rc1__py3-none-any.whl

vanna/integrations/__init__.py

@@ -0,0 +1,17 @@
+"""
+Integrations module.
+
+This package contains concrete implementations of core abstractions and capabilities.
+"""
+
+from .local import MemoryConversationStore
+from .mock import MockLlmService
+from .plotly import PlotlyChartGenerator
+from .sqlite import SqliteRunner
+
+__all__ = [
+    "MockLlmService",
+    "MemoryConversationStore",
+    "SqliteRunner",
+    "PlotlyChartGenerator",
+]

vanna/integrations/anthropic/__init__.py

@@ -0,0 +1,9 @@
+"""
+Anthropic integration.
+
+This module provides Anthropic LLM service implementation.
+"""
+
+from .llm import AnthropicLlmService
+
+__all__ = ["AnthropicLlmService"]

vanna/integrations/anthropic/llm.py

@@ -0,0 +1,270 @@
+"""
+Anthropic LLM service implementation.
+
+Implements the LlmService interface using Anthropic's Messages API
+(anthropic>=0.8.0). Supports non-streaming and streaming text output.
+Tool-calls (tool_use blocks) are surfaced at the end of a stream or after a
+non-streaming call as ToolCall entries.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+from vanna.core.llm import (
+    LlmService,
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+)
+from vanna.core.tool import ToolCall, ToolSchema
+
+
+class AnthropicLlmService(LlmService):
+    """Anthropic Messages-backed LLM service.
+
+    Args:
+        model: Anthropic model name (e.g., "claude-sonnet-4-5", "claude-opus-4").
+            Defaults to "claude-sonnet-4-5". Can also be set via ANTHROPIC_MODEL env var.
+        api_key: API key; falls back to env `ANTHROPIC_API_KEY`.
+        base_url: Optional custom base URL; env `ANTHROPIC_BASE_URL` if unset.
+        extra_client_kwargs: Extra kwargs forwarded to `anthropic.Anthropic()`.
+    """
+
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        **extra_client_kwargs: Any,
+    ) -> None:
+        try:
+            import anthropic
+        except Exception as e:  # pragma: no cover
+            raise ImportError(
+                "anthropic package is required. Install with: pip install 'vanna[anthropic]'"
+            ) from e
+
+        # Model selection - use environment variable or default
+        self.model = model or os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5")
+        api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
+        base_url = base_url or os.getenv("ANTHROPIC_BASE_URL")
+
+        client_kwargs: Dict[str, Any] = {**extra_client_kwargs}
+        if api_key:
+            client_kwargs["api_key"] = api_key
+        if base_url:
+            client_kwargs["base_url"] = base_url
+
+        self._client = anthropic.Anthropic(**client_kwargs)
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Send a non-streaming request to Anthropic and return the response."""
+        payload = self._build_payload(request)
+
+        resp = self._client.messages.create(**payload)
+
+        logger.info(f"Anthropic response: {resp}")
+
+        text_content, tool_calls = self._parse_message_content(resp)
+
+        usage: Dict[str, int] = {}
+        if getattr(resp, "usage", None):
+            try:
+                usage = {
+                    "input_tokens": int(resp.usage.input_tokens),
+                    "output_tokens": int(resp.usage.output_tokens),
+                }
+            except Exception:
+                pass
+
+        return LlmResponse(
+            content=text_content or None,
+            tool_calls=tool_calls or None,
+            finish_reason=getattr(resp, "stop_reason", None),
+            usage=usage or None,
+        )
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Stream a request to Anthropic.
+
+        Yields text chunks as they arrive. Emits tool-calls at the end by
+        inspecting the final message.
+        """
+        payload = self._build_payload(request)
+
+        logger.info(f"Anthropic streaming payload: {payload}")
+
+        # SDK provides a streaming context manager with a text_stream iterator.
+        with self._client.messages.stream(**payload) as stream:
+            for text in stream.text_stream:
+                if text:
+                    yield LlmStreamChunk(content=text)
+
+            final = stream.get_final_message()
+            logger.info(f"Anthropic stream response: {final}")
+            _, tool_calls = self._parse_message_content(final)
+            if tool_calls:
+                yield LlmStreamChunk(
+                    tool_calls=tool_calls,
+                    finish_reason=getattr(final, "stop_reason", None),
+                )
+            else:
+                yield LlmStreamChunk(
+                    finish_reason=getattr(final, "stop_reason", None) or "stop"
+                )
+
+    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
+        """Basic validation of tool schemas for Anthropic."""
+        errors: List[str] = []
+        for t in tools:
+            if not t.name:
+                errors.append("Tool name is required")
+        return errors
+
+    # Internal helpers
+    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
+        # Anthropic requires messages content as list of content blocks per message
+        # We need to group consecutive tool messages into single user messages
+        messages: List[Dict[str, Any]] = []
+        i = 0
+
+        while i < len(request.messages):
+            m = request.messages[i]
+
+            if m.role == "tool":
+                # Group consecutive tool messages into one user message
+                tool_content_blocks = []
+                while i < len(request.messages) and request.messages[i].role == "tool":
+                    tool_msg = request.messages[i]
+                    if tool_msg.tool_call_id:
+                        tool_content_blocks.append(
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": tool_msg.tool_call_id,
+                                "content": tool_msg.content,
+                            }
+                        )
+                    i += 1
+
+                if tool_content_blocks:
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": tool_content_blocks,
+                        }
+                    )
+            else:
+                # Handle non-tool messages normally
+                content_blocks = []
+
+                # Handle text content - only add if not empty
+                if m.content and m.content.strip():
+                    content_blocks.append({"type": "text", "text": m.content})
+
+                # Handle tool_calls for assistant messages (convert to tool_use blocks)
+                if m.role == "assistant" and m.tool_calls:
+                    for tc in m.tool_calls:
+                        content_blocks.append(
+                            {
+                                "type": "tool_use",
+                                "id": tc.id,
+                                "name": tc.name,
+                                "input": tc.arguments,  # type: ignore[dict-item]
+                            }
+                        )
+
+                # Ensure we have at least one content block for text messages
+                if not content_blocks and m.role in {"user", "assistant"}:
+                    content_blocks.append({"type": "text", "text": m.content or ""})
+
+                if content_blocks:
+                    role = m.role if m.role in {"user", "assistant"} else "user"
+                    messages.append(
+                        {
+                            "role": role,
+                            "content": content_blocks,
+                        }
+                    )
+
+                i += 1
+
+        tools_payload: Optional[List[Dict[str, Any]]] = None
+        if request.tools:
+            tools_payload = [
+                {
+                    "name": t.name,
+                    "description": t.description,
+                    "input_schema": t.parameters,
+                }
+                for t in request.tools
+            ]
+
+        payload: Dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            # Anthropic requires max_tokens; default if not provided
+            "max_tokens": request.max_tokens if request.max_tokens is not None else 512,
+            "temperature": request.temperature,
+        }
+        if tools_payload:
+            payload["tools"] = tools_payload
+            payload["tool_choice"] = {"type": "auto"}
+
+        # Add system prompt if provided
+        if request.system_prompt:
+            payload["system"] = request.system_prompt
+
+        return payload
+
+    def _parse_message_content(self, msg: Any) -> Tuple[str, List[ToolCall]]:
+        text_parts: List[str] = []
+        tool_calls: List[ToolCall] = []
+
+        content_list = getattr(msg, "content", []) or []
+        for block in content_list:
+            btype = getattr(block, "type", None) or (
+                block.get("type") if isinstance(block, dict) else None
+            )
+            if btype == "text":
+                # SDK returns block.text for typed object; dict uses {"text": ...}
+                text = getattr(block, "text", None)
+                if text is None and isinstance(block, dict):
+                    text = block.get("text")
+                if text:
+                    text_parts.append(str(text))
+            elif btype == "tool_use":
+                # Tool call with name and input
+                name = getattr(block, "name", None) or (
+                    block.get("name") if isinstance(block, dict) else None
+                )
+                tc_id = getattr(block, "id", None) or (
+                    block.get("id") if isinstance(block, dict) else None
+                )
+                input_data = getattr(block, "input", None) or (
+                    block.get("input") if isinstance(block, dict) else None
+                )
+                if name:
+                    try:
+                        # input_data should be a dict already
+                        args = (
+                            input_data
+                            if isinstance(input_data, dict)
+                            else {"_raw": input_data}
+                        )
+                    except Exception:
+                        args = {"_raw": str(input_data)}
+                    tool_calls.append(
+                        ToolCall(
+                            id=str(tc_id or "tool_call"), name=str(name), arguments=args
+                        )
+                    )
+
+        text_content = "".join(text_parts)
+        return text_content, tool_calls

vanna/integrations/azureopenai/__init__.py

@@ -0,0 +1,9 @@
+"""
+Azure OpenAI integration.
+
+This module provides Azure OpenAI LLM service implementations.
+"""
+
+from .llm import AzureOpenAILlmService
+
+__all__ = ["AzureOpenAILlmService"]

vanna/integrations/azureopenai/llm.py

@@ -0,0 +1,329 @@
+"""
+Azure OpenAI LLM service implementation.
+
+Provides an `LlmService` backed by Azure OpenAI Chat Completions (openai>=1.0.0)
+with support for streaming, deployment-scoped models, and Azure-specific
+authentication flows.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional, Set
+
+from vanna.core.llm import (
+    LlmService,
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+)
+from vanna.core.tool import ToolCall, ToolSchema
+
+
+# Models that don't support temperature and other sampling parameters
+REASONING_MODELS: Set[str] = {
+    "o1",
+    "o1-mini",
+    "o1-preview",
+    "o3-mini",
+    "gpt-5",
+    "gpt-5-mini",
+    "gpt-5-nano",
+    "gpt-5-pro",
+    "gpt-5-codex",
+}
+
+
+def _is_reasoning_model(model: str) -> bool:
+    """Return True when the deployment targets a reasoning-only model."""
+    model_lower = model.lower()
+    return any(reasoning_model in model_lower for reasoning_model in REASONING_MODELS)
+
+
+class AzureOpenAILlmService(LlmService):
+    """Azure OpenAI Chat Completions-backed LLM service.
+
+    Wraps `openai.AzureOpenAI` so Vanna can talk to deployment-scoped models
+    and either API key or Microsoft Entra ID authentication.
+
+    Args:
+        model: Deployment name in Azure OpenAI (required).
+        api_key: API key; falls back to `AZURE_OPENAI_API_KEY`.
+        azure_endpoint: Azure OpenAI endpoint URL; falls back to
+            `AZURE_OPENAI_ENDPOINT`.
+        api_version: API version; defaults to "2024-10-21" or
+            `AZURE_OPENAI_API_VERSION`.
+        azure_ad_token_provider: Optional bearer token provider for Entra ID.
+        **extra_client_kwargs: Additional keyword arguments forwarded to the
+            underlying client.
+    """
+
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        azure_endpoint: Optional[str] = None,
+        api_version: Optional[str] = None,
+        azure_ad_token_provider: Optional[Any] = None,
+        **extra_client_kwargs: Any,
+    ) -> None:
+        try:
+            from openai import AzureOpenAI
+        except Exception as e:  # pragma: no cover
+            raise ImportError(
+                "openai package is required. Install with: pip install 'vanna[azureopenai]' "
+                "or 'pip install openai'"
+            ) from e
+
+        # Model/deployment name is required for Azure OpenAI
+        self.model = model or os.getenv("AZURE_OPENAI_MODEL")
+        if not self.model:
+            raise ValueError(
+                "model parameter (deployment name) is required for Azure OpenAI. "
+                "Provide it as argument or set AZURE_OPENAI_MODEL environment variable."
+            )
+
+        # Azure endpoint is required
+        azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
+        if not azure_endpoint:
+            raise ValueError(
+                "azure_endpoint is required for Azure OpenAI. "
+                "Provide it as argument or set AZURE_OPENAI_ENDPOINT environment variable."
+            )
+
+        # API version - use latest stable GA version by default
+        api_version = api_version or os.getenv("AZURE_OPENAI_API_VERSION", "2024-10-21")
+
+        # Build client kwargs
+        client_kwargs: Dict[str, Any] = {
+            "azure_endpoint": azure_endpoint,
+            "api_version": api_version,
+            **extra_client_kwargs,
+        }
+
+        # Authentication: prefer Azure AD token provider, fallback to API key
+        if azure_ad_token_provider is not None:
+            client_kwargs["azure_ad_token_provider"] = azure_ad_token_provider
+        else:
+            api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "Authentication required: provide either api_key or azure_ad_token_provider. "
+                    "API key can also be set via AZURE_OPENAI_API_KEY environment variable."
+                )
+            client_kwargs["api_key"] = api_key
+
+        self._client = AzureOpenAI(**client_kwargs)
+        self._is_reasoning_model = _is_reasoning_model(self.model)
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Send a non-streaming request to Azure OpenAI and return the response."""
+        payload = self._build_payload(request)
+
+        # Call the API synchronously; this function is async but we can block here.
+        resp = self._client.chat.completions.create(**payload, stream=False)
+
+        if not resp.choices:
+            return LlmResponse(content=None, tool_calls=None, finish_reason=None)
+
+        choice = resp.choices[0]
+        content: Optional[str] = getattr(choice.message, "content", None)
+        tool_calls = self._extract_tool_calls_from_message(choice.message)
+
+        usage: Dict[str, int] = {}
+        if getattr(resp, "usage", None):
+            usage = {
+                k: int(v)
+                for k, v in {
+                    "prompt_tokens": getattr(resp.usage, "prompt_tokens", 0),
+                    "completion_tokens": getattr(resp.usage, "completion_tokens", 0),
+                    "total_tokens": getattr(resp.usage, "total_tokens", 0),
+                }.items()
+            }
+
+        return LlmResponse(
+            content=content,
+            tool_calls=tool_calls or None,
+            finish_reason=getattr(choice, "finish_reason", None),
+            usage=usage or None,
+        )
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """
+        Stream a request to Azure OpenAI.
+
+        Emits `LlmStreamChunk` for textual deltas as they arrive. Tool-calls are
+        accumulated and emitted in a final chunk when the stream ends.
+        """
+        payload = self._build_payload(request)
+
+        # Synchronous streaming iterator; iterate within async context.
+        stream = self._client.chat.completions.create(**payload, stream=True)
+
+        # Builders for streamed tool-calls (index -> partial)
+        tc_builders: Dict[int, Dict[str, Optional[str]]] = {}
+        last_finish: Optional[str] = None
+
+        for event in stream:
+            if not getattr(event, "choices", None):
+                continue
+
+            choice = event.choices[0]
+            delta = getattr(choice, "delta", None)
+            if delta is None:
+                # Some SDK versions use `event.choices[0].message` on the final packet
+                last_finish = getattr(choice, "finish_reason", last_finish)
+                continue
+
+            # Text content
+            content_piece: Optional[str] = getattr(delta, "content", None)
+            if content_piece:
+                yield LlmStreamChunk(content=content_piece)
+
+            # Tool calls (streamed)
+            streamed_tool_calls = getattr(delta, "tool_calls", None)
+            if streamed_tool_calls:
+                for tc in streamed_tool_calls:
+                    idx = getattr(tc, "index", 0) or 0
+                    b = tc_builders.setdefault(
+                        idx, {"id": None, "name": None, "arguments": ""}
+                    )
+                    if getattr(tc, "id", None):
+                        b["id"] = tc.id
+                    fn = getattr(tc, "function", None)
+                    if fn is not None:
+                        if getattr(fn, "name", None):
+                            b["name"] = fn.name
+                        if getattr(fn, "arguments", None):
+                            b["arguments"] = (b["arguments"] or "") + fn.arguments
+
+            last_finish = getattr(choice, "finish_reason", last_finish)
+
+        # Emit final tool-calls chunk if any
+        final_tool_calls: List[ToolCall] = []
+        for b in tc_builders.values():
+            if not b.get("name"):
+                continue
+            args_raw = b.get("arguments") or "{}"
+            try:
+                loaded = json.loads(args_raw)
+                if isinstance(loaded, dict):
+                    args_dict: Dict[str, Any] = loaded
+                else:
+                    args_dict = {"args": loaded}
+            except Exception:
+                args_dict = {"_raw": args_raw}
+            final_tool_calls.append(
+                ToolCall(
+                    id=b.get("id") or "tool_call",
+                    name=b["name"] or "tool",
+                    arguments=args_dict,
+                )
+            )
+
+        if final_tool_calls:
+            yield LlmStreamChunk(tool_calls=final_tool_calls, finish_reason=last_finish)
+        else:
+            # Still emit a terminal chunk to signal completion
+            yield LlmStreamChunk(finish_reason=last_finish or "stop")
+
+    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
+        """Validate tool schemas. Returns a list of error messages."""
+        errors: List[str] = []
+        # Basic checks; Azure OpenAI will enforce further validation server-side.
+        for t in tools:
+            if not t.name or len(t.name) > 64:
+                errors.append(f"Invalid tool name: {t.name!r}")
+        return errors
+
+    # Internal helpers
+    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
+        """Build the API payload from LlmRequest."""
+        messages: List[Dict[str, Any]] = []
+
+        # Add system prompt as first message if provided
+        if request.system_prompt:
+            messages.append({"role": "system", "content": request.system_prompt})
+
+        for m in request.messages:
+            msg: Dict[str, Any] = {"role": m.role, "content": m.content}
+            if m.role == "tool" and m.tool_call_id:
+                msg["tool_call_id"] = m.tool_call_id
+            elif m.role == "assistant" and m.tool_calls:
+                # Convert tool calls to OpenAI format
+                tool_calls_payload = []
+                for tc in m.tool_calls:
+                    tool_calls_payload.append(
+                        {
+                            "id": tc.id,
+                            "type": "function",
+                            "function": {
+                                "name": tc.name,
+                                "arguments": json.dumps(tc.arguments),
+                            },
+                        }
+                    )
+                msg["tool_calls"] = tool_calls_payload
+            messages.append(msg)
+
+        tools_payload: Optional[List[Dict[str, Any]]] = None
+        if request.tools:
+            tools_payload = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": t.name,
+                        "description": t.description,
+                        "parameters": t.parameters,
+                    },
+                }
+                for t in request.tools
+            ]
+
+        payload: Dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+        }
+
+        # Add temperature only for non-reasoning models
+        # Reasoning models (GPT-5, o1, o3-mini) don't support temperature parameter
+        if not self._is_reasoning_model:
+            payload["temperature"] = request.temperature
+
+        if request.max_tokens is not None:
+            payload["max_tokens"] = request.max_tokens
+
+        if tools_payload:
+            payload["tools"] = tools_payload
+            payload["tool_choice"] = "auto"
+
+        return payload
+
+    def _extract_tool_calls_from_message(self, message: Any) -> List[ToolCall]:
+        """Extract tool calls from OpenAI message object."""
+        tool_calls: List[ToolCall] = []
+        raw_tool_calls = getattr(message, "tool_calls", None) or []
+        for tc in raw_tool_calls:
+            fn = getattr(tc, "function", None)
+            if not fn:
+                continue
+            args_raw = getattr(fn, "arguments", "{}")
+            try:
+                loaded = json.loads(args_raw)
+                if isinstance(loaded, dict):
+                    args_dict: Dict[str, Any] = loaded
+                else:
+                    args_dict = {"args": loaded}
+            except Exception:
+                args_dict = {"_raw": args_raw}
+            tool_calls.append(
+                ToolCall(
+                    id=getattr(tc, "id", "tool_call"),
+                    name=getattr(fn, "name", "tool"),
+                    arguments=args_dict,
+                )
+            )
+        return tool_calls

vanna/integrations/azuresearch/__init__.py

@@ -0,0 +1,7 @@
+"""
+Azure AI Search integration for Vanna Agents.
+"""
+
+from .agent_memory import AzureAISearchAgentMemory
+
+__all__ = ["AzureAISearchAgentMemory"]