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

vanna/{vannadb → legacy/vannadb}/vannadb_vector.py

@@ -8,15 +8,15 @@ import requests
 from ..advanced import VannaAdvanced
 from ..base import VannaBase
 from ..types import (
-  DataFrameJSON,
-  NewOrganization,
-  OrganizationList,
-  Question,
-  QuestionSQLPair,
-  Status,
-  StatusWithId,
-  StringData,
-  TrainingData,
+    DataFrameJSON,
+    NewOrganization,
+    OrganizationList,
+    Question,
+    QuestionSQLPair,
+    Status,
+    StatusWithId,
+    StringData,
+    TrainingData,
 )
 from ..utils import sanitize_model_name
 
@@ -85,18 +85,25 @@ class VannaDB_VectorStore(VannaBase, VannaAdvanced):
             }
         """
 
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query})
+        response = requests.post(
+            self._graphql_endpoint, headers=self._graphql_headers, json={"query": query}
+        )
         response_json = response.json()
-        if response.status_code == 200 and 'data' in response_json and 'get_all_sql_functions' in response_json['data']:
-            self.log(response_json['data']['get_all_sql_functions'])
-            resp = response_json['data']['get_all_sql_functions']
+        if (
+            response.status_code == 200
+            and "data" in response_json
+            and "get_all_sql_functions" in response_json["data"]
+        ):
+            self.log(response_json["data"]["get_all_sql_functions"])
+            resp = response_json["data"]["get_all_sql_functions"]
 
             print(resp)
 
             return resp
         else:
-            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
-
+            raise Exception(
+                f"Query failed to run by returning code of {response.status_code}. {response.text}"
+            )
 
     def get_function(self, question: str, additional_data: dict = {}) -> dict:
         query = """
@@ -121,21 +128,38 @@ class VannaDB_VectorStore(VannaBase, VannaAdvanced):
             }
         }
         """
-        static_function_arguments = [{"name": key, "value": str(value)} for key, value in additional_data.items()]
-        variables = {"question": question, "staticFunctionArguments": static_function_arguments}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        static_function_arguments = [
+            {"name": key, "value": str(value)} for key, value in additional_data.items()
+        ]
+        variables = {
+            "question": question,
+            "staticFunctionArguments": static_function_arguments,
+        }
+        response = requests.post(
+            self._graphql_endpoint,
+            headers=self._graphql_headers,
+            json={"query": query, "variables": variables},
+        )
         response_json = response.json()
-        if response.status_code == 200 and 'data' in response_json and 'get_and_instantiate_function' in response_json['data']:
-            self.log(response_json['data']['get_and_instantiate_function'])
-            resp = response_json['data']['get_and_instantiate_function']
+        if (
+            response.status_code == 200
+            and "data" in response_json
+            and "get_and_instantiate_function" in response_json["data"]
+        ):
+            self.log(response_json["data"]["get_and_instantiate_function"])
+            resp = response_json["data"]["get_and_instantiate_function"]
 
             print(resp)
 
             return resp
         else:
-            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
+            raise Exception(
+                f"Query failed to run by returning code of {response.status_code}. {response.text}"
+            )
 
-    def create_function(self, question: str, sql: str, plotly_code: str, **kwargs) -> dict:
+    def create_function(
+        self, question: str, sql: str, plotly_code: str, **kwargs
+    ) -> dict:
         query = """
         mutation CreateFunction($question: String!, $sql: String!, $plotly_code: String!) {
             generate_and_create_sql_function(question: $question, sql: $sql, post_processing_code: $plotly_code) {
@@ -153,16 +177,27 @@ class VannaDB_VectorStore(VannaBase, VannaAdvanced):
         }
         """
         variables = {"question": question, "sql": sql, "plotly_code": plotly_code}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        response = requests.post(
+            self._graphql_endpoint,
+            headers=self._graphql_headers,
+            json={"query": query, "variables": variables},
+        )
         response_json = response.json()
-        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'generate_and_create_sql_function' in response_json['data']:
-            resp = response_json['data']['generate_and_create_sql_function']
+        if (
+            response.status_code == 200
+            and "data" in response_json
+            and response_json["data"] is not None
+            and "generate_and_create_sql_function" in response_json["data"]
+        ):
+            resp = response_json["data"]["generate_and_create_sql_function"]
 
             print(resp)
 
             return resp
         else:
-            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
+            raise Exception(
+                f"Query failed to run by returning code of {response.status_code}. {response.text}"
+            )
 
     def update_function(self, old_function_name: str, updated_function: dict) -> bool:
         """
@@ -187,41 +222,64 @@ class VannaDB_VectorStore(VannaBase, VannaAdvanced):
         """
 
         SQLFunctionUpdate = {
-            'function_name', 'description', 'arguments', 'sql_template', 'post_processing_code_template'
+            "function_name",
+            "description",
+            "arguments",
+            "sql_template",
+            "post_processing_code_template",
         }
 
         # Define the expected keys for each argument in the arguments list
-        ArgumentKeys = {'name', 'general_type', 'description', 'is_user_editable', 'available_values'}
+        ArgumentKeys = {
+            "name",
+            "general_type",
+            "description",
+            "is_user_editable",
+            "available_values",
+        }
 
         # Function to validate and transform arguments
         def validate_arguments(args):
             return [
-                {key: arg[key] for key in arg if key in ArgumentKeys}
-                for arg in args
+                {key: arg[key] for key in arg if key in ArgumentKeys} for arg in args
             ]
 
         # Keep only the keys that conform to the SQLFunctionUpdate GraphQL input type
-        updated_function = {key: value for key, value in updated_function.items() if key in SQLFunctionUpdate}
+        updated_function = {
+            key: value
+            for key, value in updated_function.items()
+            if key in SQLFunctionUpdate
+        }
 
         # Special handling for 'arguments' to ensure they conform to the spec
-        if 'arguments' in updated_function:
-            updated_function['arguments'] = validate_arguments(updated_function['arguments'])
+        if "arguments" in updated_function:
+            updated_function["arguments"] = validate_arguments(
+                updated_function["arguments"]
+            )
 
         variables = {
-            "input": {
-                "old_function_name": old_function_name,
-                **updated_function
-            }
+            "input": {"old_function_name": old_function_name, **updated_function}
         }
 
         print("variables", variables)
 
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response = requests.post(
+            self._graphql_endpoint,
+            headers=self._graphql_headers,
+            json={"query": mutation, "variables": variables},
+        )
         response_json = response.json()
-        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'update_sql_function' in response_json['data']:
-            return response_json['data']['update_sql_function']
+        if (
+            response.status_code == 200
+            and "data" in response_json
+            and response_json["data"] is not None
+            and "update_sql_function" in response_json["data"]
+        ):
+            return response_json["data"]["update_sql_function"]
         else:
-            raise Exception(f"Mutation failed to run by returning code of {response.status_code}. {response.text}")
+            raise Exception(
+                f"Mutation failed to run by returning code of {response.status_code}. {response.text}"
+            )
 
     def delete_function(self, function_name: str) -> bool:
         mutation = """
@@ -230,12 +288,23 @@ class VannaDB_VectorStore(VannaBase, VannaAdvanced):
         }
         """
         variables = {"function_name": function_name}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response = requests.post(
+            self._graphql_endpoint,
+            headers=self._graphql_headers,
+            json={"query": mutation, "variables": variables},
+        )
         response_json = response.json()
-        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'delete_sql_function' in response_json['data']:
-            return response_json['data']['delete_sql_function']
+        if (
+            response.status_code == 200
+            and "data" in response_json
+            and response_json["data"] is not None
+            and "delete_sql_function" in response_json["data"]
+        ):
+            return response_json["data"]["delete_sql_function"]
         else:
-            raise Exception(f"Mutation failed to run by returning code of {response.status_code}. {response.text}")
+            raise Exception(
+                f"Mutation failed to run by returning code of {response.status_code}. {response.text}"
+            )
 
     def create_model(self, model: str, **kwargs) -> bool:
         """

vanna/{vllm → legacy/vllm}/vllm.py

@@ -80,13 +80,12 @@ class Vllm(VannaBase):
         }
 
         if self.auth_key is not None:
-            headers = { 
-            'Content-Type': 'application/json', 
-            'Authorization': f'Bearer {self.auth_key}' 
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.auth_key}",
             }
 
-            response = requests.post(url, headers=headers,json=data)
-
+            response = requests.post(url, headers=headers, json=data)
 
         else:
             response = requests.post(url, json=data)
@@ -95,4 +94,4 @@ class Vllm(VannaBase):
 
         self.log(response.text)
 
-        return response_dict['choices'][0]['message']['content']
+        return response_dict["choices"][0]["message"]["content"]

vanna/{weaviate → legacy/weaviate}/weaviate_vector.py

@@ -6,7 +6,6 @@ from vanna.base import VannaBase
 
 
 class WeaviateDatabase(VannaBase):
-
     def __init__(self, config=None):
         """
         Initialize the VannaEnhanced class with the provided configuration.
@@ -42,30 +41,35 @@ class WeaviateDatabase(VannaBase):
         self.training_data_cluster = {
             "sql": "SQLTrainingDataEntry",
             "ddl": "DDLEntry",
-            "doc": "DocumentationEntry"
+            "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["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["doc"]: [
+                wvc.config.Property(
+                    name="description", data_type=wvc.config.DataType.TEXT
+                ),
             ],
-            self.training_data_cluster['sql']: [
+            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),
-            ]
+                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
+                    name=cluster, properties=properties
                 )
 
     def _initialize_weaviate_client(self):
@@ -74,28 +78,26 @@ class WeaviateDatabase(VannaBase):
                 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
+                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
+                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()
-
+        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
-        )
+        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
 
@@ -103,31 +105,37 @@ class WeaviateDatabase(VannaBase):
         data_object = {
             "description": ddl,
         }
-        response = self._insert_data('ddl', data_object, self.generate_embedding(ddl))
-        return f'{response}-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'
+        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'
+        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:
+    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])
+        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
+            return_properties=return_properties,
         )
         response_list = [item.properties for item in response.objects]
         self.weaviate_client.close()
@@ -135,18 +143,23 @@ class WeaviateDatabase(VannaBase):
 
     def get_related_ddl(self, question: str, **kwargs) -> list:
         vector_input = self.generate_embedding(question)
-        response_list = self._query_collection('ddl', vector_input, ["description"])
+        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"])
+        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]
+        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()
@@ -163,13 +176,19 @@ class WeaviateDatabase(VannaBase):
         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)
+            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)
+            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)
+            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/{xinference → legacy/xinference}/xinference.py

@@ -1,6 +1,6 @@
 from xinference_client.client.restful.restful_client import (
-  Client,
-  RESTfulChatModelHandle,
+    Client,
+    RESTfulChatModelHandle,
 )
 
 from ..base import VannaBase
@@ -43,11 +43,11 @@ class Xinference(VannaBase):
 
         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)"
-            )
+            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")
+            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]}"