Adding retrieval augmented generation

cli/extra_tools/RetrievalAugmentedEmbedder.py → cli/extra_tools/RAG/RetrievalAugmentedEmbedder.py

@@ -172,7 +172,11 @@ def embedding_pipeline(self, url:str) -> None:
             self.add_function({"signature": func_definition, "embedding": embedding_content})
         else:
             console.log(f"[yellow] Embedding for url {url} exists.")
-
+
+    def embedding_pipeline_functions(self):
+        for i in range(len(self.functions)):
+            embedding_content = self.functions[i]["embedding"]
+            self.add_embedding(embedding_content)
 
     @staticmethod
     def cosine_similarity(A: np.ndarray, B: List[np.ndarray]) -> List[float]:
@@ -277,58 +281,10 @@ def clear(self) -> None:
 
 if __name__ == "__main__":
     rag = RetrievalAugmentedEmbedder()
-    urls = [
-    "https://docs.scvi-tools.org/en/stable/api/reference/scvi.model.SCVI.html",
-    "https://docs.scvi-tools.org/en/stable/api/reference/scvi.model.SCANVI.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.calculate_qc_metrics.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.filter_cells.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.filter_genes.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.normalize_total.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.log1p.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.highly_variable_genes.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.regress_out.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.scale.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.pca.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.normalize_pearson_residuals.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.flag_gene_family.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.filter_highly_variable.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.harmony_integrate.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.scrublet.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.scrublet_simulate_doublets.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.neighbors.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.bbknn.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.umap.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.tsne.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.diffmap.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.draw_graph.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.mde.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.embedding_density.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.louvain.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.leiden.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.kmeans.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.score_genes.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.score_genes_cell_cycle.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.tl.rank_genes_groups_logreg.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.get.aggregate.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.get.anndata_to_GPU.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.get.anndata_to_CPU.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.get.X_to_GPU.html",
-    "https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.get.X_to_CPU.html",
-
-    # CellTypist
-    "https://celltypist.readthedocs.io/en/latest/celltypist.train.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.annotate.html#celltypist.annotate",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.dotplot.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.models.download_models.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.samples.downsample_adata.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.classifier.AnnotationResult.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.classifier.Classifier.html",
-    "https://celltypist.readthedocs.io/en/latest/celltypist.models.Model.html",
-    ]
-
-    for url in urls:
-        rag.embedding_pipeline(url)    
-    print(rag.query("What is pca"))
+    #rag.embedding_pipeline_functions()
+    print(rag.query("Find a function to download the model"))
+    print(rag.query("AttributeError: module 'celltypist.models' has no attribute 'download_model'"))  
+    rag.cosine_distance_heatmap()