asksql 2.1.0__py3-none-any.whl

vanna/legacy/weaviate/weaviate_vector.py

@@ -0,0 +1,194 @@
+import weaviate
+import weaviate.classes as wvc
+from fastembed import TextEmbedding
+
+from vanna.base import VannaBase
+
+
+class WeaviateDatabase(VannaBase):
+    def __init__(self, config=None):
+        """
+        Initialize the VannaEnhanced class with the provided configuration.
+
+        :param config: Dictionary containing configuration parameters.
+
+        params:
+        weaviate_url (str): Weaviate cluster URL while using weaviate cloud,
+        weaviate_api_key (str): Weaviate API key while using weaviate cloud,
+        weaviate_port (num): Weaviate port while using local weaviate,
+        weaviate_grpc (num): Weaviate gRPC port while using local weaviate,
+        fastembed_model (str): Fastembed model name for text embeddings. BAAI/bge-small-en-v1.5 by default.
+
+        """
+        super().__init__(config=config)
+
+        if config is None:
+            raise ValueError("config is required")
+
+        self.n_results = config.get("n_results", 3)
+        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
+        self.weaviate_api_key = config.get("weaviate_api_key")
+        self.weaviate_url = config.get("weaviate_url")
+        self.weaviate_port = config.get("weaviate_port")
+        self.weaviate_grpc_port = config.get("weaviate_grpc", 50051)
+
+        if not self.weaviate_api_key and not self.weaviate_port:
+            raise ValueError("Add proper credentials to connect to weaviate")
+
+        self.weaviate_client = self._initialize_weaviate_client()
+        self.embeddings = TextEmbedding(model_name=self.fastembed_model)
+
+        self.training_data_cluster = {
+            "sql": "SQLTrainingDataEntry",
+            "ddl": "DDLEntry",
+            "doc": "DocumentationEntry",
+        }
+
+        self._create_collections_if_not_exist()
+
+    def _create_collections_if_not_exist(self):
+        properties_dict = {
+            self.training_data_cluster["ddl"]: [
+                wvc.config.Property(
+                    name="description", data_type=wvc.config.DataType.TEXT
+                ),
+            ],
+            self.training_data_cluster["doc"]: [
+                wvc.config.Property(
+                    name="description", data_type=wvc.config.DataType.TEXT
+                ),
+            ],
+            self.training_data_cluster["sql"]: [
+                wvc.config.Property(name="sql", data_type=wvc.config.DataType.TEXT),
+                wvc.config.Property(
+                    name="natural_language_question", data_type=wvc.config.DataType.TEXT
+                ),
+            ],
+        }
+
+        for cluster, properties in properties_dict.items():
+            if not self.weaviate_client.collections.exists(cluster):
+                self.weaviate_client.collections.create(
+                    name=cluster, properties=properties
+                )
+
+    def _initialize_weaviate_client(self):
+        if self.weaviate_api_key:
+            return weaviate.connect_to_wcs(
+                cluster_url=self.weaviate_url,
+                auth_credentials=weaviate.auth.AuthApiKey(self.weaviate_api_key),
+                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
+                skip_init_checks=True,
+            )
+        else:
+            return weaviate.connect_to_local(
+                port=self.weaviate_port,
+                grpc_port=self.weaviate_grpc_port,
+                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
+                skip_init_checks=True,
+            )
+
+    def generate_embedding(self, data: str, **kwargs):
+        embedding_model = TextEmbedding(model_name=self.fastembed_model)
+        embedding = next(embedding_model.embed(data))
+        return embedding.tolist()
+
+    def _insert_data(self, cluster_key: str, data_object: dict, vector: list) -> str:
+        self.weaviate_client.connect()
+        response = self.weaviate_client.collections.get(
+            self.training_data_cluster[cluster_key]
+        ).data.insert(properties=data_object, vector=vector)
+        self.weaviate_client.close()
+        return response
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        data_object = {
+            "description": ddl,
+        }
+        response = self._insert_data("ddl", data_object, self.generate_embedding(ddl))
+        return f"{response}-ddl"
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        data_object = {
+            "description": doc,
+        }
+        response = self._insert_data("doc", data_object, self.generate_embedding(doc))
+        return f"{response}-doc"
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        data_object = {
+            "sql": sql,
+            "natural_language_question": question,
+        }
+        response = self._insert_data(
+            "sql", data_object, self.generate_embedding(question)
+        )
+        return f"{response}-sql"
+
+    def _query_collection(
+        self, cluster_key: str, vector_input: list, return_properties: list
+    ) -> list:
+        self.weaviate_client.connect()
+        collection = self.weaviate_client.collections.get(
+            self.training_data_cluster[cluster_key]
+        )
+        response = collection.query.near_vector(
+            near_vector=vector_input,
+            limit=self.n_results,
+            return_properties=return_properties,
+        )
+        response_list = [item.properties for item in response.objects]
+        self.weaviate_client.close()
+        return response_list
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection("ddl", vector_input, ["description"])
+        return [item["description"] for item in response_list]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection("doc", vector_input, ["description"])
+        return [item["description"] for item in response_list]
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection(
+            "sql", vector_input, ["sql", "natural_language_question"]
+        )
+        return [
+            {"question": item["natural_language_question"], "sql": item["sql"]}
+            for item in response_list
+        ]
+
+    def get_training_data(self, **kwargs) -> list:
+        self.weaviate_client.connect()
+        combined_response_list = []
+        for collection_name in self.training_data_cluster.values():
+            if self.weaviate_client.collections.exists(collection_name):
+                collection = self.weaviate_client.collections.get(collection_name)
+                response_list = [item.properties for item in collection.iterator()]
+                combined_response_list.extend(response_list)
+        self.weaviate_client.close()
+        return combined_response_list
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        self.weaviate_client.connect()
+        success = False
+        if id.endswith("-sql"):
+            id = id.replace("-sql", "")
+            success = self.weaviate_client.collections.get(
+                self.training_data_cluster["sql"]
+            ).data.delete_by_id(id)
+        elif id.endswith("-ddl"):
+            id = id.replace("-ddl", "")
+            success = self.weaviate_client.collections.get(
+                self.training_data_cluster["ddl"]
+            ).data.delete_by_id(id)
+        elif id.endswith("-doc"):
+            id = id.replace("-doc", "")
+            success = self.weaviate_client.collections.get(
+                self.training_data_cluster["doc"]
+            ).data.delete_by_id(id)
+        self.weaviate_client.close()
+        return success

vanna/legacy/xinference/__init__.py

@@ -0,0 +1 @@
+from .xinference import Xinference

vanna/legacy/xinference/xinference.py

@@ -0,0 +1,53 @@
+from xinference_client.client.restful.restful_client import (
+    Client,
+    RESTfulChatModelHandle,
+)
+
+from ..base import VannaBase
+
+
+class Xinference(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        if not config or "base_url" not in config:
+            raise ValueError("config must contain at least Xinference base_url")
+
+        base_url = config["base_url"]
+        api_key = config.get("api_key", "not empty")
+        self.xinference_client = Client(base_url=base_url, api_key=api_key)
+
+    def system_message(self, message: str) -> any:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> any:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> any:
+        return {"role": "assistant", "content": message}
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        if prompt is None:
+            raise Exception("Prompt is None")
+
+        if len(prompt) == 0:
+            raise Exception("Prompt is empty")
+
+        num_tokens = 0
+        for message in prompt:
+            num_tokens += len(message["content"]) / 4
+
+        model_uid = kwargs.get("model_uid") or self.config.get("model_uid", None)
+        if model_uid is None:
+            raise ValueError("model_uid is required")
+
+        xinference_model = self.xinference_client.get_model(model_uid)
+        if isinstance(xinference_model, RESTfulChatModelHandle):
+            print(f"Using model_uid {model_uid} for {num_tokens} tokens (approx)")
+
+            response = xinference_model.chat(prompt)
+            return response["choices"][0]["message"]["content"]
+        else:
+            raise NotImplementedError(
+                f"Xinference model handle type {type(xinference_model)} is not supported, required RESTfulChatModelHandle"
+            )

vanna/servers/__init__.py

@@ -0,0 +1,16 @@
+"""
+Server implementations for the Vanna Agents framework.
+
+This module provides Flask and FastAPI server factories for serving
+Vanna agents over HTTP with SSE, WebSocket, and polling endpoints.
+"""
+
+from .base import ChatHandler, ChatRequest, ChatStreamChunk
+from .cli.server_runner import ExampleAgentLoader
+
+__all__ = [
+    "ChatHandler",
+    "ChatRequest",
+    "ChatStreamChunk",
+    "ExampleAgentLoader",
+]

vanna/servers/__main__.py

@@ -0,0 +1,8 @@
+"""
+Entry point for running Vanna Agents servers.
+"""
+
+from .cli.server_runner import main
+
+if __name__ == "__main__":
+    main()

vanna/servers/base/__init__.py

@@ -0,0 +1,18 @@
+"""
+Base server components for the Vanna Agents framework.
+
+This module provides framework-agnostic components for handling chat
+requests and responses.
+"""
+
+from .chat_handler import ChatHandler
+from .models import ChatRequest, ChatStreamChunk, ChatResponse
+from .templates import INDEX_HTML
+
+__all__ = [
+    "ChatHandler",
+    "ChatRequest",
+    "ChatStreamChunk",
+    "ChatResponse",
+    "INDEX_HTML",
+]

vanna/servers/base/chat_handler.py

@@ -0,0 +1,65 @@
+"""
+Framework-agnostic chat handling logic.
+"""
+
+import uuid
+from typing import AsyncGenerator, List
+
+from ...core import Agent
+from .models import ChatRequest, ChatResponse, ChatStreamChunk
+
+
+class ChatHandler:
+    """Core chat handling logic - framework agnostic."""
+
+    def __init__(
+        self,
+        agent: Agent,
+    ):
+        """Initialize chat handler.
+
+        Args:
+            agent: The agent to handle chat requests
+        """
+        self.agent = agent
+
+    async def handle_stream(
+        self, request: ChatRequest
+    ) -> AsyncGenerator[ChatStreamChunk, None]:
+        """Stream chat responses.
+
+        Args:
+            request: Chat request
+
+        Yields:
+            Chat stream chunks
+        """
+        conversation_id = request.conversation_id or self._generate_conversation_id()
+        # Use request_id from client for tracking, or use the one generated internally
+        request_id = request.request_id or str(uuid.uuid4())
+
+        async for component in self.agent.send_message(
+            request_context=request.request_context,
+            message=request.message,
+            conversation_id=conversation_id,
+        ):
+            yield ChatStreamChunk.from_component(component, conversation_id, request_id)
+
+    async def handle_poll(self, request: ChatRequest) -> ChatResponse:
+        """Handle polling-based chat.
+
+        Args:
+            request: Chat request
+
+        Returns:
+            Complete chat response
+        """
+        chunks = []
+        async for chunk in self.handle_stream(request):
+            chunks.append(chunk)
+
+        return ChatResponse.from_chunks(chunks)
+
+    def _generate_conversation_id(self) -> str:
+        """Generate new conversation ID."""
+        return f"conv_{uuid.uuid4().hex[:8]}"

vanna/servers/base/models.py

@@ -0,0 +1,111 @@
+"""
+Request and response models for server endpoints.
+"""
+
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from ...components import UiComponent, RichComponent
+from ...core.component_manager import ComponentUpdate
+from ...core.user.request_context import RequestContext
+
+
+class ChatRequest(BaseModel):
+    """Request model for chat endpoints."""
+
+    message: str = Field(description="User message")
+    conversation_id: Optional[str] = Field(default=None, description="Conversation ID")
+    request_id: Optional[str] = Field(
+        default=None, description="Request ID for tracing"
+    )
+    request_context: RequestContext = Field(
+        default_factory=RequestContext,
+        description="Request context for user resolution",
+    )
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional metadata"
+    )
+
+
+class ChatStreamChunk(BaseModel):
+    """Single chunk in a streaming chat response."""
+
+    rich: Dict[str, Any] = Field(description="Rich component data for advanced UIs")
+    simple: Optional[Dict[str, Any]] = Field(
+        default=None, description="Simple component data for basic UIs"
+    )
+
+    # Stream metadata
+    conversation_id: str = Field(description="Conversation ID")
+    request_id: str = Field(description="Request ID")
+    timestamp: float = Field(default_factory=time.time, description="Timestamp")
+
+    @classmethod
+    def from_component(
+        cls,
+        component: Union[UiComponent, RichComponent],
+        conversation_id: str,
+        request_id: str,
+    ) -> "ChatStreamChunk":
+        """Create chunk from UI component or rich component."""
+
+        if isinstance(component, UiComponent):
+            # Full UiComponent with both rich and simple
+            rich_data = component.rich_component.serialize_for_frontend()
+            simple_data = None
+            if component.simple_component:
+                simple_data = component.simple_component.serialize_for_frontend()
+
+            return cls(
+                rich=rich_data,
+                simple=simple_data,
+                conversation_id=conversation_id,
+                request_id=request_id,
+            )
+
+        # Rich component only (no simple fallback)
+        rich_data = component.serialize_for_frontend()
+        return cls(
+            rich=rich_data,
+            simple=None,
+            conversation_id=conversation_id,
+            request_id=request_id,
+        )
+
+    @classmethod
+    def from_component_update(
+        cls, update: ComponentUpdate, conversation_id: str, request_id: str
+    ) -> "ChatStreamChunk":
+        """Create chunk from component update."""
+        update_payload = update.serialize_for_frontend()
+        return cls(
+            rich=update_payload,
+            simple=None,  # Component updates don't have simple representations
+            conversation_id=conversation_id,
+            request_id=request_id,
+        )
+
+
+class ChatResponse(BaseModel):
+    """Complete chat response for polling endpoints."""
+
+    chunks: List[ChatStreamChunk] = Field(description="Response chunks")
+    conversation_id: str = Field(description="Conversation ID")
+    request_id: str = Field(description="Request ID")
+    total_chunks: int = Field(description="Total number of chunks")
+
+    @classmethod
+    def from_chunks(cls, chunks: List[ChatStreamChunk]) -> "ChatResponse":
+        """Create response from chunks."""
+        if not chunks:
+            return cls(chunks=[], conversation_id="", request_id="", total_chunks=0)
+
+        return cls(
+            chunks=chunks,
+            conversation_id=chunks[0].conversation_id,
+            request_id=chunks[0].request_id,
+            total_chunks=len(chunks),
+        )

vanna/servers/base/rich_chat_handler.py

@@ -0,0 +1,141 @@
+# """
+# Rich component-aware chat handling logic.
+# """
+
+# import uuid
+# from typing import AsyncGenerator, Callable, List, Optional, Union
+
+# from ...core import Agent, User
+# from ...core.rich_components import RichComponent
+# from ...core.component_manager import ComponentManager, ComponentUpdate
+# from .models import ChatRequest, ChatResponse, ChatStreamChunk
+
+
+# class RichChatHandler:
+#     """Rich component-aware chat handling logic."""
+
+#     def __init__(
+#         self,
+#         agent: Agent,
+#         default_user_factory: Optional[Callable[[Optional[str]], User]] = None,
+#     ):
+#         """Initialize rich chat handler.
+
+#         Args:
+#             agent: The agent to handle chat requests
+#             default_user_factory: Function to create default user from user_id
+#         """
+#         self.agent = agent
+#         self.default_user_factory = default_user_factory or self._create_default_user
+#         self.component_managers: dict[str, ComponentManager] = {}  # Per conversation
+
+#     async def handle_stream(
+#         self, request: ChatRequest
+#     ) -> AsyncGenerator[ChatStreamChunk, None]:
+#         """Stream chat responses with rich component support.
+
+#         Args:
+#             request: Chat request
+
+#         Yields:
+#             Chat stream chunks including rich component updates
+#         """
+#         user = self._resolve_user(request.user_id)
+#         conversation_id = request.conversation_id or self._generate_conversation_id()
+#         request_id = request.request_id or str(uuid.uuid4())
+
+#         # Get or create component manager for this conversation
+#         if conversation_id not in self.component_managers:
+#             self.component_managers[conversation_id] = ComponentManager()
+
+#         component_manager = self.component_managers[conversation_id]
+
+#         async for component in self.agent.send_message(
+#             conversation_id=conversation_id,
+#             user=user,
+#             message=request.message,
+#             request_id=request_id,
+#         ):
+#             if isinstance(component, RichComponent):
+#                 # Handle rich component through manager
+#                 update = component_manager.emit(component)
+#                 yield ChatStreamChunk.from_component_update(update, conversation_id, request_id)
+#             else:
+#                 # Handle legacy components
+#                 yield ChatStreamChunk.from_component(component, conversation_id, request_id)
+
+#     async def handle_poll(self, request: ChatRequest) -> ChatResponse:
+#         """Handle polling request with rich component support.
+
+#         Args:
+#             request: Chat request
+
+#         Returns:
+#             Complete chat response with all components
+#         """
+#         chunks: List[ChatStreamChunk] = []
+
+#         async for chunk in self.handle_stream(request):
+#             chunks.append(chunk)
+
+#         return ChatResponse.from_chunks(chunks)
+
+#     def get_component_manager(self, conversation_id: str) -> Optional[ComponentManager]:
+#         """Get the component manager for a conversation."""
+#         return self.component_managers.get(conversation_id)
+
+#     def get_component(self, conversation_id: str, component_id: str) -> Optional[RichComponent]:
+#         """Get a specific component from a conversation."""
+#         manager = self.get_component_manager(conversation_id)
+#         return manager.get_component(component_id) if manager else None
+
+#     def get_all_components(self, conversation_id: str) -> List[RichComponent]:
+#         """Get all components in a conversation."""
+#         manager = self.get_component_manager(conversation_id)
+#         return manager.get_all_components() if manager else []
+
+#     def update_component(
+#         self,
+#         conversation_id: str,
+#         component_id: str,
+#         **updates
+#     ) -> Optional[ComponentUpdate]:
+#         """Update a component in a conversation."""
+#         manager = self.get_component_manager(conversation_id)
+#         return manager.update_component(component_id, **updates) if manager else None
+
+#     def remove_component(
+#         self,
+#         conversation_id: str,
+#         component_id: str
+#     ) -> Optional[ComponentUpdate]:
+#         """Remove a component from a conversation."""
+#         manager = self.get_component_manager(conversation_id)
+#         return manager.remove_component(component_id) if manager else None
+
+#     def clear_conversation_components(self, conversation_id: str):
+#         """Clear all components for a conversation."""
+#         if conversation_id in self.component_managers:
+#             del self.component_managers[conversation_id]
+
+#     def _resolve_user(self, user_id: Optional[str]) -> User:
+#         """Resolve user from ID or create default."""
+#         if user_id:
+#             # In a real implementation, you'd fetch from a user store
+#             return User(id=user_id, username=f"user_{user_id}", email="", permissions=[])
+
+#         return self.default_user_factory(user_id)
+
+#     def _create_default_user(self, user_id: Optional[str]) -> User:
+#         """Create a default user."""
+#         user_id = user_id or "anonymous"
+#         return User(
+#             id=user_id,
+#             username=f"user_{user_id}",
+#             email="",
+#             permissions=[]
+#         )
+
+#     def _generate_conversation_id(self) -> str:
+#         """Generate a new conversation ID."""
+#         return str(uuid.uuid4())