tooluniverse 0.2.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

tooluniverse/embedding_sync.py

@@ -0,0 +1,396 @@
+"""
+Embedding Sync Tool for ToolUniverse
+
+Synchronize embedding databases with HuggingFace Hub for sharing and collaboration.
+Supports uploading local databases to HuggingFace and downloading databases from HuggingFace.
+"""
+
+import os
+import json
+import shutil
+from pathlib import Path
+from typing import Dict
+from datetime import datetime
+
+try:
+    from huggingface_hub import HfApi, upload_folder, snapshot_download
+    from huggingface_hub.utils import HfHubHTTPError
+except ImportError:
+    raise ImportError(
+        "huggingface_hub is required. Install with: pip install huggingface_hub"
+    )
+
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+from .logging_config import get_logger
+
+
+@register_tool("EmbeddingSync")
+class EmbeddingSync(BaseTool):
+    """
+    Sync embedding databases with HuggingFace Hub.
+    Supports uploading local databases and downloading shared databases.
+    """
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        self.logger = get_logger("EmbeddingSync")
+
+        # HuggingFace configuration
+        hf_config = tool_config.get("configs", {}).get("huggingface_config", {})
+        self.hf_token = hf_config.get("token") or os.getenv("HF_TOKEN")
+        self.hf_endpoint = hf_config.get("endpoint", "https://huggingface.co")
+
+        if not self.hf_token:
+            self.logger.warning(
+                "HuggingFace token not found. Some operations may fail."
+            )
+
+        # Initialize HF API
+        self.hf_api = HfApi(endpoint=self.hf_endpoint, token=self.hf_token)
+
+        # Storage configuration
+        storage_config = tool_config.get("configs", {}).get("storage_config", {})
+        self.data_dir = Path(storage_config.get("data_dir", "./data/embeddings"))
+        self.export_dir = Path(storage_config.get("export_dir", "./exports"))
+
+        # Ensure directories exist
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        self.export_dir.mkdir(parents=True, exist_ok=True)
+
+    def run(self, arguments):
+        """Main entry point for the tool"""
+        action = arguments.get("action")
+
+        if action == "upload":
+            return self._upload_to_huggingface(arguments)
+        elif action == "download":
+            return self._download_from_huggingface(arguments)
+        else:
+            return {"error": f"Unknown action: {action}"}
+
+    def _upload_to_huggingface(self, arguments):
+        """Upload local database to HuggingFace Hub"""
+        database_name = arguments.get("database_name")
+        repository = arguments.get("repository")
+        description = arguments.get("description", "")
+        private = arguments.get("private", False)
+        commit_message = arguments.get(
+            "commit_message", f"Upload {database_name} database"
+        )
+
+        if not database_name:
+            return {"error": "database_name is required"}
+        if not repository:
+            return {"error": "repository is required (format: username/repo-name)"}
+        if not self.hf_token:
+            return {"error": "HuggingFace token required for upload operations"}
+
+        try:
+            # Check if local database exists
+            db_path = self.data_dir / "embeddings.db"
+            index_path = self.data_dir / f"{database_name}.faiss"
+
+            if not db_path.exists():
+                return {"error": "Local embeddings database not found"}
+            if not index_path.exists():
+                return {
+                    "error": f"FAISS index for database '{database_name}' not found"
+                }
+
+            # Create export directory for this upload
+            export_path = (
+                self.export_dir
+                / f"{database_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            )
+            export_path.mkdir(parents=True, exist_ok=True)
+
+            # Copy database files to export directory
+            shutil.copy2(db_path, export_path / "embeddings.db")
+            shutil.copy2(index_path, export_path / f"{database_name}.faiss")
+
+            # Create database info file
+            db_info = self._get_database_info(database_name)
+            if not db_info:
+                return {
+                    "error": f"Database '{database_name}' not found in local storage"
+                }
+
+            info_file = {
+                "database_name": database_name,
+                "description": description,
+                "embedding_model": db_info.get("embedding_model"),
+                "embedding_dimensions": db_info.get("embedding_dimensions"),
+                "document_count": db_info.get("document_count"),
+                "created_at": db_info.get("created_at"),
+                "uploaded_at": datetime.now().isoformat(),
+                "version": "1.0.0",
+                "format": "tooluniverse_embedding_db",
+            }
+
+            with open(export_path / "database_info.json", "w") as f:
+                json.dump(info_file, f, indent=2)
+
+            # Create README file
+            readme_content = self._generate_readme(database_name, description, db_info)
+            with open(export_path / "README.md", "w") as f:
+                f.write(readme_content)
+
+            # Create repository if it doesn't exist
+            try:
+                self.hf_api.repo_info(repository, repo_type="dataset")
+                self.logger.info(f"Repository {repository} already exists")
+            except HfHubHTTPError:
+                self.logger.info(f"Creating new repository: {repository}")
+                self.hf_api.create_repo(
+                    repo_id=repository, repo_type="dataset", private=private
+                )
+
+            # Upload files to HuggingFace
+            self.logger.info(f"Uploading database to {repository}")
+            upload_folder(
+                folder_path=str(export_path),
+                repo_id=repository,
+                repo_type="dataset",
+                token=self.hf_token,
+                commit_message=commit_message,
+            )
+
+            # Clean up export directory
+            shutil.rmtree(export_path)
+
+            return {
+                "status": "success",
+                "database_name": database_name,
+                "repository": repository,
+                "document_count": db_info.get("document_count"),
+                "upload_url": f"{self.hf_endpoint}/datasets/{repository}",
+            }
+
+        except Exception as e:
+            self.logger.error(f"Error uploading to HuggingFace: {str(e)}")
+            return {"error": f"Failed to upload: {str(e)}"}
+
+    def _download_from_huggingface(self, arguments):
+        """Download database from HuggingFace Hub"""
+        repository = arguments.get("repository")
+        local_name = arguments.get("local_name")
+        overwrite = arguments.get("overwrite", False)
+
+        if not repository:
+            return {"error": "repository is required (format: username/repo-name)"}
+        if not local_name:
+            local_name = repository.split("/")[-1]  # Use repo name as default
+
+        try:
+            # Check if local database already exists
+            if self._local_database_exists(local_name) and not overwrite:
+                return {
+                    "error": f"Local database '{local_name}' already exists. Use overwrite=true to replace."
+                }
+
+            # Download repository to temporary directory
+            temp_dir = (
+                self.export_dir / f"download_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            )
+
+            self.logger.info(f"Downloading database from {repository}")
+            snapshot_download(
+                repo_id=repository,
+                repo_type="dataset",
+                local_dir=str(temp_dir),
+                token=self.hf_token,
+            )
+
+            # Verify required files exist
+            db_file = temp_dir / "embeddings.db"
+            info_file = temp_dir / "database_info.json"
+
+            if not db_file.exists():
+                shutil.rmtree(temp_dir)
+                return {"error": "Downloaded repository does not contain embeddings.db"}
+
+            if not info_file.exists():
+                shutil.rmtree(temp_dir)
+                return {
+                    "error": "Downloaded repository does not contain database_info.json"
+                }
+
+            # Load database info
+            with open(info_file) as f:
+                db_info = json.load(f)
+
+            original_name = db_info.get("database_name")
+            faiss_file = temp_dir / f"{original_name}.faiss"
+
+            if not faiss_file.exists():
+                shutil.rmtree(temp_dir)
+                return {
+                    "error": f"FAISS index file {original_name}.faiss not found in download"
+                }
+
+            # Copy files to local storage with new name
+            local_db_path = self.data_dir / "embeddings.db"
+            local_index_path = self.data_dir / f"{local_name}.faiss"
+
+            # Handle database file (merge or replace)
+            if local_db_path.exists() and not overwrite:
+                # Merge databases (simplified approach - copy tables)
+                self._merge_databases(
+                    str(db_file), str(local_db_path), original_name, local_name
+                )
+            else:
+                shutil.copy2(db_file, local_db_path)
+                self._rename_database_in_db(
+                    str(local_db_path), original_name, local_name
+                )
+
+            # Copy FAISS index
+            shutil.copy2(faiss_file, local_index_path)
+
+            # Clean up
+            shutil.rmtree(temp_dir)
+
+            return {
+                "status": "success",
+                "repository": repository,
+                "local_name": local_name,
+                "document_count": db_info.get("document_count"),
+                "embedding_model": db_info.get("embedding_model"),
+                "downloaded_at": datetime.now().isoformat(),
+            }
+
+        except Exception as e:
+            self.logger.error(f"Error downloading from HuggingFace: {str(e)}")
+            # Clean up on error
+            if "temp_dir" in locals() and temp_dir.exists():
+                shutil.rmtree(temp_dir)
+            return {"error": f"Failed to download: {str(e)}"}
+
+    def _get_database_info(self, database_name: str) -> Dict:
+        """Get database information from local SQLite"""
+        import sqlite3
+
+        db_path = self.data_dir / "embeddings.db"
+        if not db_path.exists():
+            return {}
+
+        try:
+            with sqlite3.connect(db_path) as conn:
+                cursor = conn.execute(
+                    """
+                    SELECT name, description, embedding_model, embedding_dimensions, document_count, created_at
+                    FROM databases WHERE name = ?
+                """,
+                    (database_name,),
+                )
+                row = cursor.fetchone()
+                if row:
+                    return {
+                        "name": row[0],
+                        "description": row[1],
+                        "embedding_model": row[2],
+                        "embedding_dimensions": row[3],
+                        "document_count": row[4],
+                        "created_at": row[5],
+                    }
+        except Exception as e:
+            self.logger.error(f"Error getting database info: {str(e)}")
+
+        return {}
+
+    def _local_database_exists(self, database_name: str) -> bool:
+        """Check if database exists locally"""
+        return bool(self._get_database_info(database_name))
+
+    def _generate_readme(
+        self, database_name: str, description: str, db_info: Dict
+    ) -> str:
+        """Generate README content for HuggingFace repository"""
+        return f"""# {database_name} - Embedding Database
+
+## Description
+{description or 'Embedding database created with ToolUniverse'}
+
+## Database Information
+- **Documents**: {db_info.get('document_count', 'Unknown')}
+- **Embedding Model**: {db_info.get('embedding_model', 'Unknown')}
+- **Dimensions**: {db_info.get('embedding_dimensions', 'Unknown')}
+- **Created**: {db_info.get('created_at', 'Unknown')}
+
+## Usage
+
+To use this database in ToolUniverse:
+
+```python
+from src.tooluniverse.execute_function import ToolUniverse
+
+# Download and load the database
+tu = ToolUniverse()
+sync = tu.init_tool("EmbeddingSync")
+
+# Download from HuggingFace
+sync.run({{
+    "action": "download",
+    "repository": "username/repo-name",
+    "local_name": "{database_name}"
+}})
+
+# Search the database
+db = tu.init_tool("EmbeddingDatabaseSearch")
+results = db.run({{
+    "database_name": "{database_name}",
+    "query": "your search query",
+    "top_k": 5
+}})
+```
+
+## Format
+This database uses the ToolUniverse embedding database format with FAISS vector index and SQLite metadata storage.
+"""
+
+    def _merge_databases(
+        self, source_db: str, target_db: str, source_name: str, target_name: str
+    ):
+        """Merge source database into target database (simplified implementation)"""
+        import sqlite3
+
+        # This is a simplified merge - in practice, you'd want more sophisticated handling
+        with sqlite3.connect(source_db) as source_conn:
+            with sqlite3.connect(target_db) as target_conn:
+                # Copy database record
+                source_conn.execute(
+                    "UPDATE databases SET name = ? WHERE name = ?",
+                    (target_name, source_name),
+                )
+
+                # Copy all records (simplified)
+                target_conn.execute("ATTACH DATABASE ? AS source_db", (source_db,))
+                target_conn.execute(
+                    """
+                    INSERT OR REPLACE INTO databases
+                    SELECT * FROM source_db.databases WHERE name = ?
+                """,
+                    (target_name,),
+                )
+                target_conn.execute(
+                    """
+                    INSERT INTO documents
+                    SELECT * FROM source_db.documents WHERE database_name = ?
+                """,
+                    (target_name,),
+                )
+                target_conn.execute("DETACH DATABASE source_db")
+
+    def _rename_database_in_db(self, db_path: str, old_name: str, new_name: str):
+        """Rename database in SQLite file"""
+        import sqlite3
+
+        with sqlite3.connect(db_path) as conn:
+            conn.execute(
+                "UPDATE databases SET name = ? WHERE name = ?", (new_name, old_name)
+            )
+            conn.execute(
+                "UPDATE documents SET database_name = ? WHERE database_name = ?",
+                (new_name, old_name),
+            )

tooluniverse/enrichr_tool.py

@@ -0,0 +1,266 @@
+import json
+import requests
+import urllib.parse
+import networkx as nx
+from .base_tool import BaseTool
+from .tool_registry import register_tool
+
+
+@register_tool("EnrichrTool")
+class EnrichrTool(BaseTool):
+    """
+    Tool to perform gene enrichment analysis using Enrichr.
+    """
+
+    def __init__(self, tool_config):
+        super().__init__(tool_config)
+        # Constants
+        self.enrichr_url = "https://maayanlab.cloud/Enrichr/addList"
+        self.enrichment_url = "https://maayanlab.cloud/Enrichr/enrich"
+
+    def run(self, arguments):
+        """Main entry point for the tool."""
+        genes = arguments.get("gene_list")
+        libs = arguments.get(
+            "libs",
+            [
+                "WikiPathways_2024_Human",
+                "Reactome_Pathways_2024",
+                "MSigDB_Hallmark_2020",
+                "GO_Molecular_Function_2023",
+                "GO_Biological_Process_2023",
+            ],
+        )
+        return self.enrichr_api(genes, libs)
+
+    def get_official_gene_name(self, gene_name):
+        """
+        Retrieve the official gene symbol for a given gene name or synonym using the MyGene.info API.
+
+        Parameters:
+            gene_name (str): The gene name or synonym to query.
+
+        Returns:
+            str: The official gene symbol if found; otherwise, raises an Exception.
+        """
+        # URL-encode the gene_name to handle special characters
+        encoded_gene_name = urllib.parse.quote(gene_name)
+        url = f"https://mygene.info/v3/query?q={encoded_gene_name}&fields=symbol,alias&species=human"
+
+        response = requests.get(url)
+        if response.status_code != 200:
+            return f"Error querying MyGene.info API: {response.status_code}"
+
+        data = response.json()
+        hits = data.get("hits", [])
+        if not hits:
+            return f"No data found for: {gene_name}. Please check the gene name and try again."
+
+        # Attempt to find an exact match in the official symbol or among aliases.
+        for hit in hits:
+            symbol = hit.get("symbol", "")
+            if symbol.upper() == gene_name.upper():
+                print(
+                    f"[enrichr_api] Using the official gene name: '{symbol}' instead of {gene_name}",
+                    flush=True,
+                )
+                return symbol
+            aliases = hit.get("alias", [])
+            if any(gene_name.upper() == alias.upper() for alias in aliases):
+                print(
+                    f"[enrichr_api] Using the official gene name: '{symbol}' instead of {gene_name}",
+                    flush=True,
+                )
+                return symbol
+
+        # If no exact match is found, return the symbol of the top hit.
+        top_hit = hits[0]
+        symbol = top_hit.get("symbol", None)
+        if symbol:
+            print(
+                f"[enrichr_api] Using the official gene name: '{symbol}' instead of {gene_name}",
+                flush=True,
+            )
+            return symbol
+        else:
+            return f"No official gene symbol found for: {gene_name}. Please ensure it is correct."
+
+    def submit_gene_list(self, gene_list):
+        """
+        Submit the gene list to Enrichr and return the user list ID.
+
+        Parameters:
+            gene_list (str): Newline-separated string of gene names.
+
+        Returns:
+            str: The user list ID from Enrichr.
+        """
+        payload = {
+            "list": (None, gene_list),
+            "description": (None, f"Gene list for {gene_list}"),
+        }
+        response = requests.post(self.enrichr_url, files=payload)
+
+        if not response.ok:
+            return "Error submitting gene list to Enrichr"
+
+        return json.loads(response.text)["userListId"]
+
+    def get_enrichment_results(self, user_list_id, library):
+        """
+        Fetch enrichment results for a specific library.
+
+        Parameters:
+            user_list_id (str): The user list ID from Enrichr.
+            library (str): The name of the enrichment library.
+
+        Returns:
+            dict: The enrichment results.
+        """
+        query_string = f"?userListId={user_list_id}&backgroundType={library}"
+        response = requests.get(self.enrichment_url + query_string)
+
+        if not response.ok:
+            return f"Error fetching enrichment results for {library}"
+
+        return json.loads(response.text)
+
+    def build_graph(self, genes, enrichment_results):
+        """
+        Initialize and build the graph with gene nodes and enriched terms.
+
+        Parameters:
+            genes (list): List of gene names.
+            enrichment_results (dict): Dictionary of enrichment results by library.
+
+        Returns:
+            networkx.Graph: The constructed graph.
+        """
+        G = nx.Graph()
+
+        # Add gene nodes
+        for gene in genes:
+            G.add_node(gene, type="gene")
+
+        # Add enriched terms and edges
+        for library, results in enrichment_results.items():
+            for term in results:
+                term_name = term[1]
+                associated_genes = term[5]
+                G.add_node(term_name, type="term", library=library)
+
+                for gene in associated_genes:
+                    if gene in genes:
+                        G.add_edge(gene, term_name, weight=round(term[4], 2))
+
+        return G
+
+    def rank_paths_by_weight(self, G, source, target):
+        """
+        Find and rank paths between source and target based on total edge weight.
+
+        Parameters:
+            G (networkx.Graph): The graph to search.
+            source (str): The source node.
+            target (str): The target node.
+
+        Returns:
+            list: List of tuples (path, weight) sorted by weight descending.
+        """
+        all_paths = list(nx.all_simple_paths(G, source=source, target=target))
+        path_weights = []
+
+        for path in all_paths:
+            total_weight = sum(
+                G[path[i]][path[i + 1]].get("weight", 1) for i in range(len(path) - 1)
+            )
+            path_weights.append((path, total_weight))
+
+        return sorted(path_weights, key=lambda x: x[1], reverse=True)
+
+    def rank_paths_to_term(self, G, gene, term):
+        """
+        Find and rank paths from each gene to a specified term based on total edge weight.
+
+        Parameters:
+            G (networkx.Graph): The graph to search.
+            gene (str): The source gene.
+            term (str): The target term.
+
+        Returns:
+            list or None: List of tuples (path, weight) sorted by weight descending, or None if no paths.
+        """
+        all_paths = list(nx.all_simple_paths(G, source=gene, target=term))
+        path_weights = []
+
+        for path in all_paths:
+            total_weight = sum(
+                G[path[i]][path[i + 1]].get("weight", 1) for i in range(len(path) - 1)
+            )
+            path_weights.append((path, total_weight))
+
+        if len(path_weights) != 0:
+            return sorted(path_weights, key=lambda x: x[1], reverse=True)
+        return None
+
+    def enrichr_api(self, genes, libs):
+        """
+        Main API function to perform gene enrichment analysis.
+
+        Parameters:
+            genes (list): List of gene names.
+            libs (list): List of enrichment libraries to use.
+
+        Returns:
+            tuple: (connected_path, connections) dictionaries.
+        """
+        # Convert each gene to its official name and log the result
+        genes = [self.get_official_gene_name(gene) for gene in genes]
+        print("Official gene names:", genes)
+
+        # Ensure at least two genes are provided for path ranking
+        if len(genes) < 2:
+            raise ValueError(
+                "At least two genes are required to rank paths between genes."
+            )
+
+        # Prepare the gene list for Enrichr submission
+        gene_list_str = "\n".join(genes)
+        user_list_id = self.submit_gene_list(gene_list_str)
+
+        # Retrieve enrichment results for each specified library
+        enrichment_results = {}
+        for library in libs:
+            results = self.get_enrichment_results(user_list_id, library)
+            # Safely get the top 5 results; if the library key isn't found, default to an empty list
+            enrichment_results[library] = results.get(library, [])[:5]
+
+        # Build the graph from the gene list and enrichment results
+        G = self.build_graph(genes, enrichment_results)
+
+        # Rank paths from the first gene to the second
+        ranked_paths = self.rank_paths_by_weight(G, genes[0], genes[1])
+        connected_path = {}
+        for path, weight in ranked_paths:
+            connected_path[f"Path: {path}"] = f"Total Weight: {weight}"
+
+        # Compute connectivity data for each gene and graph node
+        connections = {}
+        for gene in genes:
+            for term in G.nodes:
+                paths_to_term = self.rank_paths_to_term(G, gene, term)
+                if paths_to_term is not None:
+                    connections[f"Connectivity: {gene} - {term}"] = paths_to_term
+
+        # Check for empty outputs and print helper messages
+        if not connected_path:
+            print(
+                f"[Enrichr] No ranked paths were found between the gene pair {genes}."
+            )
+
+        if not connections:
+            print(
+                f"[Enrichr] No connection between genes and terms in the enriched graph of {genes}."
+            )
+
+        return connected_path, connections