Vector Store

class dbt_llm_tools.VectorStore(openai_api_key: str, embedding_model_name: str = 'text-embedding-3-large', vector_db_path: str = '.local_storage/chroma.db', test_mode: bool = False)

A class representing a vector store for dbt models.

Methods:

get_client: Returns the client object for the vector store. upsert_models: Upsert the models into the vector store. reset_collection: Clear the collection of all documents.

get_client() → PersistentClient

Returns the client object for the vector store.

Returns:

chromadb.PersistentClient: The client object for the vector store.

get_models(model_ids: list[str] | None = None) → list[DbtModel]

Get the models from the vector store.

Args:

model_ids (list[str], optional): A list of model ids to be retrieved from the vector store.

Returns:

list[DbtModel]: A list of dbt model objects retrieved from the vector store.

query_collection(query: str, n_results: int = 3) → list[ParsedSearchResult]

Query the collection for the k nearest neighbours to the query.

Args:

query (str): The query to be used for nearest neighbour search. n_results (int, optional): The number of nearest neighbours to be returned. Defaults to 3.

Returns:

list[ParsedSearchResult]: A list of parsed search results.

reset_collection() → None

Clear the collection of all documents.

Returns:

None

set_embedding_fn(embedding_model_name: str) → None

Set the embedding function for the vector store.

Args:

embedding_model_name (str): The name of the OpenAI embedding model to be used.

upsert_models(models: list[DbtModel]) → None

Upsert the models into the vector store.

Args:

models (list[DbtModel]): A list of dbt model objects to be upserted into the vector store.

Returns:

None