File size: 4,045 Bytes
cd0c7a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import httpx
import asyncio
import hashlib
import json
from typing import Any
from app.tools.base import BaseTool
from app.config import settings
from app.services.cache import ttl_cache
class BlastTool(BaseTool):
name = "blast"
POLL_INTERVAL = 3.0
MAX_POLL_TIME = 180
@ttl_cache(ttl=86400, prefix="blast")
async def run(self, input: dict) -> dict:
sequence = input.get("sequence", "").strip()
database = input.get("database", "uniprotkb_swissprot")
program = input.get("program", "blastp")
max_hits = input.get("max_hits", 10)
job_id = await self._submit(sequence, program, database)
status = await self._poll(job_id)
if status != "FINISHED":
return {"error": f"BLAST job {job_id} ended with status {status}", "hits": []}
hits = await self._fetch_results(job_id)
parsed = self._parse_hits(hits, max_hits)
return {"hits": parsed, "count": len(parsed), "source": "EBI BLAST", "database": database}
async def _submit(self, sequence: str, program: str, database: str) -> str:
stype = "protein" if program == "blastp" else "dna"
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
f"{settings.EBI_BASE_URL}/run",
data={"email": settings.NCBI_EMAIL, "sequence": sequence, "program": program, "database": database, "stype": stype},
)
resp.raise_for_status()
return resp.text.strip()
async def _poll(self, job_id: str) -> str:
start = asyncio.get_event_loop().time()
async with httpx.AsyncClient(timeout=15) as client:
consecutive_failures = 0
while True:
elapsed = asyncio.get_event_loop().time() - start
if elapsed > self.MAX_POLL_TIME:
return "TIMEOUT"
try:
resp = await client.get(f"{settings.EBI_BASE_URL}/status/{job_id}")
resp.raise_for_status()
status = resp.text.strip()
consecutive_failures = 0
except Exception as e:
consecutive_failures += 1
if consecutive_failures >= 5:
return "ERROR"
await asyncio.sleep(self.POLL_INTERVAL)
continue
if status in ("FINISHED", "ERROR", "FAILED"):
return status
await asyncio.sleep(self.POLL_INTERVAL)
async def _fetch_results(self, job_id: str) -> list[dict]:
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.get(f"{settings.EBI_BASE_URL}/result/{job_id}/json")
resp.raise_for_status()
data = resp.json()
return data.get("hits", [])
def _parse_hits(self, raw_hits: list[dict], max_hits: int) -> list[dict]:
parsed = []
for hit in raw_hits[:max_hits]:
hsps = hit.get("hsps", [{}])[0] if hit.get("hsps") else {}
desc = hit.get("hit_desc", "")
organism = ""
if "[" in desc and "]" in desc:
organism = desc.split("[")[-1].rstrip("]")
desc = desc.split("[")[0].strip()
parsed.append({
"accession": hit.get("hit_acc", ""),
"id": hit.get("hit_id", ""),
"description": desc,
"organism": organism,
"evalue": hsps.get("hsp_expect", 0),
"bit_score": hsps.get("hsp_bit_score", 0),
"identity_pct": hsps.get("hsp_identity", 0),
"alignment_length": hsps.get("hsp_align_len", 0),
"query_coverage_pct": 0,
"query_from": hsps.get("hsp_query_from", 0),
"query_to": hsps.get("hsp_query_to", 0),
"hit_from": hsps.get("hsp_hit_from", 0),
"hit_to": hsps.get("hsp_hit_to", 0),
})
return parsed
|