vanna 0.7.0__tar.gz → 0.7.1__tar.gz

{vanna-0.7.0 → vanna-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.7.0
+Version: 0.7.1
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9

{vanna-0.7.0 → vanna-0.7.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.7.0"
+version = "0.7.1"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]

{vanna-0.7.0 → vanna-0.7.1}/src/vanna/weaviate/weaviate_vector.py

@@ -26,6 +26,7 @@ class WeaviateDatabase(VannaBase):
         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")
@@ -120,12 +121,12 @@ class WeaviateDatabase(VannaBase):
         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, limit: int = 3) -> 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])
         response = collection.query.near_vector(
             near_vector=vector_input,
-            limit=limit,
+            limit=self.n_results,
             return_properties=return_properties
         )
         response_list = [item.properties for item in response.objects]