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

vanna/servers/cli/server_runner.py

@@ -0,0 +1,204 @@
+"""
+CLI for running Vanna Agents servers with example agents.
+"""
+
+import importlib
+import json
+from typing import Dict, Optional, Any, cast, TextIO, Union
+
+import click
+
+from ...core import Agent
+
+
+class ExampleAgentLoader:
+    """Loads example agents for the CLI."""
+
+    @staticmethod
+    def list_available_examples() -> Dict[str, str]:
+        """Return available examples with descriptions."""
+        return {
+            "mock_quickstart": "Basic agent with mock LLM service",
+            "anthropic_quickstart": "Agent configured for Anthropic's Claude API",
+            "openai_quickstart": "Agent configured for OpenAI's GPT models",
+            "mock_custom_tool": "Agent with custom tool demonstration (mock LLM)",
+            "mock_quota_example": "Agent with usage quota management (mock LLM)",
+            "mock_rich_components_demo": "Rich components demonstration with cards, tasks, and progress (mock LLM)",
+            "coding_agent_example": "Coding agent with file system tools (list, read, write files)",
+            "email_auth_example": "Email-based authentication demonstration (mock LLM)",
+            "claude_sqlite_example": "Claude agent with SQLite database querying capabilities",
+            "mock_sqlite_example": "Mock agent with SQLite database demonstration",
+        }
+
+    @staticmethod
+    def load_example_agent(example_name: str) -> Agent:
+        """Load an example agent by name.
+
+        Args:
+            example_name: Name of the example to load
+
+        Returns:
+            Configured agent instance
+
+        Raises:
+            ValueError: If example not found or failed to load
+        """
+        try:
+            # Import the example module
+            module = importlib.import_module(f"vanna.examples.{example_name}")
+
+            # Look for standard factory functions
+            factory_functions = [
+                "create_demo_agent",
+                "create_agent",
+                "create_basic_demo",
+            ]
+
+            for func_name in factory_functions:
+                if hasattr(module, func_name):
+                    factory = getattr(module, func_name)
+                    return cast(Agent, factory())
+
+            # Look for module-level agent instances
+            if hasattr(module, "main_agent"):
+                return cast(Agent, module.main_agent)
+
+            raise AttributeError(f"No agent factory found in {example_name}")
+
+        except ImportError as e:
+            raise ValueError(f"Example '{example_name}' not found: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to load example '{example_name}': {e}")
+
+
+@click.command()
+@click.option(
+    "--framework",
+    type=click.Choice(["flask", "fastapi"]),
+    default="fastapi",
+    help="Web framework to use",
+)
+@click.option("--port", default=8000, help="Port to run server on")
+@click.option("--host", default="0.0.0.0", help="Host to bind server to")
+@click.option(
+    "--example", help="Example agent to use (use --list-examples to see options)"
+)
+@click.option("--list-examples", is_flag=True, help="List available example agents")
+@click.option(
+    "--config", type=click.File("r"), help="JSON config file for server settings"
+)
+@click.option("--debug", is_flag=True, help="Enable debug mode")
+@click.option(
+    "--dev",
+    is_flag=True,
+    help="Enable development mode (load components from local assets)",
+)
+@click.option(
+    "--static-folder", default=None, help="Static folder path for development mode"
+)
+@click.option(
+    "--cdn-url",
+    default="https://img.vanna.ai/vanna-components.js",
+    help="CDN URL for web components",
+)
+def main(
+    framework: str,
+    port: int,
+    host: str,
+    example: Optional[str],
+    list_examples: bool,
+    config: Optional[click.File],
+    debug: bool,
+    dev: bool,
+    static_folder: Optional[str],
+    cdn_url: str,
+) -> None:
+    """Run Vanna Agents server with optional example agent."""
+
+    if list_examples:
+        click.echo("Available example agents:")
+        examples = ExampleAgentLoader.list_available_examples()
+        for name, description in examples.items():
+            click.echo(f"  {name:20} - {description}")
+        return
+
+    # Load configuration
+    server_config = {}
+    if config:
+        server_config = json.load(cast(TextIO, config))
+
+    # Set default static folder based on dev mode
+    if static_folder is None:
+        static_folder = "frontend/webcomponent/static" if dev else "static"
+
+    # Add CLI options to config
+    server_config.update(
+        {
+            "dev_mode": dev,
+            "static_folder": static_folder,
+            "cdn_url": cdn_url,
+            "api_base_url": "",  # Can be overridden in config file
+        }
+    )
+
+    # Create agent
+    if example:
+        try:
+            agent = ExampleAgentLoader.load_example_agent(example)
+            click.echo(f"✓ Loaded example agent: {example}")
+        except ValueError as e:
+            click.echo(f"Error: {e}", err=True)
+            return
+    else:
+        # Fallback to basic agent
+        try:
+            from ...agents import create_basic_agent
+            from ...integrations.mock import MockLlmService
+
+            llm_service = MockLlmService(
+                response_content="Hello! I'm a Vanna Agents demo server. How can I help you?"
+            )
+            agent = create_basic_agent(llm_service)
+            click.echo(
+                "✓ Using basic demo agent (use --example to specify different agent)"
+            )
+        except ImportError as e:
+            click.echo(f"Error: Could not create basic agent: {e}", err=True)
+            return
+
+    from ..flask.app import VannaFlaskServer
+    from ..fastapi.app import VannaFastAPIServer
+
+    # Create and run server
+    server: Union[VannaFlaskServer, VannaFastAPIServer]
+    if framework == "flask":
+        server = VannaFlaskServer(agent, config=server_config)
+        click.echo(f"🚀 Starting Flask server on http://{host}:{port}")
+        if dev:
+            click.echo(
+                f"📦 Development mode: loading web components from ./{static_folder}/"
+            )
+        else:
+            click.echo(f"🌍 Production mode: loading web components from CDN")
+        try:
+            server.run(host=host, port=port, debug=debug)
+        except KeyboardInterrupt:
+            click.echo("\n👋 Server stopped")
+    else:
+        server = VannaFastAPIServer(agent, config=server_config)
+        click.echo(f"🚀 Starting FastAPI server on http://{host}:{port}")
+        click.echo(f"📖 API docs available at http://{host}:{port}/docs")
+        if dev:
+            click.echo(
+                f"📦 Development mode: loading web components from ./{static_folder}/"
+            )
+        else:
+            click.echo(f"🌍 Production mode: loading web components from CDN")
+        try:
+            server.run(host=host, port=port)
+        except KeyboardInterrupt:
+            click.echo("\n👋 Server stopped")
+
+
+if __name__ == "__main__":
+    main()

vanna/servers/fastapi/__init__.py

@@ -0,0 +1,7 @@
+"""
+FastAPI server implementation for Vanna Agents.
+"""
+
+from .app import VannaFastAPIServer
+
+__all__ = ["VannaFastAPIServer"]

vanna/servers/fastapi/app.py

@@ -0,0 +1,163 @@
+"""
+FastAPI server factory for Vanna Agents.
+"""
+
+from typing import Any, Dict, Optional
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+from ...core import Agent
+from ..base import ChatHandler
+from .routes import register_chat_routes
+
+
+class VannaFastAPIServer:
+    """FastAPI server factory for Vanna Agents."""
+
+    def __init__(self, agent: Agent, config: Optional[Dict[str, Any]] = None):
+        """Initialize FastAPI server.
+
+        Args:
+            agent: The agent to serve (must have user_resolver configured)
+            config: Optional server configuration
+        """
+        self.agent = agent
+        self.config = config or {}
+        self.chat_handler = ChatHandler(agent)
+
+    def create_app(self) -> FastAPI:
+        """Create configured FastAPI app.
+
+        Returns:
+            Configured FastAPI application
+        """
+        # Create FastAPI app
+        app_config = self.config.get("fastapi", {})
+        app = FastAPI(
+            title="Vanna Agents API",
+            description="API server for Vanna Agents framework",
+            version="0.1.0",
+            **app_config,
+        )
+
+        # Configure CORS if enabled
+        cors_config = self.config.get("cors", {})
+        if cors_config.get("enabled", True):
+            cors_params = {k: v for k, v in cors_config.items() if k != "enabled"}
+
+            # Set sensible defaults
+            cors_params.setdefault("allow_origins", ["*"])
+            cors_params.setdefault("allow_credentials", True)
+            cors_params.setdefault("allow_methods", ["*"])
+            cors_params.setdefault("allow_headers", ["*"])
+
+            app.add_middleware(CORSMiddleware, **cors_params)
+
+        # Add static file serving in dev mode
+        dev_mode = self.config.get("dev_mode", False)
+        if dev_mode:
+            static_folder = self.config.get("static_folder", "static")
+            try:
+                import os
+
+                if os.path.exists(static_folder):
+                    app.mount(
+                        "/static", StaticFiles(directory=static_folder), name="static"
+                    )
+            except Exception:
+                pass  # Static files not available
+
+        # Register routes
+        register_chat_routes(app, self.chat_handler, self.config)
+
+        # Add health check
+        @app.get("/health")
+        async def health_check() -> Dict[str, str]:
+            return {"status": "healthy", "service": "vanna"}
+
+        return app
+
+    def run(self, **kwargs: Any) -> None:
+        """Run the FastAPI server.
+
+        This method automatically detects if running in an async environment
+        (Jupyter, Colab, IPython, etc.) and:
+        - Uses appropriate async handling for existing event loops
+        - Sets up port forwarding if in Google Colab
+        - Displays the correct URL for accessing the app
+
+        Args:
+            **kwargs: Arguments passed to uvicorn configuration
+        """
+        import sys
+        import asyncio
+        import uvicorn
+
+        # Check if we're in an environment with a running event loop FIRST
+        in_async_env = False
+        try:
+            asyncio.get_running_loop()
+            in_async_env = True
+        except RuntimeError:
+            in_async_env = False
+
+        # If in async environment, apply nest_asyncio BEFORE creating the app
+        if in_async_env:
+            try:
+                import nest_asyncio
+
+                nest_asyncio.apply()
+            except ImportError:
+                print("Warning: nest_asyncio not installed. Installing...")
+                import subprocess
+
+                subprocess.check_call(
+                    [sys.executable, "-m", "pip", "install", "nest_asyncio"]
+                )
+                import nest_asyncio
+
+                nest_asyncio.apply()
+
+        # Now create the app after nest_asyncio is applied
+        app = self.create_app()
+
+        # Set defaults
+        run_kwargs = {"host": "0.0.0.0", "port": 8000, "log_level": "info", **kwargs}
+
+        # Get the port and other config from run_kwargs
+        port = run_kwargs.get("port", 8000)
+        host = run_kwargs.get("host", "0.0.0.0")
+        log_level = run_kwargs.get("log_level", "info")
+
+        # Check if we're specifically in Google Colab for port forwarding
+        in_colab = "google.colab" in sys.modules
+
+        if in_colab:
+            try:
+                from google.colab import output
+
+                output.serve_kernel_port_as_window(port)
+                from google.colab.output import eval_js
+
+                print("Your app is running at:")
+                print(eval_js(f"google.colab.kernel.proxyPort({port})"))
+            except Exception as e:
+                print(f"Warning: Could not set up Colab port forwarding: {e}")
+                print(f"Your app is running at: http://localhost:{port}")
+        else:
+            print("Your app is running at:")
+            print(f"http://localhost:{port}")
+
+        if in_async_env:
+            # In Jupyter/Colab, create config with loop="asyncio" and use asyncio.run()
+            # This matches the working pattern from Colab
+            config = uvicorn.Config(
+                app, host=host, port=port, log_level=log_level, loop="asyncio"
+            )
+            server = uvicorn.Server(config)
+            asyncio.run(server.serve())
+        else:
+            # Normal execution outside of Jupyter/Colab
+            uvicorn.run(app, **run_kwargs)

vanna/servers/fastapi/routes.py

@@ -0,0 +1,183 @@
+"""
+FastAPI route implementations for Vanna Agents.
+"""
+
+import json
+import traceback
+from typing import Any, AsyncGenerator, Dict, Optional
+
+from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
+from fastapi.responses import StreamingResponse, HTMLResponse
+
+from ..base import ChatHandler, ChatRequest, ChatResponse
+from ..base.templates import get_index_html
+from ...core.user.request_context import RequestContext
+
+
+def register_chat_routes(
+    app: FastAPI, chat_handler: ChatHandler, config: Optional[Dict[str, Any]] = None
+) -> None:
+    """Register chat routes on FastAPI app.
+
+    Args:
+        app: FastAPI application
+        chat_handler: Chat handler instance
+        config: Server configuration
+    """
+    config = config or {}
+
+    @app.get("/", response_class=HTMLResponse)
+    async def index() -> str:
+        """Serve the main chat interface."""
+        dev_mode = config.get("dev_mode", False)
+        cdn_url = config.get("cdn_url", "https://img.vanna.ai/vanna-components.js")
+        api_base_url = config.get("api_base_url", "")
+
+        return get_index_html(
+            dev_mode=dev_mode, cdn_url=cdn_url, api_base_url=api_base_url
+        )
+
+    @app.post("/api/vanna/v2/chat_sse")
+    async def chat_sse(
+        chat_request: ChatRequest, http_request: Request
+    ) -> StreamingResponse:
+        """Server-Sent Events endpoint for streaming chat."""
+        # Extract request context for user resolution
+        chat_request.request_context = RequestContext(
+            cookies=dict(http_request.cookies),
+            headers=dict(http_request.headers),
+            remote_addr=http_request.client.host if http_request.client else None,
+            query_params=dict(http_request.query_params),
+            metadata=chat_request.metadata,
+        )
+
+        async def generate() -> AsyncGenerator[str, None]:
+            """Generate SSE stream."""
+            try:
+                async for chunk in chat_handler.handle_stream(chat_request):
+                    chunk_json = chunk.model_dump_json()
+                    yield f"data: {chunk_json}\n\n"
+                yield "data: [DONE]\n\n"
+            except Exception as e:
+                traceback.print_stack()
+                traceback.print_exc()
+                error_data = {
+                    "type": "error",
+                    "data": {"message": str(e)},
+                    "conversation_id": chat_request.conversation_id or "",
+                    "request_id": chat_request.request_id or "",
+                }
+                yield f"data: {json.dumps(error_data)}\n\n"
+
+        return StreamingResponse(
+            generate(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",  # Disable nginx buffering
+            },
+        )
+
+    @app.websocket("/api/vanna/v2/chat_websocket")
+    async def chat_websocket(websocket: WebSocket) -> None:
+        """WebSocket endpoint for real-time chat."""
+        await websocket.accept()
+
+        try:
+            while True:
+                # Receive message
+                try:
+                    data = await websocket.receive_json()
+
+                    # Extract request context for user resolution
+                    metadata = data.get("metadata", {})
+                    data["request_context"] = RequestContext(
+                        cookies=dict(websocket.cookies),
+                        headers=dict(websocket.headers),
+                        remote_addr=websocket.client.host if websocket.client else None,
+                        query_params=dict(websocket.query_params),
+                        metadata=metadata,
+                    )
+
+                    chat_request = ChatRequest(**data)
+                except Exception as e:
+                    traceback.print_stack()
+                    traceback.print_exc()
+                    await websocket.send_json(
+                        {
+                            "type": "error",
+                            "data": {"message": f"Invalid request: {str(e)}"},
+                        }
+                    )
+                    continue
+
+                # Stream response
+                try:
+                    async for chunk in chat_handler.handle_stream(chat_request):
+                        await websocket.send_json(chunk.model_dump())
+
+                    # Send completion signal
+                    await websocket.send_json(
+                        {
+                            "type": "completion",
+                            "data": {"status": "done"},
+                            "conversation_id": chunk.conversation_id
+                            if "chunk" in locals()
+                            else "",
+                            "request_id": chunk.request_id
+                            if "chunk" in locals()
+                            else "",
+                        }
+                    )
+
+                except Exception as e:
+                    traceback.print_stack()
+                    traceback.print_exc()
+                    await websocket.send_json(
+                        {
+                            "type": "error",
+                            "data": {"message": str(e)},
+                            "conversation_id": chat_request.conversation_id or "",
+                            "request_id": chat_request.request_id or "",
+                        }
+                    )
+
+        except WebSocketDisconnect:
+            pass
+        except Exception as e:
+            traceback.print_stack()
+            traceback.print_exc()
+            try:
+                await websocket.send_json(
+                    {
+                        "type": "error",
+                        "data": {"message": f"WebSocket error: {str(e)}"},
+                    }
+                )
+            except Exception:
+                pass
+            finally:
+                await websocket.close()
+
+    @app.post("/api/vanna/v2/chat_poll")
+    async def chat_poll(
+        chat_request: ChatRequest, http_request: Request
+    ) -> ChatResponse:
+        """Polling endpoint for chat."""
+        # Extract request context for user resolution
+        chat_request.request_context = RequestContext(
+            cookies=dict(http_request.cookies),
+            headers=dict(http_request.headers),
+            remote_addr=http_request.client.host if http_request.client else None,
+            query_params=dict(http_request.query_params),
+            metadata=chat_request.metadata,
+        )
+
+        try:
+            result = await chat_handler.handle_poll(chat_request)
+            return result
+        except Exception as e:
+            traceback.print_stack()
+            traceback.print_exc()
+            raise HTTPException(status_code=500, detail=f"Chat failed: {str(e)}")

vanna/servers/flask/__init__.py

@@ -0,0 +1,7 @@
+"""
+Flask server implementation for Vanna Agents.
+"""
+
+from .app import VannaFlaskServer
+
+__all__ = ["VannaFlaskServer"]

vanna/servers/flask/app.py

@@ -0,0 +1,132 @@
+"""
+Flask server factory for Vanna Agents.
+"""
+
+import asyncio
+from typing import Any, Dict, Optional
+
+from flask import Flask
+from flask_cors import CORS
+
+from ...core import Agent
+from ..base import ChatHandler
+from .routes import register_chat_routes
+
+
+class VannaFlaskServer:
+    """Flask server factory for Vanna Agents."""
+
+    def __init__(self, agent: Agent, config: Optional[Dict[str, Any]] = None):
+        """Initialize Flask server.
+
+        Args:
+            agent: The agent to serve (must have user_resolver configured)
+            config: Optional server configuration
+        """
+        self.agent = agent
+        self.config = config or {}
+        self.chat_handler = ChatHandler(agent)
+
+    def create_app(self) -> Flask:
+        """Create configured Flask app.
+
+        Returns:
+            Configured Flask application
+        """
+        # Check if dev mode is enabled
+        dev_mode = self.config.get("dev_mode", False)
+        static_folder = self.config.get("static_folder", "static") if dev_mode else None
+
+        app = Flask(__name__, static_folder=static_folder, static_url_path="/static")
+
+        # Apply configuration
+        app.config.update(self.config.get("flask", {}))
+
+        # Enable CORS if configured
+        cors_config = self.config.get("cors", {})
+        if cors_config.get("enabled", True):
+            CORS(app, **{k: v for k, v in cors_config.items() if k != "enabled"})
+
+        # Register routes
+        register_chat_routes(app, self.chat_handler, self.config)
+
+        # Add health check
+        @app.route("/health")
+        def health_check() -> Dict[str, str]:
+            return {"status": "healthy", "service": "vanna"}
+
+        return app
+
+    def run(self, **kwargs: Any) -> None:
+        """Run the Flask server.
+
+        This method automatically detects if running in an async environment
+        (Jupyter, Colab, IPython, etc.) and:
+        - Installs and applies nest_asyncio to handle existing event loops
+        - Sets up port forwarding if in Google Colab
+        - Displays the correct URL for accessing the app
+
+        Args:
+            **kwargs: Arguments passed to Flask.run()
+        """
+        import sys
+
+        app = self.create_app()
+
+        # Set defaults
+        run_kwargs = {"host": "0.0.0.0", "port": 5000, "debug": False, **kwargs}
+
+        # Get the port from run_kwargs
+        port = run_kwargs.get("port", 5000)
+
+        # Check if we're in an environment with a running event loop
+        # (Jupyter, Colab, IPython, VS Code notebooks, etc.)
+        in_async_env = False
+        try:
+            import asyncio
+
+            try:
+                asyncio.get_running_loop()
+                in_async_env = True
+            except RuntimeError:
+                in_async_env = False
+        except Exception:
+            pass
+
+        if in_async_env:
+            # Apply nest_asyncio to allow nested event loops
+            try:
+                import nest_asyncio
+
+                nest_asyncio.apply()
+            except ImportError:
+                print("Warning: nest_asyncio not installed. Installing...")
+                import subprocess
+
+                subprocess.check_call(
+                    [sys.executable, "-m", "pip", "install", "nest_asyncio"]
+                )
+                import nest_asyncio
+
+                nest_asyncio.apply()
+
+        # Check if we're specifically in Google Colab for port forwarding
+        in_colab = "google.colab" in sys.modules
+
+        if in_colab:
+            try:
+                from google.colab import output
+
+                output.serve_kernel_port_as_window(port)
+                from google.colab.output import eval_js
+
+                print("Your app is running at:")
+                print(eval_js(f"google.colab.kernel.proxyPort({port})"))
+            except Exception as e:
+                print(f"Warning: Could not set up Colab port forwarding: {e}")
+                print(f"Your app is running at: http://localhost:{port}")
+        else:
+            print("Your app is running at:")
+            print(f"http://localhost:{port}")
+
+        app.run(**run_kwargs)