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

vanna/integrations/ollama/__init__.py

@@ -0,0 +1,7 @@
+"""
+Ollama integration for Vanna Agents.
+"""
+
+from .llm import OllamaLlmService
+
+__all__ = ["OllamaLlmService"]

vanna/integrations/ollama/llm.py

@@ -0,0 +1,252 @@
+"""
+Ollama LLM service implementation.
+
+This module provides an implementation of the LlmService interface backed by
+Ollama's local LLM API. It supports non-streaming responses and streaming
+of text content. Tool calling support depends on the Ollama model being used.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+from vanna.core.llm import (
+    LlmService,
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+)
+from vanna.core.tool import ToolCall, ToolSchema
+
+
+class OllamaLlmService(LlmService):
+    """Ollama-backed LLM service for local model inference.
+
+    Args:
+        model: Ollama model name (e.g., "gpt-oss:20b").
+        host: Ollama server URL; defaults to "http://localhost:11434" or env `OLLAMA_HOST`.
+        timeout: Request timeout in seconds; defaults to 240.
+        num_ctx: Context window size; defaults to 8192.
+        temperature: Sampling temperature; defaults to 0.7.
+        extra_options: Additional options passed to Ollama (e.g., num_predict, top_k, top_p).
+    """
+
+    def __init__(
+        self,
+        model: str,
+        host: Optional[str] = None,
+        timeout: float = 240.0,
+        num_ctx: int = 8192,
+        temperature: float = 0.7,
+        **extra_options: Any,
+    ) -> None:
+        try:
+            import ollama
+        except ImportError as e:
+            raise ImportError(
+                "ollama package is required. Install with: pip install 'vanna[ollama]' or pip install ollama"
+            ) from e
+
+        if not model:
+            raise ValueError("model parameter is required for Ollama")
+
+        self.model = model
+        self.host = host or os.getenv("OLLAMA_HOST", "http://localhost:11434")
+        self.timeout = timeout
+        self.num_ctx = num_ctx
+        self.temperature = temperature
+        self.extra_options = extra_options
+
+        # Create Ollama client
+        self._client = ollama.Client(host=self.host, timeout=timeout)
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Send a non-streaming request to Ollama and return the response."""
+        payload = self._build_payload(request)
+
+        # Call the Ollama API
+        try:
+            resp = self._client.chat(**payload)
+        except Exception as e:
+            raise RuntimeError(f"Ollama request failed: {str(e)}") from e
+
+        # Extract message from response
+        message = resp.get("message", {})
+        content = message.get("content")
+        tool_calls = self._extract_tool_calls_from_message(message)
+
+        # Extract usage information if available
+        usage: Dict[str, int] = {}
+        if "prompt_eval_count" in resp or "eval_count" in resp:
+            usage = {
+                "prompt_tokens": resp.get("prompt_eval_count", 0),
+                "completion_tokens": resp.get("eval_count", 0),
+                "total_tokens": resp.get("prompt_eval_count", 0)
+                + resp.get("eval_count", 0),
+            }
+
+        return LlmResponse(
+            content=content,
+            tool_calls=tool_calls or None,
+            finish_reason=resp.get("done_reason")
+            or ("stop" if resp.get("done") else None),
+            usage=usage or None,
+        )
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Stream a request to Ollama.
+
+        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)
+
+        # Ollama streaming
+        try:
+            stream = self._client.chat(**payload, stream=True)
+        except Exception as e:
+            raise RuntimeError(f"Ollama streaming request failed: {str(e)}") from e
+
+        # Accumulate tool calls if present
+        accumulated_tool_calls: List[ToolCall] = []
+        last_finish: Optional[str] = None
+
+        for chunk in stream:
+            message = chunk.get("message", {})
+
+            # Yield text content
+            content = message.get("content")
+            if content:
+                yield LlmStreamChunk(content=content)
+
+            # Accumulate tool calls
+            tool_calls = self._extract_tool_calls_from_message(message)
+            if tool_calls:
+                accumulated_tool_calls.extend(tool_calls)
+
+            # Track finish reason
+            if chunk.get("done"):
+                last_finish = chunk.get("done_reason", "stop")
+
+        # Emit final chunk with tool calls if any
+        if accumulated_tool_calls:
+            yield LlmStreamChunk(
+                tool_calls=accumulated_tool_calls, finish_reason=last_finish or "stop"
+            )
+        else:
+            # Emit 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 validation; Ollama model support for tools varies
+        for t in tools:
+            if not t.name:
+                errors.append(f"Tool must have a name")
+            if not t.description:
+                errors.append(f"Tool '{t.name}' should have a description")
+        return errors
+
+    # Internal helpers
+    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
+        """Build the Ollama chat 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})
+
+        # Convert messages to Ollama format
+        for m in request.messages:
+            msg: Dict[str, Any] = {"role": m.role, "content": m.content or ""}
+
+            # Handle tool calls in assistant messages
+            if m.role == "assistant" and m.tool_calls:
+                # Some Ollama models support tool_calls in message
+                tool_calls_payload = []
+                for tc in m.tool_calls:
+                    tool_calls_payload.append(
+                        {"function": {"name": tc.name, "arguments": tc.arguments}}
+                    )
+                msg["tool_calls"] = tool_calls_payload
+
+            messages.append(msg)
+
+        # Build tools array if tools are provided
+        tools_payload: Optional[List[Dict[str, Any]]] = None
+        if request.tools:
+            tools_payload = []
+            for t in request.tools:
+                tools_payload.append(
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": t.name,
+                            "description": t.description,
+                            "parameters": t.parameters,
+                        },
+                    }
+                )
+
+        # Build options
+        options: Dict[str, Any] = {
+            "num_ctx": self.num_ctx,
+            "temperature": self.temperature,
+            **self.extra_options,
+        }
+
+        # Build final payload
+        payload: Dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            "options": options,
+        }
+
+        # Add tools if provided (note: not all Ollama models support tools)
+        if tools_payload:
+            payload["tools"] = tools_payload
+
+        return payload
+
+    def _extract_tool_calls_from_message(
+        self, message: Dict[str, Any]
+    ) -> List[ToolCall]:
+        """Extract tool calls from Ollama message."""
+        tool_calls: List[ToolCall] = []
+
+        # Check for tool_calls in message
+        raw_tool_calls = message.get("tool_calls", [])
+        if not raw_tool_calls:
+            return tool_calls
+
+        for idx, tc in enumerate(raw_tool_calls):
+            fn = tc.get("function", {})
+            name = fn.get("name")
+            if not name:
+                continue
+
+            # Parse arguments
+            arguments = fn.get("arguments", {})
+            if isinstance(arguments, str):
+                try:
+                    arguments = json.loads(arguments)
+                except Exception:
+                    arguments = {"_raw": arguments}
+
+            if not isinstance(arguments, dict):
+                arguments = {"args": arguments}
+
+            tool_calls.append(
+                ToolCall(
+                    id=tc.get("id", f"tool_call_{idx}"),
+                    name=name,
+                    arguments=arguments,
+                )
+            )
+
+        return tool_calls

vanna/integrations/openai/__init__.py

@@ -0,0 +1,10 @@
+"""
+OpenAI integration.
+
+This module provides OpenAI LLM service implementations.
+"""
+
+from .llm import OpenAILlmService
+from .responses import OpenAIResponsesService
+
+__all__ = ["OpenAILlmService", "OpenAIResponsesService"]

vanna/integrations/openai/llm.py

@@ -0,0 +1,267 @@
+"""
+OpenAI LLM service implementation.
+
+This module provides an implementation of the LlmService interface backed by
+OpenAI's Chat Completions API (openai>=1.0.0). It supports non-streaming
+responses and best-effort streaming of text content. Tool/function calling is
+passed through when tools are provided, but full tool-call conversation
+round-tripping may require adding assistant tool-call messages to the
+conversation upstream.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional, cast
+
+from vanna.core.llm import (
+    LlmService,
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+)
+from vanna.core.tool import ToolCall, ToolSchema
+
+
+class OpenAILlmService(LlmService):
+    """OpenAI Chat Completions-backed LLM service.
+
+    Args:
+        model: OpenAI model name (e.g., "gpt-5").
+        api_key: API key; falls back to env `OPENAI_API_KEY`.
+        organization: Optional org; env `OPENAI_ORG` if unset.
+        base_url: Optional custom base URL; env `OPENAI_BASE_URL` if unset.
+        extra_client_kwargs: Extra kwargs forwarded to `openai.OpenAI()`.
+    """
+
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        organization: Optional[str] = None,
+        base_url: Optional[str] = None,
+        **extra_client_kwargs: Any,
+    ) -> None:
+        try:
+            from openai import OpenAI
+        except Exception as e:  # pragma: no cover - import-time error surface
+            raise ImportError(
+                "openai package is required. Install with: pip install 'vanna[openai]'"
+            ) from e
+
+        self.model = model or os.getenv("OPENAI_MODEL", "gpt-5")
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        organization = organization or os.getenv("OPENAI_ORG")
+        base_url = base_url or os.getenv("OPENAI_BASE_URL")
+
+        client_kwargs: Dict[str, Any] = {**extra_client_kwargs}
+        if api_key:
+            client_kwargs["api_key"] = api_key
+        if organization:
+            client_kwargs["organization"] = organization
+        if base_url:
+            client_kwargs["base_url"] = base_url
+
+        self._client = OpenAI(**client_kwargs)
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Send a non-streaming request to 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 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; 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]:
+        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,
+        }
+        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]:
+        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