Parameterize number of documentation documents retrieved

[GiacomoPracucci]

I am working with ChromaDB with many small documentation documents and need to increase the number of retrieved documentation documents from the vectorstore.

I modified the chromadb_vector.py file in this way, trying to keep the n_results referring to sql documents and the n_results referring to documentation documents separate.

class ChromaDB_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if config is not None:
            path = config.get("path", ".")
            self.embedding_function = config.get("embedding_function", default_ef)
            curr_client = config.get("client", "persistent")
            self.n_results_sql = config.get("n_results_sql", 10)
            self.n_results_documentation = config.get("n_results_documentation", 20)
        else:
            path = "."
            self.embedding_function = default_ef
            curr_client = "persistent"  # defaults to persistent storage
            self.n_results_sql = 10  # defaults to 10 sql documents
            self.n_results_documentation = 20 # defaults to 20 documentation documents

and

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        return ChromaDB_VectorStore._extract_documents(
            self.sql_collection.query(
                query_texts=[question],
                n_results=self.n_results_sql,
            )
        )

    def get_related_documentation(self, question: str, **kwargs) -> list:
        return ChromaDB_VectorStore._extract_documents(
            self.documentation_collection.query(
                query_texts=[question],
                n_results=self.n_results_documentation,
            )
        )

The code correctly returns the number of sql and documentation documents given in the parameters. Does the procedure look correct to you? Could it be worth opening a PR?

Best regards

[andreped]

@GiacomoPracucci I just had this issue myself. The solution is technically sound to me. Feel free to open a PR and others can review and test the solution there :]

[andreped]

Would be great if you could make a feature request under Issues about it, such that it is easier to track which PR addresses which issues.

[Mioooooo]

@GiacomoPracucci Does the latest version of vanna already include this feature?

[GiacomoPracucci]

I would say yes, unless there were changes in later releases that changed this