Foldseek Structural Search

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs Foldseek search for a PDB/mmCIF file against different databases."""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

import argparse
import json
import os
import sys
import time
import uuid

from science_skills.scienceskillscommon import http_client

ALLOWED_DATABASES = [
    "afdb50",
    "afdb-swissprot",
    "pdb100",
    "BFVD",
    "mgnify_esm30",
    "cath50",
    "gmgcl_id",
    "bfmd",
    "afdb-proteome",
]

MAX_ALIGNMENT_HITS = 300
DEFAULT_EVALUE = 1000

# Respect 0.1 queries per second requirement.
CLIENT = http_client.HttpClient("https://search.foldseek.com", qps=1)


def build_multipart_payload(fields, files):
  """Manually constructs a multipart/form-data byte payload."""
  boundary = uuid.uuid4().hex
  body = bytearray()

  # Add standard form fields (handling lists for multiple databases)
  for key, values in fields.items():
    if not isinstance(values, list):
      values = [values]
    for value in values:
      body.extend(f"--{boundary}\r\n".encode("utf-8"))
      body.extend(
          f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode(
              "utf-8"
          )
      )
      body.extend(f"{value}\r\n".encode("utf-8"))

  # Add file data
  for key, filepath in files.items():
    filename = os.path.basename(filepath)
    with open(filepath, "rb") as f:
      content = f.read()
    body.extend(f"--{boundary}\r\n".encode("utf-8"))
    body.extend(
        f'Content-Disposition: form-data; name="{key}";'
        f' filename="{filename}"\r\n'.encode("utf-8")
    )
    body.extend(b"Content-Type: application/octet-stream\r\n\r\n")
    body.extend(content)
    body.extend(b"\r\n")

  body.extend(f"--{boundary}--\r\n".encode("utf-8"))
  return boundary, body


# ---------------------------------------------------


def main():
  # Set up command line argument parsing
  parser = argparse.ArgumentParser(
      description="Query Foldseek with a PDB/mmCIF file and save the results."
  )
  parser.add_argument("input_file", help="Path to the mmCIF or PDB file")
  parser.add_argument(
      "-o",
      "--output",
      help="Path to save the output JSON file",
      default="foldseek_results.json",
  )
  parser.add_argument(
      "--databases",
      help="Comma-separated list of databases to search",
      default="pdb100,afdb50",
  )
  args = parser.parse_args()

  file_path = args.input_file
  output_path = args.output
  ticket_url = "https://search.foldseek.com/api/ticket"

  # Process and validate the databases
  selected_dbs = [db.strip() for db in args.databases.split(",")]
  invalid_dbs = [db for db in selected_dbs if db not in ALLOWED_DATABASES]

  if invalid_dbs:
    print(f"[!] Error: Invalid database(s) provided: {', '.join(invalid_dbs)}")
    print(f"[*] Allowed databases are: {', '.join(ALLOWED_DATABASES)}")
    sys.exit(1)

  # Standard headers for all requests
  headers = {
      "User-Agent": "",
      "Accept": "application/json",
  }

  print(f"[*] Submitting {file_path} to Foldseek API...")
  print(f"[*] Searching databases: {', '.join(selected_dbs)}")

  # 1. Submit Ticket
  try:
    if not os.path.exists(file_path):
      raise FileNotFoundError(f"No such file: '{file_path}'")

    boundary, body = build_multipart_payload(
        fields={"mode": "3diaa", "database[]": se