nat.plugins.vanna.vanna_utils#

Attributes#

logger

Classes#

`VannaLangChainLLM`	LangChain LLM integration for Vanna framework.
`MilvusVectorStore`	Extended Milvus vector store for Vanna.
`VannaLangChain`	Combined Vanna implementation with Milvus and LangChain LLM.
`VannaSingleton`	Singleton manager for Vanna instances.

Functions#

`extract_json_from_string`(→ dict)	Extract JSON from a string that may contain additional content.
`remove_think_tags`(→ str)	Remove think tags from reasoning model output based on model type.
`to_langchain_msgs`(msgs)	Convert message dicts to LangChain message objects.
`train_vanna`(vn[, auto_train])	Train Vanna with DDL, documentation, and question-SQL examples.

Module Contents#

logger#

extract_json_from_string(content: str) → dict#

Extract JSON from a string that may contain additional content.

Args:: content: String containing JSON data
Returns:: Parsed JSON as dictionary
Raises:: ValueError: If no valid JSON found

remove_think_tags( text: str, model_name: str, reasoning_models: set[str], ) → str#

Remove think tags from reasoning model output based on model type.

Args:: text: Text potentially containing think tags model_name: Name of the model reasoning_models: Set of model names that require think tag removal
Returns:: Text with think tags removed if applicable

to_langchain_msgs(msgs)#: Convert message dicts to LangChain message objects.

class VannaLangChainLLM(client=None, config=None)#

Bases: vanna.legacy.base.VannaBase

LangChain LLM integration for Vanna framework.

client = None#

config#

dialect#

model = 'unknown'#

milvus_search_limit#

reasoning_models#

chat_models#

system_message(message: str) → dict#: Create system message.

user_message(message: str) → dict#: Create user message.

assistant_message(message: str) → dict#: Create assistant message.

get_training_sql_prompt(ddl_list: list, doc_list: list) → list#: Generate prompt for synthetic question-SQL pairs.

get_sql_prompt( initial_prompt: str | None, question: str, question_sql_list: list, ddl_list: list, doc_list: list, error_message: dict | None = None, \*\*kwargs, ) → list#: Generate prompt for SQL generation.

async submit_prompt(prompt, \*\*kwargs) → str#: Submit prompt to LLM.

class MilvusVectorStore(config=None)#

Bases: vanna.legacy.milvus.Milvus_VectorStore

Extended Milvus vector store for Vanna.

async _ensure_collections_created()#: Ensure all necessary Milvus collections are created (async).

async _create_sql_collection(name: str)#: Create SQL collection using async client.

async _create_ddl_collection(name: str)#: Create DDL collection using async client.

async _create_doc_collection(name: str)#: Create documentation collection using async client.

async add_question_sql(question: str, sql: str, \*\*kwargs) → str#: Add question-SQL pair to collection using async client.

async add_ddl(ddl: str, \*\*kwargs) → str#: Add DDL to collection using async client.

async add_documentation(documentation: str, \*\*kwargs) → str#: Add documentation to collection using async client.

async get_related_record(collection_name: str) → list#: Retrieve all related records using async client.

async get_similar_question_sql(question: str, \*\*kwargs) → list#: Get similar question-SQL pairs using async client.

async get_training_data(\*\*kwargs)#: Get all training data using async client.

async close()#: Close async Milvus client connection.

class VannaLangChain(client, config=None)#

Bases: MilvusVectorStore, VannaLangChainLLM

Combined Vanna implementation with Milvus and LangChain LLM.

Initialize VannaLangChain.

Args:: client: LangChain LLM client config: Configuration dict for Milvus vector store and LLM settings

db_engine = None#

async generate_sql( question: str, allow_llm_to_see_data: bool = False, error_message: dict | None = None, \*\*kwargs, ) → dict[str, str | None]#

Generate SQL using the LLM.

Args:: question: Natural language question to convert to SQL allow_llm_to_see_data: Whether to allow LLM to see actual data error_message: Optional error message from previous SQL execution kwargs: Additional keyword arguments
Returns:: Dictionary with ‘sql’ and optional ‘explanation’ keys

class VannaSingleton#

Singleton manager for Vanna instances.

_instance: VannaLangChain | None = None#

_lock: asyncio.Lock | None = None#

classmethod _get_lock() → asyncio.Lock#: Get or create the lock in the current event loop.

classmethod instance() → VannaLangChain | None#

Get current instance without creating one.

Returns:: Current Vanna instance or None if not initialized

classmethod get_instance( llm_client, embedder_client, async_milvus_client, dialect: str = 'SQLite', initial_prompt: str | None = None, n_results: int = 5, sql_collection: str = 'vanna_sql', ddl_collection: str = 'vanna_ddl', doc_collection: str = 'vanna_documentation', milvus_search_limit: int = 1000, reasoning_models: set[str] | None = None, chat_models: set[str] | None = None, create_collections: bool = True, ) → VannaLangChain#

Async:

Get or create a singleton Vanna instance.

Args:: llm_client: LangChain LLM client for SQL generation embedder_client: LangChain embedder for vector operations async_milvus_client: Async Milvus client dialect: SQL dialect (e.g., ‘databricks’, ‘postgres’, ‘mysql’) initial_prompt: Optional custom system prompt n_results: Number of similar examples to retrieve sql_collection: Collection name for SQL examples ddl_collection: Collection name for DDL doc_collection: Collection name for documentation milvus_search_limit: Maximum limit size for vector search operations reasoning_models: Models requiring special handling for think tags chat_models: Models using standard response handling create_collections: Whether to create Milvus collections if they don’t exist (default True)
Returns:: Initialized Vanna instance

classmethod reset()#

Async:

Reset the singleton Vanna instance.

Useful for testing or when configuration changes. Properly disposes of database engine if present.

async train_vanna(vn: VannaLangChain, auto_train: bool = False)#

Train Vanna with DDL, documentation, and question-SQL examples.

Args:: vn: Vanna instance auto_train: Whether to automatically train Vanna (auto-extract DDL and generate training data from database)