"""Use Carbon-X as a tool for any chat model with OpenAI-style tool calling.

Serve Carbon-X first: vllm serve HuggingFaceBio/Carbon-X --trust-remote-code
Then set the agent client and model below and run: python agent_tool.py
"""
import json

from openai import OpenAI

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


def ask_carbon(question: str, sequences: dict[str, str], context: str = "") -> str:
    """Ask Carbon-X a question about one or more named DNA sequences."""
    parts = [context, question] if context else [question]
    parts += [f"{name}: <dna>{seq}</dna>" for name, seq in sequences.items()]
    response = 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


sequences = {"reference.fa": "ATGGTGCATCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGG",
             "variant.fa": "ATGGTGCATCTGACTCCTGTGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGG"}

carbon_tool = {
    "type": "function",
    "function": {
        "name": "ask_carbon",
        "description": (
            "Ask Carbon-X, a DNA language model, a question about DNA sequences: variant effects, "
            "regulatory elements, organism of origin, sequence features. Refer to sequences by name; "
            "do not paste DNA. Its answers are model predictions, not verified facts."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "question": {"type": "string", "description": "One self-contained question."},
                "sequences": {"type": "array", "items": {"type": "string", "enum": list(sequences)}},
                "context": {"type": "string", "description": "Known gene, coordinates or annotations, if any."},
            },
            "required": ["question", "sequences"],
        },
    },
}

agent = OpenAI()  # the agent model's provider
messages = [
    {"role": "system", "content": "For questions about DNA, consult Carbon-X with the ask_carbon tool before answering."},
    {"role": "user", "content": "variant.fa differs from reference.fa by one base. Is the change likely benign or pathogenic?"},
]
while True:
    reply = agent.chat.completions.create(model="<agent model>", messages=messages, tools=[carbon_tool])
    message = reply.choices[0].message
    messages.append(message)
    if not message.tool_calls:
        break
    for call in message.tool_calls:
        args = json.loads(call.function.arguments)
        answer = ask_carbon(args["question"], {name: sequences[name] for name in args["sequences"]}, args.get("context", ""))
        messages.append({"role": "tool", "tool_call_id": call.id, "content": answer})
print(message.content)