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

vanna/{google → legacy/google}/bigquery_vector.py

@@ -2,10 +2,7 @@ import datetime
 import os
 import uuid
 from typing import List, Optional
-from vertexai.language_models import (
-  TextEmbeddingInput,
-  TextEmbeddingModel
-)
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
 
 import pandas as pd
 from google.cloud import bigquery
@@ -18,7 +15,9 @@ class BigQuery_VectorStore(VannaBase):
         self.config = config
 
         self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
-        self.n_results_documentation = config.get("n_results_documentation", config.get("n_results", 10))
+        self.n_results_documentation = config.get(
+            "n_results_documentation", config.get("n_results", 10)
+        )
         self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
 
         if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
@@ -47,7 +46,7 @@ class BigQuery_VectorStore(VannaBase):
 
         self.conn = bigquery.Client(project=self.project_id)
 
-        dataset_name = self.config.get('bigquery_dataset_name', 'vanna_managed')
+        dataset_name = self.config.get("bigquery_dataset_name", "vanna_managed")
         self.dataset_id = f"{self.project_id}.{dataset_name}"
         dataset = bigquery.Dataset(self.dataset_id)
 
@@ -101,21 +100,35 @@ class BigQuery_VectorStore(VannaBase):
         # except Exception as e:
         #     print(f"Failed to create vector index: {e}")
 
-    def store_training_data(self, training_data_type: str, question: str, content: str, embedding: List[float], **kwargs) -> str:
+    def store_training_data(
+        self,
+        training_data_type: str,
+        question: str,
+        content: str,
+        embedding: List[float],
+        **kwargs,
+    ) -> str:
         id = str(uuid.uuid4())
         created_at = datetime.datetime.now()
-        self.conn.insert_rows_json(self.table_id, [{
-            "id": id,
-            "training_data_type": training_data_type,
-            "question": question,
-            "content": content,
-            "embedding": embedding,
-            "created_at": created_at.isoformat()
-        }])
+        self.conn.insert_rows_json(
+            self.table_id,
+            [
+                {
+                    "id": id,
+                    "training_data_type": training_data_type,
+                    "question": question,
+                    "content": content,
+                    "embedding": embedding,
+                    "created_at": created_at.isoformat(),
+                }
+            ],
+        )
 
         return id
 
-    def fetch_similar_training_data(self, training_data_type: str, question: str, n_results, **kwargs) -> pd.DataFrame:
+    def fetch_similar_training_data(
+        self, training_data_type: str, question: str, n_results, **kwargs
+    ) -> pd.DataFrame:
         question_embedding = self.generate_question_embedding(question)
 
         query = f"""
@@ -145,29 +158,28 @@ class BigQuery_VectorStore(VannaBase):
         embeddings = None
 
         if self.type == "VERTEX_AI":
-          input = [TextEmbeddingInput(data, task)]
-          model = TextEmbeddingModel.from_pretrained("text-embedding-004")
+            input = [TextEmbeddingInput(data, task)]
+            model = TextEmbeddingModel.from_pretrained("text-embedding-004")
 
-          result = model.get_embeddings(input)
+            result = model.get_embeddings(input)
 
-          if len(result) > 0:
-              embeddings = result[0].values
+            if len(result) > 0:
+                embeddings = result[0].values
         else:
-          # Use Gemini Consumer API
-          result = self.genai.embed_content(
-            model="models/text-embedding-004",
-            content=data,
-            task_type=task)
+            # Use Gemini Consumer API
+            result = self.genai.embed_content(
+                model="models/text-embedding-004", content=data, task_type=task
+            )
 
-          if 'embedding' in result:
-            embeddings = result['embedding']
+            if "embedding" in result:
+                embeddings = result["embedding"]
 
         return embeddings
 
     def generate_question_embedding(self, data: str, **kwargs) -> List[float]:
         result = self.get_embeddings(data, "RETRIEVAL_QUERY")
 
-        if result != None:
+        if result is not None:
             return result
         else:
             raise ValueError("No embeddings returned")
@@ -175,7 +187,7 @@ class BigQuery_VectorStore(VannaBase):
     def generate_storage_embedding(self, data: str, **kwargs) -> List[float]:
         result = self.get_embeddings(data, "RETRIEVAL_DOCUMENT")
 
-        if result != None:
+        if result is not None:
             return result
         else:
             raise ValueError("No embeddings returned")
@@ -195,45 +207,66 @@ class BigQuery_VectorStore(VannaBase):
         return self.generate_storage_embedding(data, **kwargs)
 
     def get_similar_question_sql(self, question: str, **kwargs) -> list:
-        df = self.fetch_similar_training_data(training_data_type="sql", question=question, n_results=self.n_results_sql)
+        df = self.fetch_similar_training_data(
+            training_data_type="sql", question=question, n_results=self.n_results_sql
+        )
 
         # Return a list of dictionaries with only question, sql fields. The content field needs to be renamed to sql
-        return df.rename(columns={"content": "sql"})[["question", "sql"]].to_dict(orient="records")
+        return df.rename(columns={"content": "sql"})[["question", "sql"]].to_dict(
+            orient="records"
+        )
 
     def get_related_ddl(self, question: str, **kwargs) -> list:
-        df = self.fetch_similar_training_data(training_data_type="ddl", question=question, n_results=self.n_results_ddl)
+        df = self.fetch_similar_training_data(
+            training_data_type="ddl", question=question, n_results=self.n_results_ddl
+        )
 
         # Return a list of strings of the content
         return df["content"].tolist()
 
     def get_related_documentation(self, question: str, **kwargs) -> list:
-        df = self.fetch_similar_training_data(training_data_type="documentation", question=question, n_results=self.n_results_documentation)
+        df = self.fetch_similar_training_data(
+            training_data_type="documentation",
+            question=question,
+            n_results=self.n_results_documentation,
+        )
 
         # Return a list of strings of the content
         return df["content"].tolist()
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
-        doc = {
-            "question": question,
-            "sql": sql
-        }
+        doc = {"question": question, "sql": sql}
 
         embedding = self.generate_embedding(str(doc))
 
-        return self.store_training_data(training_data_type="sql", question=question, content=sql, embedding=embedding)
+        return self.store_training_data(
+            training_data_type="sql",
+            question=question,
+            content=sql,
+            embedding=embedding,
+        )
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
         embedding = self.generate_embedding(ddl)
 
-        return self.store_training_data(training_data_type="ddl", question="", content=ddl, embedding=embedding)
+        return self.store_training_data(
+            training_data_type="ddl", question="", content=ddl, embedding=embedding
+        )
 
     def add_documentation(self, documentation: str, **kwargs) -> str:
         embedding = self.generate_embedding(documentation)
 
-        return self.store_training_data(training_data_type="documentation", question="", content=documentation, embedding=embedding)
+        return self.store_training_data(
+            training_data_type="documentation",
+            question="",
+            content=documentation,
+            embedding=embedding,
+        )
 
     def get_training_data(self, **kwargs) -> pd.DataFrame:
-        query = f"SELECT id, training_data_type, question, content FROM `{self.table_id}`"
+        query = (
+            f"SELECT id, training_data_type, question, content FROM `{self.table_id}`"
+        )
 
         return self.conn.query(query).result().to_dataframe()
 

vanna/{google → legacy/google}/gemini_chat.py

@@ -35,14 +35,18 @@ class GoogleGeminiChat(VannaBase):
             import vertexai
             from vertexai.generative_models import GenerativeModel
 
-            json_file_path = config.get("google_credentials")  # Assuming the JSON file path is provided in the config
+            json_file_path = config.get(
+                "google_credentials"
+            )  # Assuming the JSON file path is provided in the config
 
             if not json_file_path or not os.path.exists(json_file_path):
-                raise FileNotFoundError(f"JSON credentials file not found at: {json_file_path}")
+                raise FileNotFoundError(
+                    f"JSON credentials file not found at: {json_file_path}"
+                )
 
             try:
                 # Validate and set the JSON file path for GOOGLE_APPLICATION_CREDENTIALS
-                os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = json_file_path
+                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = json_file_path
 
                 # Initialize VertexAI with the credentials
                 credentials, _ = google.auth.default()

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

@@ -61,7 +61,6 @@ class Hf(VannaBase):
         return self.extract_sql_query(sql)
 
     def submit_prompt(self, prompt, **kwargs) -> str:
-
         input_ids = self.tokenizer.apply_chat_template(
             prompt, add_generation_prompt=True, return_tensors="pt"
         ).to(self.model.device)

vanna/{milvus → legacy/milvus}/milvus_vector.py

@@ -33,6 +33,7 @@ class Milvus_VectorStore(VannaBase):
                 For more models, please refer to:
                 https://milvus.io/docs/embeddings.md
     """
+
     def __init__(self, config=None):
         VannaBase.__init__(self, config=config)
 
@@ -45,7 +46,9 @@ class Milvus_VectorStore(VannaBase):
             self.embedding_function = config.get("embedding_function")
         else:
             self.embedding_function = model.DefaultEmbeddingFunction()
-        self._embedding_dim = self.embedding_function.encode_documents(["foo"])[0].shape[0]
+        self._embedding_dim = self.embedding_function.encode_documents(["foo"])[
+            0
+        ].shape[0]
         self._create_collections()
         self.n_results = config.get("n_results", 10)
 
@@ -54,21 +57,32 @@ class Milvus_VectorStore(VannaBase):
         self._create_ddl_collection("vannaddl")
         self._create_doc_collection("vannadoc")
 
-
     def generate_embedding(self, data: str, **kwargs) -> List[float]:
         return self.embedding_function.encode_documents(data).tolist()
 
-
     def _create_sql_collection(self, name: str):
         if not self.milvus_client.has_collection(collection_name=name):
             vannasql_schema = MilvusClient.create_schema(
                 auto_id=False,
                 enable_dynamic_field=False,
             )
-            vannasql_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
-            vannasql_schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
-            vannasql_schema.add_field(field_name="sql", datatype=DataType.VARCHAR, max_length=65535)
-            vannasql_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+            vannasql_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannasql_schema.add_field(
+                field_name="text", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannasql_schema.add_field(
+                field_name="sql", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannasql_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
 
             vannasql_index_params = self.milvus_client.prepare_index_params()
             vannasql_index_params.add_index(
@@ -81,7 +95,7 @@ class Milvus_VectorStore(VannaBase):
                 collection_name=name,
                 schema=vannasql_schema,
                 index_params=vannasql_index_params,
-                consistency_level="Strong"
+                consistency_level="Strong",
             )
 
     def _create_ddl_collection(self, name: str):
@@ -90,9 +104,20 @@ class Milvus_VectorStore(VannaBase):
                 auto_id=False,
                 enable_dynamic_field=False,
             )
-            vannaddl_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
-            vannaddl_schema.add_field(field_name="ddl", datatype=DataType.VARCHAR, max_length=65535)
-            vannaddl_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+            vannaddl_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannaddl_schema.add_field(
+                field_name="ddl", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannaddl_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
 
             vannaddl_index_params = self.milvus_client.prepare_index_params()
             vannaddl_index_params.add_index(
@@ -105,7 +130,7 @@ class Milvus_VectorStore(VannaBase):
                 collection_name=name,
                 schema=vannaddl_schema,
                 index_params=vannaddl_index_params,
-                consistency_level="Strong"
+                consistency_level="Strong",
             )
 
     def _create_doc_collection(self, name: str):
@@ -114,9 +139,20 @@ class Milvus_VectorStore(VannaBase):
                 auto_id=False,
                 enable_dynamic_field=False,
             )
-            vannadoc_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
-            vannadoc_schema.add_field(field_name="doc", datatype=DataType.VARCHAR, max_length=65535)
-            vannadoc_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+            vannadoc_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannadoc_schema.add_field(
+                field_name="doc", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannadoc_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
 
             vannadoc_index_params = self.milvus_client.prepare_index_params()
             vannadoc_index_params.add_index(
@@ -129,7 +165,7 @@ class Milvus_VectorStore(VannaBase):
                 collection_name=name,
                 schema=vannadoc_schema,
                 index_params=vannadoc_index_params,
-                consistency_level="Strong"
+                consistency_level="Strong",
             )
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
@@ -139,12 +175,7 @@ class Milvus_VectorStore(VannaBase):
         embedding = self.embedding_function.encode_documents([question])[0]
         self.milvus_client.insert(
             collection_name="vannasql",
-            data={
-                "id": _id,
-                "text": question,
-                "sql": sql,
-                "vector": embedding
-            }
+            data={"id": _id, "text": question, "sql": sql, "vector": embedding},
         )
         return _id
 
@@ -155,11 +186,7 @@ class Milvus_VectorStore(VannaBase):
         embedding = self.embedding_function.encode_documents([ddl])[0]
         self.milvus_client.insert(
             collection_name="vannaddl",
-            data={
-                "id": _id,
-                "ddl": ddl,
-                "vector": embedding
-            }
+            data={"id": _id, "ddl": ddl, "vector": embedding},
         )
         return _id
 
@@ -170,11 +197,7 @@ class Milvus_VectorStore(VannaBase):
         embedding = self.embedding_function.encode_documents([documentation])[0]
         self.milvus_client.insert(
             collection_name="vannadoc",
-            data={
-                "id": _id,
-                "doc": documentation,
-                "vector": embedding
-            }
+            data={"id": _id, "doc": documentation, "vector": embedding},
         )
         return _id
 
@@ -237,7 +260,7 @@ class Milvus_VectorStore(VannaBase):
             data=embeddings,
             limit=self.n_results,
             output_fields=["text", "sql"],
-            search_params=search_params
+            search_params=search_params,
         )
         res = res[0]
 
@@ -261,7 +284,7 @@ class Milvus_VectorStore(VannaBase):
             data=embeddings,
             limit=self.n_results,
             output_fields=["ddl"],
-            search_params=search_params
+            search_params=search_params,
         )
         res = res[0]
 
@@ -282,7 +305,7 @@ class Milvus_VectorStore(VannaBase):
             data=embeddings,
             limit=self.n_results,
             output_fields=["doc"],
-            search_params=search_params
+            search_params=search_params,
         )
         res = res[0]
 

vanna/{mock → legacy/mock}/llm.py

@@ -1,4 +1,3 @@
-
 from ..base import VannaBase
 
 

vanna/legacy/mock/vectordb.py

@@ -0,0 +1,67 @@
+import pandas as pd
+
+from ..base import VannaBase
+
+
+class MockVectorDB(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def _get_id(self, value: str, **kwargs) -> str:
+        # Hash the value and return the ID
+        return str(hash(value))
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        return self._get_id(ddl)
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        return self._get_id(doc)
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        return self._get_id(question)
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        return pd.DataFrame(
+            {
+                "id": {
+                    0: "19546-ddl",
+                    1: "91597-sql",
+                    2: "133976-sql",
+                    3: "59851-doc",
+                    4: "73046-sql",
+                },
+                "training_data_type": {
+                    0: "ddl",
+                    1: "sql",
+                    2: "sql",
+                    3: "documentation",
+                    4: "sql",
+                },
+                "question": {
+                    0: None,
+                    1: "What are the top selling genres?",
+                    2: "What are the low 7 artists by sales?",
+                    3: None,
+                    4: "What is the total sales for each customer?",
+                },
+                "content": {
+                    0: "CREATE TABLE [Invoice]\n(\n    [InvoiceId] INTEGER  NOT NULL,\n    [CustomerId] INTEGER  NOT NULL,\n    [InvoiceDate] DATETIME  NOT NULL,\n    [BillingAddress] NVARCHAR(70),\n    [BillingCity] NVARCHAR(40),\n    [BillingState] NVARCHAR(40),\n    [BillingCountry] NVARCHAR(40),\n    [BillingPostalCode] NVARCHAR(10),\n    [Total] NUMERIC(10,2)  NOT NULL,\n    CONSTRAINT [PK_Invoice] PRIMARY KEY  ([InvoiceId]),\n    FOREIGN KEY ([CustomerId]) REFERENCES [Customer] ([CustomerId]) \n\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n)",
+                    1: "SELECT g.Name AS Genre, SUM(il.Quantity) AS TotalSales\nFROM Genre g\nJOIN Track t ON g.GenreId = t.GenreId\nJOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY g.GenreId, g.Name\nORDER BY TotalSales DESC;",
+                    2: "SELECT a.ArtistId, a.Name, SUM(il.Quantity) AS TotalSales\nFROM Artist a\nINNER JOIN Album al ON a.ArtistId = al.ArtistId\nINNER JOIN Track t ON al.AlbumId = t.AlbumId\nINNER JOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY a.ArtistId, a.Name\nORDER BY TotalSales ASC\nLIMIT 7;",
+                    3: "This is a SQLite database. For dates rememeber to use SQLite syntax.",
+                    4: "SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) AS TotalSales\nFROM Customer c\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nGROUP BY c.CustomerId, c.FirstName, c.LastName;",
+                },
+            }
+        )
+
+    def remove_training_data(id: str, **kwargs) -> bool:
+        return True

vanna/legacy/ollama/ollama.py

@@ -0,0 +1,110 @@
+import json
+import re
+
+from httpx import Timeout
+
+from ..base import VannaBase
+from ..exceptions import DependencyError
+
+
+class Ollama(VannaBase):
+    def __init__(self, config=None):
+        try:
+            ollama = __import__("ollama")
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method, run command:"
+                " \npip install ollama"
+            )
+
+        if not config:
+            raise ValueError("config must contain at least Ollama model")
+        if "model" not in config.keys():
+            raise ValueError("config must contain at least Ollama model")
+        self.host = config.get("ollama_host", "http://localhost:11434")
+        self.model = config["model"]
+        if ":" not in self.model:
+            self.model += ":latest"
+
+        self.ollama_timeout = config.get("ollama_timeout", 240.0)
+
+        self.ollama_client = ollama.Client(
+            self.host, timeout=Timeout(self.ollama_timeout)
+        )
+        self.keep_alive = config.get("keep_alive", None)
+        self.ollama_options = config.get("options", {})
+        self.num_ctx = self.ollama_options.get("num_ctx", 2048)
+        self.__pull_model_if_ne(self.ollama_client, self.model)
+
+    @staticmethod
+    def __pull_model_if_ne(ollama_client, model):
+        model_response = ollama_client.list()
+        model_lists = [
+            model_element["model"] for model_element in model_response.get("models", [])
+        ]
+        if model not in model_lists:
+            ollama_client.pull(model)
+
+    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 extract_sql(self, llm_response):
+        """
+        Extracts the first SQL statement after the word 'select', ignoring case,
+        matches until the first semicolon, three backticks, or the end of the string,
+        and removes three backticks if they exist in the extracted string.
+
+        Args:
+        - llm_response (str): The string to search within for an SQL statement.
+
+        Returns:
+        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
+        """
+        # Remove ollama-generated extra characters
+        llm_response = llm_response.replace("\\_", "_")
+        llm_response = llm_response.replace("\\", "")
+
+        # Regular expression to find ```sql' and capture until '```'
+        sql = re.search(r"```sql\n((.|\n)*?)(?=;|\[|```)", llm_response, re.DOTALL)
+        # Regular expression to find 'select, with (ignoring case) and capture until ';', [ (this happens in case of mistral) or end of string
+        select_with = re.search(
+            r"(select|(with.*?as \())(.*?)(?=;|\[|```)",
+            llm_response,
+            re.IGNORECASE | re.DOTALL,
+        )
+        if sql:
+            self.log(f"Output from LLM: {llm_response} \nExtracted SQL: {sql.group(1)}")
+            return sql.group(1).replace("```", "")
+        elif select_with:
+            self.log(
+                f"Output from LLM: {llm_response} \nExtracted SQL: {select_with.group(0)}"
+            )
+            return select_with.group(0)
+        else:
+            return llm_response
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        self.log(
+            f"Ollama parameters:\n"
+            f"model={self.model},\n"
+            f"options={self.ollama_options},\n"
+            f"keep_alive={self.keep_alive}"
+        )
+        self.log(f"Prompt Content:\n{json.dumps(prompt, ensure_ascii=False)}")
+        response_dict = self.ollama_client.chat(
+            model=self.model,
+            messages=prompt,
+            stream=False,
+            options=self.ollama_options,
+            keep_alive=self.keep_alive,
+        )
+
+        self.log(f"Ollama Response:\n{str(response_dict)}")
+
+        return response_dict["message"]["content"]

vanna/{openai → legacy/openai}/openai_chat.py

@@ -65,9 +65,7 @@ class OpenAI_Chat(VannaBase):
 
         if kwargs.get("model", None) is not None:
             model = kwargs.get("model", None)
-            print(
-                f"Using model {model} for {num_tokens} tokens (approx)"
-            )
+            print(f"Using model {model} for {num_tokens} tokens (approx)")
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
@@ -76,9 +74,7 @@ class OpenAI_Chat(VannaBase):
             )
         elif kwargs.get("engine", None) is not None:
             engine = kwargs.get("engine", None)
-            print(
-                f"Using model {engine} for {num_tokens} tokens (approx)"
-            )
+            print(f"Using model {engine} for {num_tokens} tokens (approx)")
             response = self.client.chat.completions.create(
                 engine=engine,
                 messages=prompt,