amina-cli 0.2.4__tar.gz → 0.2.6__tar.gz

{amina_cli-0.2.4 → amina_cli-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amina-cli
-Version: 0.2.4
+Version: 0.2.6
 Summary: CLI for AminoAnalytica protein engineering platform
 Project-URL: Homepage, https://aminoanalytica.com
 Project-URL: Documentation, https://docs.aminoanalytica.com

{amina_cli-0.2.4 → amina_cli-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "amina-cli"
-version = "0.2.4"
+version = "0.2.6"
 description = "CLI for AminoAnalytica protein engineering platform"
 readme = "README.md"
 license = {text = "Apache-2.0"}

{amina_cli-0.2.4 → amina_cli-0.2.6}/src/amina_cli/__init__.py

@@ -9,4 +9,4 @@ Quick start:
     amina run esmfold --sequence "MKFLILLFNILCLFPVLAADNH"
 """
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"

{amina_cli-0.2.4 → amina_cli-0.2.6}/src/amina_cli/auth.py

@@ -336,3 +336,34 @@ def update_job_status(job_id: str, status: str) -> None:
     flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
     with os.fdopen(os.open(JOBS_FILE, flags, 0o600), "w") as f:
         f.write(content)
+
+
+def update_job_info(job_id: str, updates: dict) -> None:
+    """
+    Merge fields into a saved job entry.
+
+    Used to persist a resolved call_id so subsequent polls skip the queue check.
+
+    Args:
+        job_id: Job ID to update
+        updates: Dict of fields to merge (e.g., {"call_id": "fc-xxx"})
+    """
+    import os
+
+    if not JOBS_FILE.exists():
+        return
+
+    try:
+        jobs = json.loads(JOBS_FILE.read_text())
+    except (json.JSONDecodeError, IOError):
+        return
+
+    for job in jobs:
+        if job.get("job_id") == job_id:
+            job.update(updates)
+            break
+
+    content = json.dumps(jobs, indent=2)
+    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
+    with os.fdopen(os.open(JOBS_FILE, flags, 0o600), "w") as f:
+        f.write(content)

{amina_cli-0.2.4 → amina_cli-0.2.6}/src/amina_cli/client.py

@@ -92,10 +92,12 @@ class InsufficientCreditsError(ClientError):
 class JobLimitError(ClientError):
     """Raised when user has reached job concurrency limits."""
 
-    def __init__(self, limit_type: str, current: int, limit: int, message: str):
+    def __init__(self, limit_type: str, current: int, limit: int, message: str, **kwargs: Any):
         self.limit_type = limit_type
         self.current = current
         self.limit = limit
+        self.running = kwargs.get("running", 0)
+        self.queued = kwargs.get("queued", 0)
         super().__init__(message)
 
 
@@ -185,6 +187,7 @@ async def _poll_queued_until_submitted(
     poll_interval: float = 5.0,  # Longer interval for queue polling
     max_poll_interval: float = 30.0,
     on_position_update: Optional[QueuePositionCallback] = None,
+    headers: Optional[dict[str, str]] = None,
 ) -> tuple[str, str]:
     """
     Poll queued_job_status until job is submitted.
@@ -229,7 +232,7 @@ async def _poll_queued_until_submitted(
         status_url = f"{queued_endpoint}?job_id={job_id}"
 
         try:
-            response = await client.get(status_url)
+            response = await client.get(status_url, headers=headers or {})
         except httpx.RequestError as e:
             # Log transient error and continue polling
             print(f"Queue poll network error (will retry): {type(e).__name__}: {e}")
@@ -396,6 +399,8 @@ async def run_tool(
                     current=detail.get("current", 0),
                     limit=detail.get("limit", 0),
                     message=detail.get("message", "Job limit reached. Please wait for running jobs to complete."),
+                    running=detail.get("running", 0),
+                    queued=detail.get("queued", 0),
                 )
             except (ValueError, KeyError):
                 raise JobLimitError(
@@ -438,6 +443,7 @@ async def run_tool(
                 poll_interval=5.0,  # Slower for queue
                 max_poll_interval=30.0,
                 on_position_update=on_queue_position_update,
+                headers=auth_headers,
             )
         else:
             # Job was submitted immediately
@@ -663,6 +669,8 @@ async def submit_tool(
                     current=detail.get("current", 0),
                     limit=detail.get("limit", 0),
                     message=detail.get("message", "Job limit reached. Please wait for running jobs to complete."),
+                    running=detail.get("running", 0),
+                    queued=detail.get("queued", 0),
                 )
             except (ValueError, KeyError):
                 raise JobLimitError(
@@ -855,12 +863,16 @@ async def check_queued_job_status(job_id: str) -> dict[str, Any]:
         - {"status": "submitted", "call_id": "..."} if spawned
         - {"status": "failed", "error": "..."} if failed/expired
     """
+    from amina_cli.auth import get_api_key
+
     queued_endpoint = get_queued_status_endpoint()
     status_url = f"{queued_endpoint}?job_id={job_id}"
+    api_key = get_api_key()
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
 
     async with httpx.AsyncClient(timeout=30.0) as client:
         try:
-            response = await client.get(status_url)
+            response = await client.get(status_url, headers=headers)
         except httpx.RequestError as e:
             raise ToolExecutionError(f"Network error checking queued job status: {str(e)}")
 

{amina_cli-0.2.4 → amina_cli-0.2.6}/src/amina_cli/commands/jobs_cmd.py

@@ -21,6 +21,65 @@ app = typer.Typer(no_args_is_help=True)
 console = Console()
 
 
+def _resolve_job_status(job_info: dict) -> dict:
+    """
+    Resolve the current status of a job, handling both spawned and queued jobs.
+
+    If call_id is present, polls Modal directly. If missing (queued job),
+    checks the queue first and resolves call_id when available.
+
+    Args:
+        job_info: Full job info dict from local history
+
+    Returns:
+        Dict with status info (same shape as check_job_status_sync)
+    """
+    from amina_cli.client import check_job_status_sync, check_queued_job_status_sync
+    from amina_cli.auth import update_job_info
+
+    call_id = job_info.get("call_id", "")
+    job_id = job_info.get("job_id", "")
+
+    if call_id:
+        # Normal path: job was spawned, poll Modal directly
+        return check_job_status_sync(
+            call_id=call_id,
+            job_id=job_id,
+            user_id=job_info.get("user_id", ""),
+            tool_name=job_info.get("tool_name", ""),
+            reserved_cost=job_info.get("reserved_cost", 0.0),
+            compute_type=job_info.get("compute_type", "cpu"),
+        )
+
+    # Queued path: no call_id yet, check queue status
+    queue_result = check_queued_job_status_sync(job_id)
+    queue_status = queue_result.get("status", "")
+
+    if queue_status in ("queued", "not_found"):
+        return queue_result
+
+    if queue_status == "failed":
+        return queue_result
+
+    # Job has been spawned (submitted/running/completed) — try to get call_id
+    resolved_call_id = queue_result.get("call_id", "")
+    if not resolved_call_id:
+        # Race condition: status changed but call_id not yet written
+        return {"status": "queued", "message": "Job is being spawned"}
+
+    # Persist call_id so future polls skip the queue check
+    update_job_info(job_id, {"call_id": resolved_call_id})
+
+    return check_job_status_sync(
+        call_id=resolved_call_id,
+        job_id=job_id,
+        user_id=job_info.get("user_id", ""),
+        tool_name=job_info.get("tool_name", ""),
+        reserved_cost=job_info.get("reserved_cost", 0.0),
+        compute_type=job_info.get("compute_type", "cpu"),
+    )
+
+
 @app.command("list")
 def list_jobs(
     limit: int = typer.Option(
@@ -101,7 +160,7 @@ def status(
         amina jobs status abc123 --json
     """
     from amina_cli.auth import get_job_info
-    from amina_cli.client import check_job_status_sync, ToolExecutionError
+    from amina_cli.client import ToolExecutionError
 
     results = []
 
@@ -121,14 +180,7 @@ def status(
 
         # Query server for current status
         try:
-            status_result = check_job_status_sync(
-                call_id=job_info.get("call_id", ""),
-                job_id=job_info.get("job_id", job_id),
-                user_id=job_info.get("user_id", ""),
-                tool_name=job_info.get("tool_name", ""),
-                reserved_cost=job_info.get("reserved_cost", 0.0),
-                compute_type=job_info.get("compute_type", "cpu"),
-            )
+            status_result = _resolve_job_status(job_info)
             results.append(
                 {
                     "job_id": job_info.get("job_id", job_id),
@@ -212,7 +264,7 @@ def wait(
         amina jobs wait abc123 --timeout 7200 --poll-interval 30
     """
     from amina_cli.auth import get_job_info
-    from amina_cli.client import check_job_status_sync, ToolExecutionError
+    from amina_cli.client import ToolExecutionError
 
     start_time = time.time()
     pending_jobs = set(job_ids)
@@ -237,14 +289,7 @@ def wait(
                 continue
 
             try:
-                status_result = check_job_status_sync(
-                    call_id=job_info.get("call_id", ""),
-                    job_id=job_info.get("job_id", job_id),
-                    user_id=job_info.get("user_id", ""),
-                    tool_name=job_info.get("tool_name", ""),
-                    reserved_cost=job_info.get("reserved_cost", 0.0),
-                    compute_type=job_info.get("compute_type", "cpu"),
-                )
+                status_result = _resolve_job_status(job_info)
 
                 if status_result.get("status") in ("completed", "failed"):
                     results[job_id] = {
@@ -312,7 +357,7 @@ def download(
         amina jobs download abc123 -o ./results/
     """
     from amina_cli.auth import get_job_info
-    from amina_cli.client import check_job_status_sync, ToolExecutionError
+    from amina_cli.client import ToolExecutionError
     from amina_cli.storage import download_results, StorageError
 
     # Look up job info
@@ -324,18 +369,16 @@ def download(
 
     # Check status
     try:
-        status_result = check_job_status_sync(
-            call_id=job_info.get("call_id", ""),
-            job_id=job_info.get("job_id", job_id),
-            user_id=job_info.get("user_id", ""),
-            tool_name=job_info.get("tool_name", ""),
-            reserved_cost=job_info.get("reserved_cost", 0.0),
-            compute_type=job_info.get("compute_type", "cpu"),
-        )
+        status_result = _resolve_job_status(job_info)
     except ToolExecutionError as e:
         console.print(f"[red]Error checking job status:[/red] {e}")
         raise typer.Exit(1)
 
+    if status_result.get("status") == "queued":
+        console.print("[yellow]Job is still queued.[/yellow] Wait for completion first:")
+        console.print(f"  amina jobs wait {job_id}")
+        raise typer.Exit(1)
+
     if status_result.get("status") == "running":
         console.print("[yellow]Job is still running.[/yellow] Wait for completion first:")
         console.print(f"  amina jobs wait {job_id}")

{amina_cli-0.2.4 → amina_cli-0.2.6}/src/amina_cli/commands/tools/properties/esm2_embedding.py

@@ -9,16 +9,20 @@ METADATA = {
     "name": "esm2-embedding",
     "display_name": "ESM2 Embedding",
     "category": "properties",
-    "description": "Extract sequence-level protein embeddings from FASTA files using ESM2 models",
+    "description": "Extract protein embeddings from FASTA files using ESM2 models",
     "modal_function_name": "esm2_embedding_worker",
     "modal_app_name": "esm2-embedding-api",
     "status": "available",
     "outputs": {
+        # Sequence mode outputs
         "csv_filepath": "CSV file with sequence embeddings (sequence_id, sequence, dim_0...dim_N)",
         "tsne_plot_filepath": "Static t-SNE visualization of embeddings (PNG)",
         "tsne_plot_html_filepath": "Interactive t-SNE visualization with zoom/pan/hover (HTML)",
         "clusters_csv_filepath": "Cluster assignments (CSV, unless --no-cluster)",
         "cluster_summary_filepath": "Cluster statistics and quality metrics (JSON, unless --no-cluster)",
+        # Residue mode outputs
+        "manifest_filepath": "Manifest JSON with file list and metadata (residue mode)",
+        "residue_N": "Per-sequence NPZ files with residue embeddings (residue mode)",
     },
 }
 
@@ -76,6 +80,17 @@ def register(app: typer.Typer):
             "--cluster-min-samples",
             help="Min samples for density estimation (auto-tuned if not specified)",
         ),
+        embedding_mode: str = typer.Option(
+            "sequence",
+            "--mode",
+            "-M",
+            help="Embedding mode: 'sequence' (default) for mean-pooled embeddings, 'residue' for per-residue embeddings",
+        ),
+        include_special_tokens: bool = typer.Option(
+            False,
+            "--include-special-tokens",
+            help="Include BOS/EOS tokens in residue embeddings (residue mode only)",
+        ),
         job_name: Optional[str] = typer.Option(
             None,
             "--job-name",
@@ -94,8 +109,11 @@ def register(app: typer.Typer):
         Generates dense vector representations (embeddings) of protein sequences
         that capture evolutionary, structural, and functional information.
 
-        By default, also performs HDBSCAN clustering to group similar sequences.
-        Use --no-cluster to disable clustering.
+        Two embedding modes:
+          - sequence (default): Mean-pooled sequence embeddings (CSV output)
+            with optional t-SNE visualization and HDBSCAN clustering
+          - residue: Per-residue embeddings (NPZ output per sequence)
+            with manifest JSON, no t-SNE or clustering
 
         Model variants (embedding dimensions):
           - 650M: 1280-dim - Highest accuracy (default)
@@ -104,11 +122,15 @@ def register(app: typer.Typer):
           - 8M:   320-dim  - Fastest, lower accuracy
 
         Examples:
+            # Sequence mode (default)
             amina run esm2-embedding --fasta sequences.fasta -o ./output/
             amina run esm2-embedding -f seqs.fa -m 8M -o ./output/  # Fast mode
             amina run esm2-embedding -f seqs.fa --no-tsne -o ./output/
-            amina run esm2-embedding -f seqs.fa --no-cluster -o ./output/  # Skip clustering
-            amina run esm2-embedding -f seqs.fa --cluster-min-size 5 -o ./output/  # Custom params
+            amina run esm2-embedding -f seqs.fa --no-cluster -o ./output/
+
+            # Residue mode
+            amina run esm2-embedding -f seqs.fa --mode residue -o ./output/
+            amina run esm2-embedding -f seqs.fa -M residue --include-special-tokens -o ./output/
         """
         # Validate required options
         if output is None and not background:
@@ -128,6 +150,26 @@ def register(app: typer.Typer):
             console.print("[red]Error:[/red] Batch size must be between 1 and 32")
             raise typer.Exit(1)
 
+        # Validate embedding mode
+        valid_modes = ["sequence", "residue"]
+        if embedding_mode.lower() not in valid_modes:
+            console.print(f"[red]Error:[/red] Invalid mode '{embedding_mode}'. Choose from: {', '.join(valid_modes)}")
+            raise typer.Exit(1)
+        embedding_mode = embedding_mode.lower()
+
+        # Validate mode-specific options
+        if embedding_mode == "sequence" and include_special_tokens:
+            console.print("[red]Error:[/red] --include-special-tokens only applies to residue mode")
+            raise typer.Exit(1)
+
+        if embedding_mode == "residue":
+            if not no_tsne:
+                console.print("[yellow]Note:[/yellow] t-SNE visualization is not available in residue mode (skipping)")
+                no_tsne = True
+            if not no_cluster:
+                console.print("[yellow]Note:[/yellow] Clustering is not available in residue mode (skipping)")
+                no_cluster = True
+
         # Read FASTA content
         fasta_content = fasta.read_text()
         if not fasta_content.strip().startswith(">"):
@@ -148,6 +190,8 @@ def register(app: typer.Typer):
             "cluster_by_embeddings": not no_cluster,
             "cluster_min_size": cluster_min_size,
             "cluster_min_samples": cluster_min_samples,
+            "embedding_mode": embedding_mode,
+            "include_special_tokens": include_special_tokens,
         }
 
         if job_name: