class MyChromaDB(ChromaDB_VectorStore):def __init__(self, config=None):super().__init__(config)
self.rerank_model = CrossEncoder('bge-reranker-large',max_length=512)

def get_similar_question_sql(self, question: str, **kwargs) -> list:documents = ChromaDB_VectorStore._extract_documents(self.sql_collection.query(query_texts=[question],n_results=self.n_results_sql,))if documents != []:scores = self.rerank_model.predict([(question, doc["question"]) for doc in documents])combined = list(zip(documents, scores))sorted_combined = sorted(combined, key=lambda x: x[1], reverse=True)documents = [item[0] for item in sorted_combined]return documents
def get_related_ddl(self, question: str, **kwargs) -> list:documents = ChromaDB_VectorStore._extract_documents(self.ddl_collection.query(query_texts=[question],n_results=self.n_results_ddl,))if documents != []:scores = self.rerank_model.predict([(question, doc) for doc in documents])combined = list(zip(documents, scores))sorted_combined = sorted(combined, key=lambda x: x[1], reverse=True)documents = [item[0] for item in sorted_combined]return documents
def get_related_documentation(self, question: str, **kwargs) -> list:documents = ChromaDB_VectorStore._extract_documents(self.documentation_collection.query(query_texts=[question],n_results=self.n_results_documentation,))if documents != []:scores = self.rerank_model.predict([(question, doc) for doc in documents])combined = list(zip(documents, scores))sorted_combined = sorted(combined, key=lambda x: x[1], reverse=True)documents = [item[0] for item in sorted_combined]return documents