"""Use Carbon-X as a tool for an Inspect react() agent.

Serve Carbon-X first: vllm serve HuggingFaceBio/Carbon-X --trust-remote-code
Then run: python inspect_agent.py
"""
from inspect_ai import Task, eval, task
from inspect_ai.agent import as_solver, react
from inspect_ai.dataset import Sample
from inspect_ai.scorer import includes
from inspect_ai.solver import solver
from inspect_ai.tool import ToolError, tool
from inspect_ai.util import store
from openai import AsyncOpenAI

carbon = AsyncOpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")


@tool
def ask_carbon():
    async def execute(question: str, sequences: list[str], context: str = "") -> str:
        """Ask Carbon-X, a DNA language model, about this task's DNA sequences. Its answers are predictions.

        Args:
            question: One self-contained question.
            sequences: Names of the sequences to send, as given in the task.
            context: Known facts, such as the gene, positions or alleles.
        """
        dna = store().get("sequences")
        if missing := [name for name in sequences if name not in dna]:
            raise ToolError(f"Unknown sequence names: {missing}. Available: {list(dna)}")
        parts = [context, question] if context else [question]
        parts += [f"{name}: <dna>{dna[name]}</dna>" for name in sequences]
        response = await carbon.chat.completions.create(
            model="HuggingFaceBio/Carbon-X",
            messages=[{"role": "user", "content": "\n\n".join(parts)}],
            max_tokens=4096,
            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
        )
        return response.choices[0].message.content

    return execute


@solver
def load_sequences():
    async def solve(state, generate):
        store().set("sequences", state.metadata["sequences"])  # the tool reads DNA from here
        return state

    return solve


@task
def carbon_agent():
    sample = Sample(
        input="Sequence `variant` differs from `reference` at position 20 (A>T). Is the variant benign or pathogenic?",
        target="pathogenic",
        metadata={"sequences": {
            "reference": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGG",
            "variant": "ATGGTGCATCTGACTCCTGTGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGG",
        }},
    )
    agent = react(prompt="Consult Carbon-X with the ask_carbon tool before you answer.", tools=[ask_carbon()])
    return Task(dataset=[sample], solver=[load_sequences(), as_solver(agent)], scorer=includes())


eval(carbon_agent(), model="openai/gpt-5")  # any Inspect model with tool calling as the agent