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

vanna/integrations/openai/responses.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, TYPE_CHECKING
+
+from vanna.core.llm import LlmService, LlmRequest, LlmResponse, LlmStreamChunk
+from vanna.core.tool import ToolCall, ToolSchema
+
+if TYPE_CHECKING:
+    from openai.types.responses import Response
+
+
+class OpenAIResponsesService(LlmService):
+    def __init__(
+        self, api_key: Optional[str] = None, model: Optional[str] = None
+    ) -> None:
+        try:
+            from openai import AsyncOpenAI
+            from openai.types.responses import Response
+        except Exception as e:  # pragma: no cover
+            raise ImportError(
+                "openai package is required. Install with: pip install 'vanna[openai]'"
+            ) from e
+
+        self.client = AsyncOpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY"))
+        self.model = model or os.getenv("OPENAI_MODEL", "gpt-5")
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        payload = self._payload(request)
+        resp: Response = await self.client.responses.create(**payload)
+        self._debug_print("response", resp)
+        text, tools, status, usage = self._extract(resp)
+        return LlmResponse(
+            content=text,
+            tool_calls=tools or None,
+            finish_reason=status,
+            usage=usage or None,
+            metadata={"request_id": getattr(resp, "id", None)},
+        )
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        payload = self._payload(request)
+        async with self.client.responses.stream(**payload) as stream:
+            async for event in stream:
+                self._debug_print("stream_event", event)
+                event_type = getattr(event, "type", None)
+                if event_type == "response.output_text.delta":
+                    delta = getattr(event, "delta", None)
+                    if delta:
+                        yield LlmStreamChunk(content=delta)
+            final: Response = await stream.get_final_response()
+            self._debug_print("final_response", final)
+
+        _text, tools, status, _usage = self._extract(final)
+        yield LlmStreamChunk(tool_calls=tools or None, finish_reason=status)
+
+    async def validate_tools(self, tools: List[Any]) -> List[str]:
+        return []  # minimal: accept whatever's passed through
+
+    # ---- helpers ----
+
+    def _payload(self, request: LlmRequest) -> Dict[str, Any]:
+        msgs = [{"role": m.role, "content": m.content} for m in request.messages]
+        p: Dict[str, Any] = {"model": self.model, "input": msgs}
+        if request.system_prompt:
+            p["instructions"] = request.system_prompt
+        if request.max_tokens:
+            p["max_output_tokens"] = request.max_tokens
+        if request.tools:
+            p["tools"] = [self._serialize_tool(t) for t in request.tools]
+        return p
+
+    def _debug_print(self, label: str, obj: Any) -> None:
+        try:
+            payload = obj.model_dump()
+        except AttributeError:
+            try:
+                payload = obj.dict()
+            except AttributeError:
+                payload = obj
+        print(f"[OpenAIResponsesService] {label}: {payload}")
+
+    def _extract(
+        self, resp: Response
+    ) -> Tuple[
+        Optional[str], Optional[List[ToolCall]], Optional[str], Optional[Dict[str, int]]
+    ]:
+        text = getattr(resp, "output_text", None)
+
+        tool_calls: List[ToolCall] = []
+        for oc in getattr(resp, "output", []) or []:
+            for item in getattr(oc, "content", []) or []:
+                if getattr(item, "type", None) == "tool_call":
+                    tc = getattr(item, "tool_call", None)
+                    if tc and getattr(tc, "type", None) == "function":
+                        fn = getattr(tc, "function", None)
+                        if fn:
+                            name = getattr(fn, "name", None)
+                            args = getattr(fn, "arguments", None)
+                            if not isinstance(args, (dict, list)):
+                                try:
+                                    args = json.loads(args) if args else {}
+                                except Exception:
+                                    args = {"_raw": args}
+                            tool_calls.append(ToolCall(name=name, arguments=args))
+
+        usage = None
+        if getattr(resp, "usage", None):
+            usage = {
+                "input_tokens": getattr(resp.usage, "input_tokens", 0) or 0,
+                "output_tokens": getattr(resp.usage, "output_tokens", 0) or 0,
+                "total_tokens": getattr(resp.usage, "total_tokens", None)
+                or (
+                    (getattr(resp.usage, "input_tokens", 0) or 0)
+                    + (getattr(resp.usage, "output_tokens", 0) or 0)
+                ),
+            }
+
+        status = getattr(resp, "status", None)  # e.g. "completed"
+        return text, (tool_calls or None), status, usage
+
+    def _serialize_tool(self, tool: Any) -> Dict[str, Any]:
+        """Convert a tool schema into the dict format expected by OpenAI Responses."""
+
+        if isinstance(tool, ToolSchema):
+            return {
+                "type": "function",
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": tool.parameters,
+                "strict": False,
+            }
+
+        # Support generic pydantic/BaseModel style objects without importing pydantic here.
+        if hasattr(tool, "model_dump"):
+            data = tool.model_dump()
+            if all(key in data for key in ("name", "description", "parameters")):
+                return {
+                    "type": "function",
+                    "name": data["name"],
+                    "description": data["description"],
+                    "parameters": data["parameters"],
+                    "strict": data.get("strict", False),
+                }
+            return data
+
+        if isinstance(tool, dict):
+            if "type" in tool:
+                return tool
+            if all(k in tool for k in ("name", "description", "parameters")):
+                return {
+                    "type": "function",
+                    "name": tool["name"],
+                    "description": tool["description"],
+                    "parameters": tool["parameters"],
+                    "strict": tool.get("strict", False),
+                }
+            return tool
+
+        raise TypeError(f"Unsupported tool schema type: {type(tool)!r}")

vanna/integrations/opensearch/__init__.py

@@ -0,0 +1,7 @@
+"""
+OpenSearch integration for Vanna Agents.
+"""
+
+from .agent_memory import OpenSearchAgentMemory
+
+__all__ = ["OpenSearchAgentMemory"]

vanna/integrations/opensearch/agent_memory.py

@@ -0,0 +1,411 @@
+"""
+OpenSearch vector database implementation of AgentMemory.
+
+This implementation uses OpenSearch for distributed search and storage of tool usage patterns.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+try:
+    from opensearchpy import OpenSearch, helpers
+
+    OPENSEARCH_AVAILABLE = True
+except ImportError:
+    OPENSEARCH_AVAILABLE = False
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class OpenSearchAgentMemory(AgentMemory):
+    """OpenSearch-based implementation of AgentMemory."""
+
+    def __init__(
+        self,
+        index_name: str = "tool_memories",
+        hosts: Optional[List[str]] = None,
+        http_auth: Optional[tuple] = None,
+        use_ssl: bool = False,
+        verify_certs: bool = False,
+        dimension: int = 384,
+    ):
+        if not OPENSEARCH_AVAILABLE:
+            raise ImportError(
+                "OpenSearch is required for OpenSearchAgentMemory. Install with: pip install opensearch-py"
+            )
+
+        self.index_name = index_name
+        self.hosts = hosts or ["localhost:9200"]
+        self.http_auth = http_auth
+        self.use_ssl = use_ssl
+        self.verify_certs = verify_certs
+        self.dimension = dimension
+        self._client = None
+        self._executor = ThreadPoolExecutor(max_workers=2)
+
+    def _get_client(self):
+        """Get or create OpenSearch client."""
+        if self._client is None:
+            self._client = OpenSearch(
+                hosts=self.hosts,
+                http_auth=self.http_auth,
+                use_ssl=self.use_ssl,
+                verify_certs=self.verify_certs,
+                ssl_show_warn=False,
+            )
+
+            # Create index if it doesn't exist
+            if not self._client.indices.exists(index=self.index_name):
+                index_body = {
+                    "settings": {
+                        "index": {"knn": True, "knn.algo_param.ef_search": 100}
+                    },
+                    "mappings": {
+                        "properties": {
+                            "memory_id": {"type": "keyword"},
+                            "question": {"type": "text"},
+                            "tool_name": {"type": "keyword"},
+                            "args": {"type": "object", "enabled": False},
+                            "timestamp": {"type": "date"},
+                            "success": {"type": "boolean"},
+                            "metadata": {"type": "object", "enabled": False},
+                            "embedding": {
+                                "type": "knn_vector",
+                                "dimension": self.dimension,
+                                "method": {
+                                    "name": "hnsw",
+                                    "space_type": "cosinesimil",
+                                    "engine": "nmslib",
+                                },
+                            },
+                        }
+                    },
+                }
+                self._client.indices.create(index=self.index_name, body=index_body)
+
+        return self._client
+
+    def _create_embedding(self, text: str) -> List[float]:
+        """Create a simple embedding from text (placeholder)."""
+        import hashlib
+
+        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
+        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]
+
+    async def save_tool_usage(
+        self,
+        question: str,
+        tool_name: str,
+        args: Dict[str, Any],
+        context: ToolContext,
+        success: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Save a tool usage pattern."""
+
+        def _save():
+            client = self._get_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(question)
+
+            document = {
+                "memory_id": memory_id,
+                "question": question,
+                "tool_name": tool_name,
+                "args": args,
+                "timestamp": timestamp,
+                "success": success,
+                "metadata": metadata or {},
+                "embedding": embedding,
+            }
+
+            client.index(
+                index=self.index_name, body=document, id=memory_id, refresh=True
+            )
+
+        await asyncio.get_event_loop().run_in_executor(self._executor, _save)
+
+    async def search_similar_usage(
+        self,
+        question: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+        tool_name_filter: Optional[str] = None,
+    ) -> List[ToolMemorySearchResult]:
+        """Search for similar tool usage patterns."""
+
+        def _search():
+            client = self._get_client()
+
+            embedding = self._create_embedding(question)
+
+            # Build query
+            must_conditions = [{"term": {"success": True}}]
+            if tool_name_filter:
+                must_conditions.append({"term": {"tool_name": tool_name_filter}})
+
+            query = {
+                "size": limit,
+                "query": {
+                    "bool": {
+                        "must": must_conditions,
+                        "filter": {
+                            "knn": {"embedding": {"vector": embedding, "k": limit}}
+                        },
+                    }
+                },
+            }
+
+            response = client.search(index=self.index_name, body=query)
+
+            search_results = []
+            for i, hit in enumerate(response["hits"]["hits"]):
+                source = hit["_source"]
+                score = hit["_score"]
+
+                # Normalize score to 0-1 range (OpenSearch scores can vary)
+                similarity_score = min(score / 10.0, 1.0)
+
+                if similarity_score >= similarity_threshold:
+                    memory = ToolMemory(
+                        memory_id=source["memory_id"],
+                        question=source["question"],
+                        tool_name=source["tool_name"],
+                        args=source["args"],
+                        timestamp=source.get("timestamp"),
+                        success=source.get("success", True),
+                        metadata=source.get("metadata", {}),
+                    )
+
+                    search_results.append(
+                        ToolMemorySearchResult(
+                            memory=memory, similarity_score=similarity_score, rank=i + 1
+                        )
+                    )
+
+            return search_results
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)
+
+    async def get_recent_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[ToolMemory]:
+        """Get recently added memories."""
+
+        def _get_recent():
+            client = self._get_client()
+
+            query = {
+                "size": limit,
+                "query": {"match_all": {}},
+                "sort": [{"timestamp": {"order": "desc"}}],
+            }
+
+            response = client.search(index=self.index_name, body=query)
+
+            memories = []
+            for hit in response["hits"]["hits"]:
+                source = hit["_source"]
+
+                memory = ToolMemory(
+                    memory_id=source["memory_id"],
+                    question=source["question"],
+                    tool_name=source["tool_name"],
+                    args=source["args"],
+                    timestamp=source.get("timestamp"),
+                    success=source.get("success", True),
+                    metadata=source.get("metadata", {}),
+                )
+                memories.append(memory)
+
+            return memories
+
+        return await asyncio.get_event_loop().run_in_executor(
+            self._executor, _get_recent
+        )
+
+    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a memory by its ID."""
+
+        def _delete():
+            client = self._get_client()
+
+            try:
+                client.delete(index=self.index_name, id=memory_id, refresh=True)
+                return True
+            except Exception:
+                return False
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)
+
+    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
+        """Save a text memory."""
+
+        def _save():
+            client = self._get_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(content)
+
+            document = {
+                "memory_id": memory_id,
+                "content": content,
+                "timestamp": timestamp,
+                "is_text_memory": True,
+                "embedding": embedding,
+            }
+
+            client.index(
+                index=self.index_name, body=document, id=memory_id, refresh=True
+            )
+
+            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)
+
+    async def search_text_memories(
+        self,
+        query: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+    ) -> List[TextMemorySearchResult]:
+        """Search for similar text memories."""
+
+        def _search():
+            client = self._get_client()
+
+            embedding = self._create_embedding(query)
+
+            query_body = {
+                "size": limit,
+                "query": {
+                    "bool": {
+                        "must": [{"term": {"is_text_memory": True}}],
+                        "filter": {
+                            "knn": {"embedding": {"vector": embedding, "k": limit}}
+                        },
+                    }
+                },
+            }
+
+            response = client.search(index=self.index_name, body=query_body)
+
+            search_results = []
+            for i, hit in enumerate(response["hits"]["hits"]):
+                source = hit["_source"]
+                score = hit["_score"]
+
+                similarity_score = min(score / 10.0, 1.0)
+
+                if similarity_score >= similarity_threshold:
+                    memory = TextMemory(
+                        memory_id=source["memory_id"],
+                        content=source.get("content", ""),
+                        timestamp=source.get("timestamp"),
+                    )
+
+                    search_results.append(
+                        TextMemorySearchResult(
+                            memory=memory, similarity_score=similarity_score, rank=i + 1
+                        )
+                    )
+
+            return search_results
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)
+
+    async def get_recent_text_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[TextMemory]:
+        """Get recently added text memories."""
+
+        def _get_recent():
+            client = self._get_client()
+
+            query = {
+                "size": limit,
+                "query": {"term": {"is_text_memory": True}},
+                "sort": [{"timestamp": {"order": "desc"}}],
+            }
+
+            response = client.search(index=self.index_name, body=query)
+
+            memories = []
+            for hit in response["hits"]["hits"]:
+                source = hit["_source"]
+
+                memory = TextMemory(
+                    memory_id=source["memory_id"],
+                    content=source.get("content", ""),
+                    timestamp=source.get("timestamp"),
+                )
+                memories.append(memory)
+
+            return memories
+
+        return await asyncio.get_event_loop().run_in_executor(
+            self._executor, _get_recent
+        )
+
+    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a text memory by its ID."""
+
+        def _delete():
+            client = self._get_client()
+
+            try:
+                client.delete(index=self.index_name, id=memory_id, refresh=True)
+                return True
+            except Exception:
+                return False
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)
+
+    async def clear_memories(
+        self,
+        context: ToolContext,
+        tool_name: Optional[str] = None,
+        before_date: Optional[str] = None,
+    ) -> int:
+        """Clear stored memories."""
+
+        def _clear():
+            client = self._get_client()
+
+            # Build query
+            must_conditions = []
+            if tool_name:
+                must_conditions.append({"term": {"tool_name": tool_name}})
+            if before_date:
+                must_conditions.append({"range": {"timestamp": {"lt": before_date}}})
+
+            if must_conditions:
+                query = {"query": {"bool": {"must": must_conditions}}}
+            else:
+                query = {"query": {"match_all": {}}}
+
+            response = client.delete_by_query(
+                index=self.index_name, body=query, refresh=True
+            )
+
+            return response.get("deleted", 0)
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)

vanna/integrations/oracle/__init__.py

@@ -0,0 +1,5 @@
+"""Oracle integration for Vanna."""
+
+from .sql_runner import OracleRunner
+
+__all__ = ["OracleRunner"]

vanna/integrations/oracle/sql_runner.py

@@ -0,0 +1,75 @@
+"""Oracle implementation of SqlRunner interface."""
+
+from typing import Optional
+import pandas as pd
+
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.core.tool import ToolContext
+
+
+class OracleRunner(SqlRunner):
+    """Oracle implementation of the SqlRunner interface."""
+
+    def __init__(self, user: str, password: str, dsn: str, **kwargs):
+        """Initialize with Oracle connection parameters.
+
+        Args:
+            user: Oracle database user name
+            password: Oracle database user password
+            dsn: Oracle database host - format: host:port/sid
+            **kwargs: Additional oracledb connection parameters
+        """
+        try:
+            import oracledb
+
+            self.oracledb = oracledb
+        except ImportError as e:
+            raise ImportError(
+                "oracledb package is required. Install with: pip install 'vanna[oracle]'"
+            ) from e
+
+        self.user = user
+        self.password = password
+        self.dsn = dsn
+        self.kwargs = kwargs
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against Oracle database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            oracledb.Error: If query execution fails
+        """
+        # Connect to the database
+        conn = self.oracledb.connect(
+            user=self.user, password=self.password, dsn=self.dsn, **self.kwargs
+        )
+
+        cursor = conn.cursor()
+
+        try:
+            # Strip and remove trailing semicolons (Oracle doesn't like them)
+            sql = args.sql.rstrip()
+            if sql.endswith(";"):
+                sql = sql[:-1]
+
+            # Execute the query
+            cursor.execute(sql)
+            results = cursor.fetchall()
+
+            # Create a pandas dataframe from the results
+            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
+            return df
+
+        except self.oracledb.Error:
+            conn.rollback()
+            raise
+        finally:
+            cursor.close()
+            conn.close()

vanna/integrations/pinecone/__init__.py

@@ -0,0 +1,7 @@
+"""
+Pinecone integration for Vanna Agents.
+"""
+
+from .agent_memory import PineconeAgentMemory
+
+__all__ = ["PineconeAgentMemory"]