vanna 0.7.4__tar.gz → 0.7.5__tar.gz

{vanna-0.7.4 → vanna-0.7.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.7.4
+Version: 0.7.5
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -52,6 +52,7 @@ Requires-Dist: boto3 ; extra == "all"
 Requires-Dist: botocore ; extra == "all"
 Requires-Dist: langchain_core ; extra == "all"
 Requires-Dist: langchain_postgres ; extra == "all"
+Requires-Dist: xinference-client ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: azure-search-documents ; extra == "azuresearch"
 Requires-Dist: azure-identity ; extra == "azuresearch"
@@ -90,6 +91,7 @@ Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
 Requires-Dist: vllm ; extra == "vllm"
 Requires-Dist: weaviate-client ; extra == "weaviate"
+Requires-Dist: xinference-client ; extra == "xinference-client"
 Requires-Dist: zhipuai ; extra == "zhipuai"
 Project-URL: Bug Tracker, https://github.com/vanna-ai/vanna/issues
 Project-URL: Homepage, https://github.com/vanna-ai/vanna
@@ -122,6 +124,7 @@ Provides-Extra: snowflake
 Provides-Extra: test
 Provides-Extra: vllm
 Provides-Extra: weaviate
+Provides-Extra: xinference-client
 Provides-Extra: zhipuai
 
 

{vanna-0.7.4 → vanna-0.7.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.7.4"
+version = "0.7.5"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -33,7 +33,7 @@ bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]
 google = ["google-generativeai", "google-cloud-aiplatform"]
-all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres", "xinference-client"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
@@ -56,3 +56,4 @@ azuresearch = ["azure-search-documents", "azure-identity", "azure-common", "fast
 pgvector = ["langchain-postgres>=0.0.12"]
 faiss-cpu = ["faiss-cpu"]
 faiss-gpu = ["faiss-gpu"]
+xinference-client = ["xinference-client"]

vanna-0.7.5/src/vanna/xinference/__init__.py

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

vanna-0.7.5/src/vanna/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")