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

vanna/__init__.py

@@ -1,400 +1,172 @@
-import dataclasses
-import json
-import os
-from dataclasses import dataclass
-from typing import Callable, List, Tuple, Union
-
-import pandas as pd
-import requests
-import plotly.graph_objs
-
-from .exceptions import (
-    OTPCodeError,
-    ValidationError,
-)
-from .types import (
-    ApiKey,
-    Status,
-    TrainingData,
-    UserEmail,
-    UserOTP,
-)
-from .utils import sanitize_model_name, validate_config_path
-
-api_key: Union[str, None] = None  # API key for Vanna.AI
-
-fig_as_img: bool = False  # Whether or not to return Plotly figures as images
-
-run_sql: Union[
-    Callable[[str], pd.DataFrame], None
-] = None  # Function to convert SQL to a Pandas DataFrame
 """
-**Example**
-```python
-vn.run_sql = lambda sql: pd.read_sql(sql, engine)
-```
-
-Set the SQL to DataFrame function for Vanna.AI. This is used in the [`vn.ask(...)`][vanna.ask] function.
-Instead of setting this directly you can also use [`vn.connect_to_snowflake(...)`][vanna.connect_to_snowflake] to set this.
+Vanna Agents - A modular framework for building LLM agents.
 
+This package provides a flexible framework for creating conversational AI agents
+with tool execution, conversation management, and user scoping.
 """
 
-__org: Union[str, None] = None  # Organization name for Vanna.AI
-
-_unauthenticated_endpoint = "https://ask.vanna.ai/unauthenticated_rpc"
-
-def error_deprecation():
-    raise Exception("""
-Please switch to the following method for initializing Vanna:
-
-from vanna.remote import VannaDefault
-
-api_key = # Your API key from https://vanna.ai/account/profile 
-vanna_model_name = # Your model name from https://vanna.ai/account/profile
-                    
-vn = VannaDefault(model=vanna_model_name, api_key=api_key)
-""")
-
-def __unauthenticated_rpc_call(method, params):
-    headers = {
-        "Content-Type": "application/json",
-    }
-    data = {"method": method, "params": [__dataclass_to_dict(obj) for obj in params]}
-
-    response = requests.post(
-        _unauthenticated_endpoint, headers=headers, data=json.dumps(data)
-    )
-    return response.json()
-
-
-
-def __dataclass_to_dict(obj):
-    return dataclasses.asdict(obj)
-
-
-def get_api_key(email: str, otp_code: Union[str, None] = None) -> str:
-    """
-    **Example:**
-    ```python
-    vn.get_api_key(email="my-email@example.com")
-    ```
-
-    Login to the Vanna.AI API.
-
-    Args:
-        email (str): The email address to login with.
-        otp_code (Union[str, None]): The OTP code to login with. If None, an OTP code will be sent to the email address.
-
-    Returns:
-        str: The API key.
-    """
-    vanna_api_key = os.environ.get("VANNA_API_KEY", None)
-
-    if vanna_api_key is not None:
-        return vanna_api_key
-
-    if email == "my-email@example.com":
-        raise ValidationError(
-            "Please replace 'my-email@example.com' with your email address."
-        )
-
-    if otp_code is None:
-        params = [UserEmail(email=email)]
-
-        d = __unauthenticated_rpc_call(method="send_otp", params=params)
-
-        if "result" not in d:
-            raise OTPCodeError("Error sending OTP code.")
-
-        status = Status(**d["result"])
-
-        if not status.success:
-            raise OTPCodeError(f"Error sending OTP code: {status.message}")
-
-        otp_code = input("Check your email for the code and enter it here: ")
-
-    params = [UserOTP(email=email, otp=otp_code)]
-
-    d = __unauthenticated_rpc_call(method="verify_otp", params=params)
-
-    if "result" not in d:
-        raise OTPCodeError("Error verifying OTP code.")
-
-    key = ApiKey(**d["result"])
-
-    if key is None:
-        raise OTPCodeError("Error verifying OTP code.")
-
-    api_key = key.key
-
-    return api_key
-
-
-def set_api_key(key: str) -> None:
-    error_deprecation()
-
-
-def get_models() -> List[str]:
-    error_deprecation()
-
-
-def create_model(model: str, db_type: str) -> bool:
-    error_deprecation()
-
-
-def add_user_to_model(model: str, email: str, is_admin: bool) -> bool:
-    error_deprecation()
-
-
-def update_model_visibility(public: bool) -> bool:
-    error_deprecation()
-
-
-def set_model(model: str):
-    error_deprecation()
-
-
-def add_sql(
-    question: str, sql: str, tag: Union[str, None] = "Manually Trained"
-) -> bool:
-    error_deprecation()
-
-
-def add_ddl(ddl: str) -> bool:
-    error_deprecation()
-
-
-def add_documentation(documentation: str) -> bool:
-    error_deprecation()
-
-
-@dataclass
-class TrainingPlanItem:
-    item_type: str
-    item_group: str
-    item_name: str
-    item_value: str
-
-    def __str__(self):
-        if self.item_type == self.ITEM_TYPE_SQL:
-            return f"Train on SQL: {self.item_group} {self.item_name}"
-        elif self.item_type == self.ITEM_TYPE_DDL:
-            return f"Train on DDL: {self.item_group} {self.item_name}"
-        elif self.item_type == self.ITEM_TYPE_IS:
-            return f"Train on Information Schema: {self.item_group} {self.item_name}"
-
-    ITEM_TYPE_SQL = "sql"
-    ITEM_TYPE_DDL = "ddl"
-    ITEM_TYPE_IS = "is"
-
-
-class TrainingPlan:
-    """
-    A class representing a training plan. You can see what's in it, and remove items from it that you don't want trained.
-
-    **Example:**
-    ```python
-    plan = vn.get_training_plan()
-
-    plan.get_summary()
-    ```
-
-    """
-
-    _plan: List[TrainingPlanItem]
-
-    def __init__(self, plan: List[TrainingPlanItem]):
-        self._plan = plan
-
-    def __str__(self):
-        return "\n".join(self.get_summary())
-
-    def __repr__(self):
-        return self.__str__()
-
-    def get_summary(self) -> List[str]:
-        """
-        **Example:**
-        ```python
-        plan = vn.get_training_plan()
-
-        plan.get_summary()
-        ```
-
-        Get a summary of the training plan.
-
-        Returns:
-            List[str]: A list of strings describing the training plan.
-        """
-
-        return [f"{item}" for item in self._plan]
-
-    def remove_item(self, item: str):
-        """
-        **Example:**
-        ```python
-        plan = vn.get_training_plan()
-
-        plan.remove_item("Train on SQL: What is the average salary of employees?")
-        ```
-
-        Remove an item from the training plan.
-
-        Args:
-            item (str): The item to remove.
-        """
-        for plan_item in self._plan:
-            if str(plan_item) == item:
-                self._plan.remove(plan_item)
-                break
-
-
-def get_training_plan_postgres(
-    filter_databases: Union[List[str], None] = None,
-    filter_schemas: Union[List[str], None] = None,
-    include_information_schema: bool = False,
-    use_historical_queries: bool = True,
-) -> TrainingPlan:
-    error_deprecation()
-
-
-def get_training_plan_generic(df) -> TrainingPlan:
-    error_deprecation()
-
-
-def get_training_plan_experimental(
-    filter_databases: Union[List[str], None] = None,
-    filter_schemas: Union[List[str], None] = None,
-    include_information_schema: bool = False,
-    use_historical_queries: bool = True,
-) -> TrainingPlan:
-    error_deprecation()
-
-
-def train(
-    question: str = None,
-    sql: str = None,
-    ddl: str = None,
-    documentation: str = None,
-    json_file: str = None,
-    sql_file: str = None,
-    plan: TrainingPlan = None,
-) -> bool:
-    error_deprecation()
-
-
-def flag_sql_for_review(
-    question: str, sql: Union[str, None] = None, error_msg: Union[str, None] = None
-) -> bool:
-    error_deprecation()
-
-
-def remove_sql(question: str) -> bool:
-    error_deprecation()
-
-
-def remove_training_data(id: str) -> bool:
-    error_deprecation()
-
-
-def generate_sql(question: str) -> str:
-    error_deprecation()
-
-
-def get_related_training_data(question: str) -> TrainingData:
-    error_deprecation()
-
-
-def generate_meta(question: str) -> str:
-    error_deprecation()
-
-
-def generate_followup_questions(question: str, df: pd.DataFrame) -> List[str]:
-    error_deprecation()
-
-
-def generate_questions() -> List[str]:
-    error_deprecation()
-
-
-def ask(
-    question: Union[str, None] = None,
-    print_results: bool = True,
-    auto_train: bool = True,
-    generate_followups: bool = True,
-) -> Union[
-    Tuple[
-        Union[str, None],
-        Union[pd.DataFrame, None],
-        Union[plotly.graph_objs.Figure, None],
-        Union[List[str], None],
-    ],
-    None,
-]:
-    error_deprecation()
-
-def generate_plotly_code(
-    question: Union[str, None],
-    sql: Union[str, None],
-    df: pd.DataFrame,
-    chart_instructions: Union[str, None] = None,
-) -> str:
-    error_deprecation()
-
-
-def get_plotly_figure(
-    plotly_code: str, df: pd.DataFrame, dark_mode: bool = True
-) -> plotly.graph_objs.Figure:
-    error_deprecation()
-
-
-def get_results(cs, default_database: str, sql: str) -> pd.DataFrame:
-    error_deprecation()
-
-
-def generate_explanation(sql: str) -> str:
-    error_deprecation()
-
-
-def generate_question(sql: str) -> str:
-    error_deprecation()
-
-
-def get_all_questions() -> pd.DataFrame:
-    error_deprecation()
-
-
-def get_training_data() -> pd.DataFrame:
-    error_deprecation()
-
-
-def connect_to_sqlite(url: str):
-    error_deprecation()
-
-
-def connect_to_snowflake(
-    account: str,
-    username: str,
-    password: str,
-    database: str,
-    schema: Union[str, None] = None,
-    role: Union[str, None] = None,
-):
-    error_deprecation()
-
-
-def connect_to_postgres(
-    host: str = None,
-    dbname: str = None,
-    user: str = None,
-    password: str = None,
-    port: int = None,
-):
-    error_deprecation()
-
-
-def connect_to_bigquery(cred_file_path: str = None, project_id: str = None):
-    error_deprecation()
+# Version information
+__version__ = "0.1.0"
+
+# Import core framework components
+from .core import (
+    # Interfaces
+    Agent,
+    ConversationStore,
+    LlmService,
+    SystemPromptBuilder,
+    Tool,
+    UserService,
+    T,
+    # Models
+    Conversation,
+    LlmMessage,
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+    Message,
+    ToolCall,
+    ToolContext,
+    ToolResult,
+    ToolSchema,
+    User,
+    # UI Components
+    UiComponent,
+    SimpleComponent,
+    SimpleComponentType,
+    SimpleTextComponent,
+    SimpleImageComponent,
+    SimpleLinkComponent,
+    # Rich Components
+    ArtifactComponent,
+    BadgeComponent,
+    CardComponent,
+    DataFrameComponent,
+    IconTextComponent,
+    LogViewerComponent,
+    NotificationComponent,
+    ProgressBarComponent,
+    ProgressDisplayComponent,
+    RichTextComponent,
+    StatusCardComponent,
+    TaskListComponent,
+    # Core implementations
+    Agent,
+    AgentConfig,
+    DefaultSystemPromptBuilder,
+    DefaultWorkflowHandler,
+    ToolRegistry,
+    # Evaluation
+    Evaluator,
+    TestCase,
+    ExpectedOutcome,
+    AgentResult,
+    EvaluationResult,
+    TestCaseResult,
+    AgentVariant,
+    EvaluationRunner,
+    TrajectoryEvaluator,
+    OutputEvaluator,
+    LLMAsJudgeEvaluator,
+    EfficiencyEvaluator,
+    EvaluationReport,
+    ComparisonReport,
+    EvaluationDataset,
+    # Exceptions
+    AgentError,
+    ConversationNotFoundError,
+    LlmServiceError,
+    PermissionError,
+    ToolExecutionError,
+    ToolNotFoundError,
+    ValidationError,
+)
 
-def connect_to_duckdb(url: str="memory", init_sql: str = None):
-    error_deprecation()
+# Import basic implementations
+from .integrations import MemoryConversationStore, MockLlmService
+
+# Main exports
+__all__ = [
+    # Version
+    "__version__",
+    # Core interfaces
+    "Agent",
+    "Tool",
+    "LlmService",
+    "ConversationStore",
+    "UserService",
+    "SystemPromptBuilder",
+    "T",
+    # Models
+    "User",
+    "Message",
+    "Conversation",
+    "ToolCall",
+    "ToolResult",
+    "ToolContext",
+    "ToolSchema",
+    "LlmMessage",
+    "LlmRequest",
+    "LlmResponse",
+    "LlmStreamChunk",
+    # UI Components
+    "UiComponent",
+    "SimpleComponent",
+    "SimpleComponentType",
+    "SimpleTextComponent",
+    "SimpleImageComponent",
+    "SimpleLinkComponent",
+    # Rich Components
+    "ArtifactComponent",
+    "BadgeComponent",
+    "CardComponent",
+    "DataFrameComponent",
+    "IconTextComponent",
+    "LogViewerComponent",
+    "NotificationComponent",
+    "ProgressBarComponent",
+    "ProgressDisplayComponent",
+    "RichTextComponent",
+    "StatusCardComponent",
+    "TaskListComponent",
+    # Core implementations
+    "Agent",
+    "AgentConfig",
+    "ToolRegistry",
+    "DefaultSystemPromptBuilder",
+    "DefaultWorkflowHandler",
+    # Evaluation
+    "Evaluator",
+    "TestCase",
+    "ExpectedOutcome",
+    "AgentResult",
+    "EvaluationResult",
+    "TestCaseResult",
+    "AgentVariant",
+    "EvaluationRunner",
+    "TrajectoryEvaluator",
+    "OutputEvaluator",
+    "LLMAsJudgeEvaluator",
+    "EfficiencyEvaluator",
+    "EvaluationReport",
+    "ComparisonReport",
+    "EvaluationDataset",
+    # Basic implementations
+    "MemoryConversationStore",
+    "MockLlmService",
+    # Server components
+    "VannaFlaskServer",
+    "VannaFastAPIServer",
+    "ChatHandler",
+    "ChatRequest",
+    "ChatStreamChunk",
+    "ExampleAgentLoader",
+    # Exceptions
+    "AgentError",
+    "ToolExecutionError",
+    "ToolNotFoundError",
+    "PermissionError",
+    "ConversationNotFoundError",
+    "LlmServiceError",
+    "ValidationError",
+]

vanna/agents/__init__.py

@@ -0,0 +1,7 @@
+"""
+Agent implementations.
+
+This package contains agent implementations and utilities.
+"""
+
+__all__: list[str] = []

vanna/capabilities/__init__.py

@@ -0,0 +1,17 @@
+"""
+Capabilities module.
+
+This package contains abstractions for tool capabilities - reusable utilities
+that tools can compose via dependency injection.
+"""
+
+from .file_system import CommandResult, FileSearchMatch, FileSystem
+from .sql_runner import RunSqlToolArgs, SqlRunner
+
+__all__ = [
+    "FileSystem",
+    "FileSearchMatch",
+    "CommandResult",
+    "SqlRunner",
+    "RunSqlToolArgs",
+]

vanna/capabilities/agent_memory/__init__.py

@@ -0,0 +1,21 @@
+"""
+Agent memory capability package.
+"""
+
+from .base import AgentMemory
+from .models import (
+    MemoryStats,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+
+__all__ = [
+    "AgentMemory",
+    "TextMemory",
+    "TextMemorySearchResult",
+    "ToolMemory",
+    "ToolMemorySearchResult",
+    "MemoryStats",
+]