Merge pull request #323 from andreped/import-tests

Add regular and shortcut import tests

pyproject.toml

@@ -31,11 +31,12 @@ mysql = ["PyMySQL"]
 bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]
-all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic", "zhipuai"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
 mistralai = ["mistralai"]
 anthropic = ["anthropic"]
 gemini = ["google-generativeai"]
 marqo = ["marqo"]
+zhipuai = ["zhipuai"]

src/vanna/ZhipuAI/ZhipuAI_Chat.py

@@ -220,19 +220,14 @@ def submit_prompt(
         if len(prompt) == 0:
             raise Exception("Prompt is empty")
 
-        client = ZhipuAI(api_key=self.api_key)  # 填写您自己的APIKey
+        client = ZhipuAI(api_key=self.api_key)
         response = client.chat.completions.create(
-            model="glm-4",  # 填写需要调用的模型名称
+            model="glm-4",
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             stop=stop,
             messages=prompt,
         )
-        # print(prompt)
-
-        # print(response)
-
-        # print(f"Cost {response.usage.total_tokens} token")
 
         return response.choices[0].message.content

tests/test_imports.py

@@ -0,0 +1,28 @@
+
+
+def test_regular_imports():
+    from vanna.anthropic.anthropic_chat import Anthropic_Chat
+    from vanna.base.base import VannaBase
+    from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
+    from vanna.local import LocalContext_OpenAI
+    from vanna.marqo.marqo import Marqo_VectorStore
+    from vanna.mistral.mistral import Mistral
+    from vanna.ollama.ollama import Ollama
+    from vanna.openai.openai_chat import OpenAI_Chat
+    from vanna.openai.openai_embeddings import OpenAI_Embeddings
+    from vanna.remote import VannaDefault
+    from vanna.vannadb.vannadb_vector import VannaDB_VectorStore
+    from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
+    from vanna.ZhipuAI.ZhipuAI_embeddings import ZhipuAI_Embeddings
+
+
+def test_shortcut_imports():
+    from vanna.anthropic import Anthropic_Chat
+    from vanna.base import VannaBase
+    from vanna.chromadb import ChromaDB_VectorStore
+    from vanna.marqo import Marqo_VectorStore
+    from vanna.mistral import Mistral
+    from vanna.ollama import Ollama
+    from vanna.openai import OpenAI_Chat, OpenAI_Embeddings
+    from vanna.vannadb import VannaDB_VectorStore
+    from vanna.ZhipuAI import ZhipuAI_Chat, ZhipuAI_Embeddings