app/tools/uniprot.py

import re
import httpx
from typing import Any
from app.tools.base import BaseTool
from app.config import settings
from app.services.cache import ttl_cache


class UniprotTool(BaseTool):
    name = "uniprot"

    @ttl_cache(ttl=86400, prefix="uniprot")
    async def run(self, input: dict) -> dict:
        accession = input.get("accession", "").strip().upper()
        if not accession:
            return {"error": "No accession provided"}

        data = await self._fetch(accession)
        if "error" in data:
            return data

        return {
            "accession": data.get("primaryAccession", ""),
            "full_name": self._extract_name(data),
            "ec_number": (data.get("proteinDescription", {}) or {}).get("ecNumbers", [{}])[0].get("ecNumber", "") if data.get("proteinDescription") else "",
            "gene_names": [g.get("geneName", {}).get("value", "") for g in (data.get("genes") or []) if g.get("geneName")],
            "organism": ((data.get("organism", {}) or {}).get("scientificName", "")),
            "functions": self._extract_functions(data),
            "keywords": [kw.get("name", "") for kw in (data.get("keywords") or [])],
            "sequence": (data.get("sequence", {}) or {}).get("value", ""),
            "sequence_length": ((data.get("sequence", {}) or {}).get("length", 0)),
            "subcellular_locations": self._extract_locations(data),
            "pdb_ids": self._extract_pdb(data),
            "features": self._extract_features(data),
            "go_terms": self._extract_go_terms(data),
            "cds_accessions": self._extract_cds_accessions(data),
        }

    async def _fetch(self, accession: str) -> dict:
        accession = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', accession)
        url = f"{settings.UNIPROT_BASE_URL}/{accession}"
        async with httpx.AsyncClient(timeout=15) as client:
            resp = await client.get(url, params={"format": "json"})
            if resp.status_code == 404:
                return {"error": f"Accession {accession} not found"}
            if resp.status_code >= 400:
                return {"error": f"UniProt returned {resp.status_code} for {accession}"}
            return resp.json()

    def _extract_name(self, data: dict) -> str:
        desc = data.get("proteinDescription", {}) or {}
        rec_name = desc.get("recommendedName", {}) or {}
        return (rec_name.get("fullName", {}) or {}).get("value", "")

    def _extract_functions(self, data: dict) -> list[str]:
        comments = data.get("comments") or []
        funcs = []
        for c in comments:
            if c.get("commentType") == "FUNCTION":
                texts = c.get("texts") or []
                for t in texts:
                    val = (t.get("value") or "").strip()
                    if val:
                        funcs.append(val)
        return funcs

    def _extract_locations(self, data: dict) -> list[str]:
        comments = data.get("comments") or []
        locs = []
        for c in comments:
            if c.get("commentType") == "SUBCELLULAR_LOCATION":
                subcels = c.get("subcellularLocations") or []
                for s in subcels:
                    loc = (s.get("location", {}) or {}).get("value", "")
                    if loc:
                        locs.append(loc)
        return locs

    def _extract_cds_accessions(self, data: dict) -> list[dict]:
        refs = data.get("uniProtKBCrossReferences") or []
        cds = []
        seen_ids = set()
        for r in refs:
            db = r.get("database", "")
            if db in ("EMBL", "GenBank", "DDBJ"):
                props = {p.get("key", ""): p.get("value", "") for p in (r.get("properties") or [])}
                acc = r.get("id", "")
                if acc and acc not in seen_ids:
                    seen_ids.add(acc)
                    cds.append({
                        "database": db,
                        "accession": acc,
                        "protein_sequence_id": props.get("protein sequence ID", ""),
                        "nucleotide_sequence_id": props.get("nucleotide sequence ID", ""),
                    })
        return cds

    def _extract_pdb(self, data: dict) -> list[str]:
        refs = data.get("uniProtKBCrossReferences") or []
        pdbs = []
        for r in refs:
            if r.get("database") == "PDB":
                pdbs.append(r.get("id", ""))
        return pdbs

    def _extract_features(self, data: dict) -> list[dict]:
        features = data.get("features") or []
        result = []
        for f in features:
            result.append({
                "type": f.get("type", ""),
                "description": f.get("description", ""),
                "begin": (f.get("location", {}) or {}).get("start", {}).get("value"),
                "end": (f.get("location", {}) or {}).get("end", {}).get("value"),
            })
        return result

    def _extract_go_terms(self, data: dict) -> list[str]:
        refs = data.get("uniProtKBCrossReferences") or []
        go = []
        for r in refs:
            if r.get("database") == "GO":
                term = r.get("properties", [{}])[0].get("value", "") if r.get("properties") else ""
                if term:
                    go.append(term)
        return go